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