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 <type_traits>
12 
13 #include "lib/jxl/ac_strategy.h"
14 #include "lib/jxl/aux_out.h"
15 #include "lib/jxl/base/compiler_specific.h"
16 #include "lib/jxl/base/padded_bytes.h"
17 #include "lib/jxl/base/profiler.h"
18 #include "lib/jxl/base/span.h"
19 #include "lib/jxl/color_encoding_internal.h"
20 #include "lib/jxl/common.h"
21 #include "lib/jxl/compressed_dc.h"
22 #include "lib/jxl/dct_scales.h"
23 #include "lib/jxl/dct_util.h"
24 #include "lib/jxl/dec_frame.h"
25 #include "lib/jxl/enc_frame.h"
26 #include "lib/jxl/enc_group.h"
27 #include "lib/jxl/enc_modular.h"
28 #include "lib/jxl/frame_header.h"
29 #include "lib/jxl/image.h"
30 #include "lib/jxl/image_bundle.h"
31 #include "lib/jxl/image_ops.h"
32 #include "lib/jxl/passes_state.h"
33 #include "lib/jxl/quantizer.h"
34 
35 namespace jxl {
36 
InitializePassesEncoder(const Image3F & opsin,ThreadPool * pool,PassesEncoderState * enc_state,ModularFrameEncoder * modular_frame_encoder,AuxOut * aux_out)37 void InitializePassesEncoder(const Image3F& opsin, ThreadPool* pool,
38                              PassesEncoderState* enc_state,
39                              ModularFrameEncoder* modular_frame_encoder,
40                              AuxOut* aux_out) {
41   PROFILER_FUNC;
42 
43   PassesSharedState& JXL_RESTRICT shared = enc_state->shared;
44 
45   enc_state->histogram_idx.resize(shared.frame_dim.num_groups);
46 
47   enc_state->x_qm_multiplier =
48       std::pow(1.25f, shared.frame_header.x_qm_scale - 2.0f);
49   enc_state->b_qm_multiplier =
50       std::pow(1.25f, shared.frame_header.b_qm_scale - 2.0f);
51 
52   if (enc_state->coeffs.size() < shared.frame_header.passes.num_passes) {
53     enc_state->coeffs.reserve(shared.frame_header.passes.num_passes);
54     for (size_t i = enc_state->coeffs.size();
55          i < shared.frame_header.passes.num_passes; i++) {
56       // Allocate enough coefficients for each group on every row.
57       enc_state->coeffs.emplace_back(make_unique<ACImageT<int32_t>>(
58           kGroupDim * kGroupDim, shared.frame_dim.num_groups));
59     }
60   }
61   while (enc_state->coeffs.size() > shared.frame_header.passes.num_passes) {
62     enc_state->coeffs.pop_back();
63   }
64 
65   Image3F dc(shared.frame_dim.xsize_blocks, shared.frame_dim.ysize_blocks);
66   RunOnPool(
67       pool, 0, shared.frame_dim.num_groups, ThreadPool::SkipInit(),
68       [&](size_t group_idx, size_t _) {
69         ComputeCoefficients(group_idx, enc_state, opsin, &dc);
70       },
71       "Compute coeffs");
72 
73   if (shared.frame_header.flags & FrameHeader::kUseDcFrame) {
74     CompressParams cparams = enc_state->cparams;
75     // Guess a distance that produces good initial results.
76     cparams.butteraugli_distance =
77         std::max(kMinButteraugliDistance,
78                  enc_state->cparams.butteraugli_distance * 0.1f);
79     cparams.dots = Override::kOff;
80     cparams.noise = Override::kOff;
81     cparams.patches = Override::kOff;
82     cparams.gaborish = Override::kOff;
83     cparams.epf = 0;
84     cparams.max_error_mode = true;
85     cparams.resampling = 1;
86     cparams.ec_resampling = 1;
87     for (size_t c = 0; c < 3; c++) {
88       cparams.max_error[c] = shared.quantizer.MulDC()[c];
89     }
90     JXL_ASSERT(cparams.progressive_dc > 0);
91     cparams.progressive_dc--;
92     // The DC frame will have alpha=0. Don't erase its contents.
93     cparams.keep_invisible = Override::kOn;
94     // No EPF or Gaborish in DC frames.
95     cparams.epf = 0;
96     cparams.gaborish = Override::kOff;
97     // Use kVarDCT in max_error_mode for intermediate progressive DC,
98     // and kModular for the smallest DC (first in the bitstream)
99     if (cparams.progressive_dc == 0) {
100       cparams.modular_mode = true;
101       cparams.quality_pair.first = cparams.quality_pair.second =
102           99.f - enc_state->cparams.butteraugli_distance * 0.2f;
103     }
104     ImageBundle ib(&shared.metadata->m);
105     // This is a lie - dc is in XYB
106     // (but EncodeFrame will skip RGB->XYB conversion anyway)
107     ib.SetFromImage(
108         std::move(dc),
109         ColorEncoding::LinearSRGB(shared.metadata->m.color_encoding.IsGray()));
110     if (!ib.metadata()->extra_channel_info.empty()) {
111       // Add dummy extra channels to the patch image: dc_level frames do not yet
112       // support extra channels, but the codec expects that the amount of extra
113       // channels in frames matches that in the metadata of the codestream.
114       std::vector<ImageF> extra_channels;
115       extra_channels.reserve(ib.metadata()->extra_channel_info.size());
116       for (size_t i = 0; i < ib.metadata()->extra_channel_info.size(); i++) {
117         extra_channels.emplace_back(ib.xsize(), ib.ysize());
118         // Must initialize the image with data to not affect blending with
119         // uninitialized memory.
120         // TODO(lode): dc_level must copy and use the real extra channels
121         // instead.
122         ZeroFillImage(&extra_channels.back());
123       }
124       ib.SetExtraChannels(std::move(extra_channels));
125     }
126     std::unique_ptr<PassesEncoderState> state =
127         jxl::make_unique<PassesEncoderState>();
128 
129     auto special_frame = std::unique_ptr<BitWriter>(new BitWriter());
130     FrameInfo dc_frame_info;
131     dc_frame_info.frame_type = FrameType::kDCFrame;
132     dc_frame_info.dc_level = shared.frame_header.dc_level + 1;
133     dc_frame_info.ib_needs_color_transform = false;
134     dc_frame_info.save_before_color_transform = true;  // Implicitly true
135     // TODO(lode): the EncodeFrame / DecodeFrame pair here is likely broken in
136     // case of dc_level >= 3, since EncodeFrame may output multiple frames
137     // to the bitwriter, while DecodeFrame reads only one.
138     JXL_CHECK(EncodeFrame(cparams, dc_frame_info, shared.metadata, ib,
139                           state.get(), pool, special_frame.get(), nullptr));
140     const Span<const uint8_t> encoded = special_frame->GetSpan();
141     enc_state->special_frames.emplace_back(std::move(special_frame));
142 
143     BitReader br(encoded);
144     ImageBundle decoded(&shared.metadata->m);
145     std::unique_ptr<PassesDecoderState> dec_state =
146         jxl::make_unique<PassesDecoderState>();
147     JXL_CHECK(dec_state->output_encoding_info.Set(
148         *shared.metadata,
149         ColorEncoding::LinearSRGB(shared.metadata->m.color_encoding.IsGray())));
150     JXL_CHECK(DecodeFrame({}, dec_state.get(), pool, &br, &decoded,
151                           *shared.metadata, /*constraints=*/nullptr));
152     // TODO(lode): shared.frame_header.dc_level should be equal to
153     // dec_state.shared->frame_header.dc_level - 1 here, since above we set
154     // dc_frame_info.dc_level = shared.frame_header.dc_level + 1, and
155     // dc_frame_info.dc_level is used by EncodeFrame. However, if EncodeFrame
156     // outputs multiple frames, this assumption could be wrong.
157     shared.dc_storage =
158         CopyImage(dec_state->shared->dc_frames[shared.frame_header.dc_level]);
159     ZeroFillImage(&shared.quant_dc);
160     shared.dc = &shared.dc_storage;
161     JXL_CHECK(br.Close());
162   } else {
163     auto compute_dc_coeffs = [&](int group_index, int /* thread */) {
164       modular_frame_encoder->AddVarDCTDC(
165           dc, group_index,
166           enc_state->cparams.butteraugli_distance >= 2.0f &&
167               enc_state->cparams.speed_tier < SpeedTier::kFalcon,
168           enc_state);
169     };
170     RunOnPool(pool, 0, shared.frame_dim.num_dc_groups, ThreadPool::SkipInit(),
171               compute_dc_coeffs, "Compute DC coeffs");
172     // TODO(veluca): this is only useful in tests and if inspection is enabled.
173     if (!(shared.frame_header.flags & FrameHeader::kSkipAdaptiveDCSmoothing)) {
174       AdaptiveDCSmoothing(shared.quantizer.MulDC(), &shared.dc_storage, pool);
175     }
176   }
177   auto compute_ac_meta = [&](int group_index, int /* thread */) {
178     modular_frame_encoder->AddACMetadata(group_index, /*jpeg_transcode=*/false,
179                                          enc_state);
180   };
181   RunOnPool(pool, 0, shared.frame_dim.num_dc_groups, ThreadPool::SkipInit(),
182             compute_ac_meta, "Compute AC Metadata");
183 
184   if (aux_out != nullptr) {
185     aux_out->InspectImage3F("compressed_image:InitializeFrameEncCache:dc_dec",
186                             shared.dc_storage);
187   }
188 }
189 
InitOnce()190 void EncCache::InitOnce() {
191   PROFILER_FUNC;
192 
193   if (num_nzeroes.xsize() == 0) {
194     num_nzeroes = Image3I(kGroupDimInBlocks, kGroupDimInBlocks);
195   }
196 }
197 
198 }  // namespace jxl
199