libjxl

FORK: libjxl patches used on blog
git clone https://git.neptards.moe/blog/libjxl.git
Log | Files | Refs | Submodules | README | LICENSE

enc_cache.cc (9929B)


      1 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
      2 //
      3 // Use of this source code is governed by a BSD-style
      4 // license that can be found in the LICENSE file.
      5 
      6 #include "lib/jxl/enc_cache.h"
      7 
      8 #include <stddef.h>
      9 #include <stdint.h>
     10 
     11 #include <memory>
     12 
     13 #include "lib/jxl/base/common.h"
     14 #include "lib/jxl/base/compiler_specific.h"
     15 #include "lib/jxl/base/span.h"
     16 #include "lib/jxl/base/status.h"
     17 #include "lib/jxl/color_encoding_internal.h"
     18 #include "lib/jxl/compressed_dc.h"
     19 #include "lib/jxl/dct_util.h"
     20 #include "lib/jxl/dec_frame.h"
     21 #include "lib/jxl/enc_aux_out.h"
     22 #include "lib/jxl/enc_frame.h"
     23 #include "lib/jxl/enc_group.h"
     24 #include "lib/jxl/enc_modular.h"
     25 #include "lib/jxl/enc_quant_weights.h"
     26 #include "lib/jxl/frame_dimensions.h"
     27 #include "lib/jxl/frame_header.h"
     28 #include "lib/jxl/image.h"
     29 #include "lib/jxl/image_bundle.h"
     30 #include "lib/jxl/image_ops.h"
     31 #include "lib/jxl/passes_state.h"
     32 #include "lib/jxl/quantizer.h"
     33 
     34 namespace jxl {
     35 
     36 Status InitializePassesEncoder(const FrameHeader& frame_header,
     37                                const Image3F& opsin, const Rect& rect,
     38                                const JxlCmsInterface& cms, ThreadPool* pool,
     39                                PassesEncoderState* enc_state,
     40                                ModularFrameEncoder* modular_frame_encoder,
     41                                AuxOut* aux_out) {
     42   PassesSharedState& JXL_RESTRICT shared = enc_state->shared;
     43 
     44   enc_state->x_qm_multiplier = std::pow(1.25f, frame_header.x_qm_scale - 2.0f);
     45   enc_state->b_qm_multiplier = std::pow(1.25f, frame_header.b_qm_scale - 2.0f);
     46 
     47   if (enc_state->coeffs.size() < frame_header.passes.num_passes) {
     48     enc_state->coeffs.reserve(frame_header.passes.num_passes);
     49     for (size_t i = enc_state->coeffs.size();
     50          i < frame_header.passes.num_passes; i++) {
     51       // Allocate enough coefficients for each group on every row.
     52       JXL_ASSIGN_OR_RETURN(
     53           std::unique_ptr<ACImageT<int32_t>> coeffs,
     54           ACImageT<int32_t>::Make(kGroupDim * kGroupDim,
     55                                   shared.frame_dim.num_groups));
     56       enc_state->coeffs.emplace_back(std::move(coeffs));
     57     }
     58   }
     59   while (enc_state->coeffs.size() > frame_header.passes.num_passes) {
     60     enc_state->coeffs.pop_back();
     61   }
     62 
     63   if (enc_state->initialize_global_state) {
     64     float scale =
     65         shared.quantizer.ScaleGlobalScale(enc_state->cparams.quant_ac_rescale);
     66     DequantMatricesScaleDC(&shared.matrices, scale);
     67     shared.quantizer.RecomputeFromGlobalScale();
     68   }
     69 
     70   JXL_ASSIGN_OR_RETURN(Image3F dc,
     71                        Image3F::Create(shared.frame_dim.xsize_blocks,
     72                                        shared.frame_dim.ysize_blocks));
     73   JXL_RETURN_IF_ERROR(RunOnPool(
     74       pool, 0, shared.frame_dim.num_groups, ThreadPool::NoInit,
     75       [&](size_t group_idx, size_t _) {
     76         ComputeCoefficients(group_idx, enc_state, opsin, rect, &dc);
     77       },
     78       "Compute coeffs"));
     79 
     80   if (frame_header.flags & FrameHeader::kUseDcFrame) {
     81     CompressParams cparams = enc_state->cparams;
     82     cparams.dots = Override::kOff;
     83     cparams.noise = Override::kOff;
     84     cparams.patches = Override::kOff;
     85     cparams.gaborish = Override::kOff;
     86     cparams.epf = 0;
     87     cparams.resampling = 1;
     88     cparams.ec_resampling = 1;
     89     // The DC frame will have alpha=0. Don't erase its contents.
     90     cparams.keep_invisible = Override::kOn;
     91     JXL_ASSERT(cparams.progressive_dc > 0);
     92     cparams.progressive_dc--;
     93     // Use kVarDCT in max_error_mode for intermediate progressive DC,
     94     // and kModular for the smallest DC (first in the bitstream)
     95     if (cparams.progressive_dc == 0) {
     96       cparams.modular_mode = true;
     97       cparams.speed_tier = static_cast<SpeedTier>(
     98           std::max(static_cast<int>(SpeedTier::kTortoise),
     99                    static_cast<int>(cparams.speed_tier) - 1));
    100       cparams.butteraugli_distance =
    101           std::max(kMinButteraugliDistance,
    102                    enc_state->cparams.butteraugli_distance * 0.02f);
    103     } else {
    104       cparams.max_error_mode = true;
    105       for (size_t c = 0; c < 3; c++) {
    106         cparams.max_error[c] = shared.quantizer.MulDC()[c];
    107       }
    108       // Guess a distance that produces good initial results.
    109       cparams.butteraugli_distance =
    110           std::max(kMinButteraugliDistance,
    111                    enc_state->cparams.butteraugli_distance * 0.1f);
    112     }
    113     ImageBundle ib(&shared.metadata->m);
    114     // This is a lie - dc is in XYB
    115     // (but EncodeFrame will skip RGB->XYB conversion anyway)
    116     ib.SetFromImage(
    117         std::move(dc),
    118         ColorEncoding::LinearSRGB(shared.metadata->m.color_encoding.IsGray()));
    119     if (!ib.metadata()->extra_channel_info.empty()) {
    120       // Add placeholder extra channels to the patch image: dc_level frames do
    121       // not yet support extra channels, but the codec expects that the amount
    122       // of extra channels in frames matches that in the metadata of the
    123       // codestream.
    124       std::vector<ImageF> extra_channels;
    125       extra_channels.reserve(ib.metadata()->extra_channel_info.size());
    126       for (size_t i = 0; i < ib.metadata()->extra_channel_info.size(); i++) {
    127         JXL_ASSIGN_OR_RETURN(ImageF ch, ImageF::Create(ib.xsize(), ib.ysize()));
    128         extra_channels.emplace_back(std::move(ch));
    129         // Must initialize the image with data to not affect blending with
    130         // uninitialized memory.
    131         // TODO(lode): dc_level must copy and use the real extra channels
    132         // instead.
    133         ZeroFillImage(&extra_channels.back());
    134       }
    135       ib.SetExtraChannels(std::move(extra_channels));
    136     }
    137     auto special_frame = std::unique_ptr<BitWriter>(new BitWriter());
    138     FrameInfo dc_frame_info;
    139     dc_frame_info.frame_type = FrameType::kDCFrame;
    140     dc_frame_info.dc_level = frame_header.dc_level + 1;
    141     dc_frame_info.ib_needs_color_transform = false;
    142     dc_frame_info.save_before_color_transform = true;  // Implicitly true
    143     AuxOut dc_aux_out;
    144     JXL_CHECK(EncodeFrame(cparams, dc_frame_info, shared.metadata, ib, cms,
    145                           pool, special_frame.get(),
    146                           aux_out ? &dc_aux_out : nullptr));
    147     if (aux_out) {
    148       for (const auto& l : dc_aux_out.layers) {
    149         aux_out->layers[kLayerDC].Assimilate(l);
    150       }
    151     }
    152     const Span<const uint8_t> encoded = special_frame->GetSpan();
    153     enc_state->special_frames.emplace_back(std::move(special_frame));
    154 
    155     ImageBundle decoded(&shared.metadata->m);
    156     std::unique_ptr<PassesDecoderState> dec_state =
    157         jxl::make_unique<PassesDecoderState>();
    158     JXL_CHECK(
    159         dec_state->output_encoding_info.SetFromMetadata(*shared.metadata));
    160     const uint8_t* frame_start = encoded.data();
    161     size_t encoded_size = encoded.size();
    162     for (int i = 0; i <= cparams.progressive_dc; ++i) {
    163       JXL_CHECK(DecodeFrame(dec_state.get(), pool, frame_start, encoded_size,
    164                             /*frame_header=*/nullptr, &decoded,
    165                             *shared.metadata));
    166       frame_start += decoded.decoded_bytes();
    167       encoded_size -= decoded.decoded_bytes();
    168     }
    169     // TODO(lode): frame_header.dc_level should be equal to
    170     // dec_state.frame_header.dc_level - 1 here, since above we set
    171     // dc_frame_info.dc_level = frame_header.dc_level + 1, and
    172     // dc_frame_info.dc_level is used by EncodeFrame. However, if EncodeFrame
    173     // outputs multiple frames, this assumption could be wrong.
    174     const Image3F& dc_frame =
    175         dec_state->shared->dc_frames[frame_header.dc_level];
    176     JXL_ASSIGN_OR_RETURN(shared.dc_storage,
    177                          Image3F::Create(dc_frame.xsize(), dc_frame.ysize()));
    178     CopyImageTo(dc_frame, &shared.dc_storage);
    179     ZeroFillImage(&shared.quant_dc);
    180     shared.dc = &shared.dc_storage;
    181     JXL_CHECK(encoded_size == 0);
    182   } else {
    183     std::atomic<bool> has_error{false};
    184     auto compute_dc_coeffs = [&](int group_index, int /* thread */) {
    185       if (has_error) return;
    186       const Rect r = enc_state->shared.frame_dim.DCGroupRect(group_index);
    187       int modular_group_index = group_index;
    188       if (enc_state->streaming_mode) {
    189         JXL_ASSERT(group_index == 0);
    190         modular_group_index = enc_state->dc_group_index;
    191       }
    192       if (!modular_frame_encoder->AddVarDCTDC(
    193               frame_header, dc, r, modular_group_index,
    194               enc_state->cparams.speed_tier < SpeedTier::kFalcon, enc_state,
    195               /*jpeg_transcode=*/false)) {
    196         has_error = true;
    197         return;
    198       }
    199     };
    200     JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, shared.frame_dim.num_dc_groups,
    201                                   ThreadPool::NoInit, compute_dc_coeffs,
    202                                   "Compute DC coeffs"));
    203     if (has_error) return JXL_FAILURE("Compute DC coeffs failed");
    204     // TODO(veluca): this is only useful in tests and if inspection is enabled.
    205     if (!(frame_header.flags & FrameHeader::kSkipAdaptiveDCSmoothing)) {
    206       JXL_RETURN_IF_ERROR(AdaptiveDCSmoothing(shared.quantizer.MulDC(),
    207                                               &shared.dc_storage, pool));
    208     }
    209   }
    210   std::atomic<bool> has_error{false};
    211   auto compute_ac_meta = [&](int group_index, int /* thread */) {
    212     const Rect r = enc_state->shared.frame_dim.DCGroupRect(group_index);
    213     int modular_group_index = group_index;
    214     if (enc_state->streaming_mode) {
    215       JXL_ASSERT(group_index == 0);
    216       modular_group_index = enc_state->dc_group_index;
    217     }
    218     if (!modular_frame_encoder->AddACMetadata(r, modular_group_index,
    219                                               /*jpeg_transcode=*/false,
    220                                               enc_state)) {
    221       has_error = true;
    222       return;
    223     }
    224   };
    225   JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, shared.frame_dim.num_dc_groups,
    226                                 ThreadPool::NoInit, compute_ac_meta,
    227                                 "Compute AC Metadata"));
    228   if (has_error) return JXL_FAILURE("Compute AC Metadata failed");
    229   return true;
    230 }
    231 
    232 }  // namespace jxl