Coverage Report

Created: 2024-05-21 06:24

/src/libjxl/lib/jxl/enc_frame.cc
Line
Count
Source (jump to first uncovered line)
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_frame.h"
7
8
#include <jxl/memory_manager.h>
9
10
#include <algorithm>
11
#include <array>
12
#include <atomic>
13
#include <cmath>
14
#include <cstddef>
15
#include <cstdint>
16
#include <memory>
17
#include <numeric>
18
#include <utility>
19
#include <vector>
20
21
#include "lib/jxl/ac_context.h"
22
#include "lib/jxl/ac_strategy.h"
23
#include "lib/jxl/base/bits.h"
24
#include "lib/jxl/base/common.h"
25
#include "lib/jxl/base/compiler_specific.h"
26
#include "lib/jxl/base/data_parallel.h"
27
#include "lib/jxl/base/override.h"
28
#include "lib/jxl/base/printf_macros.h"
29
#include "lib/jxl/base/rect.h"
30
#include "lib/jxl/base/status.h"
31
#include "lib/jxl/chroma_from_luma.h"
32
#include "lib/jxl/coeff_order.h"
33
#include "lib/jxl/coeff_order_fwd.h"
34
#include "lib/jxl/color_encoding_internal.h"
35
#include "lib/jxl/common.h"  // kMaxNumPasses
36
#include "lib/jxl/dct_util.h"
37
#include "lib/jxl/dec_external_image.h"
38
#include "lib/jxl/enc_ac_strategy.h"
39
#include "lib/jxl/enc_adaptive_quantization.h"
40
#include "lib/jxl/enc_ans.h"
41
#include "lib/jxl/enc_aux_out.h"
42
#include "lib/jxl/enc_bit_writer.h"
43
#include "lib/jxl/enc_cache.h"
44
#include "lib/jxl/enc_chroma_from_luma.h"
45
#include "lib/jxl/enc_coeff_order.h"
46
#include "lib/jxl/enc_context_map.h"
47
#include "lib/jxl/enc_entropy_coder.h"
48
#include "lib/jxl/enc_external_image.h"
49
#include "lib/jxl/enc_fields.h"
50
#include "lib/jxl/enc_group.h"
51
#include "lib/jxl/enc_heuristics.h"
52
#include "lib/jxl/enc_modular.h"
53
#include "lib/jxl/enc_noise.h"
54
#include "lib/jxl/enc_params.h"
55
#include "lib/jxl/enc_patch_dictionary.h"
56
#include "lib/jxl/enc_photon_noise.h"
57
#include "lib/jxl/enc_quant_weights.h"
58
#include "lib/jxl/enc_splines.h"
59
#include "lib/jxl/enc_toc.h"
60
#include "lib/jxl/enc_xyb.h"
61
#include "lib/jxl/fields.h"
62
#include "lib/jxl/frame_dimensions.h"
63
#include "lib/jxl/frame_header.h"
64
#include "lib/jxl/image.h"
65
#include "lib/jxl/image_bundle.h"
66
#include "lib/jxl/image_ops.h"
67
#include "lib/jxl/jpeg/enc_jpeg_data.h"
68
#include "lib/jxl/loop_filter.h"
69
#include "lib/jxl/modular/options.h"
70
#include "lib/jxl/quant_weights.h"
71
#include "lib/jxl/quantizer.h"
72
#include "lib/jxl/splines.h"
73
#include "lib/jxl/toc.h"
74
75
namespace jxl {
76
77
52
Status ParamsPostInit(CompressParams* p) {
78
52
  if (!p->manual_noise.empty() &&
79
52
      p->manual_noise.size() != NoiseParams::kNumNoisePoints) {
80
0
    return JXL_FAILURE("Invalid number of noise lut entries");
81
0
  }
82
52
  if (!p->manual_xyb_factors.empty() && p->manual_xyb_factors.size() != 3) {
83
0
    return JXL_FAILURE("Invalid number of XYB quantization factors");
84
0
  }
85
52
  if (!p->modular_mode && p->butteraugli_distance == 0.0) {
86
0
    p->butteraugli_distance = kMinButteraugliDistance;
87
0
  }
88
52
  if (p->original_butteraugli_distance == -1.0) {
89
52
    p->original_butteraugli_distance = p->butteraugli_distance;
90
52
  }
91
52
  if (p->resampling <= 0) {
92
52
    p->resampling = 1;
93
    // For very low bit rates, using 2x2 resampling gives better results on
94
    // most photographic images, with an adjusted butteraugli score chosen to
95
    // give roughly the same amount of bits per pixel.
96
52
    if (!p->already_downsampled && p->butteraugli_distance >= 20) {
97
0
      p->resampling = 2;
98
0
      p->butteraugli_distance = 6 + ((p->butteraugli_distance - 20) * 0.25);
99
0
    }
100
52
  }
101
52
  if (p->ec_resampling <= 0) {
102
52
    p->ec_resampling = p->resampling;
103
52
  }
104
52
  return true;
105
52
}
106
107
namespace {
108
109
template <typename T>
110
uint32_t GetBitDepth(JxlBitDepth bit_depth, const T& metadata,
111
52
                     JxlPixelFormat format) {
112
52
  if (bit_depth.type == JXL_BIT_DEPTH_FROM_PIXEL_FORMAT) {
113
52
    return BitsPerChannel(format.data_type);
114
52
  } else if (bit_depth.type == JXL_BIT_DEPTH_FROM_CODESTREAM) {
115
0
    return metadata.bit_depth.bits_per_sample;
116
0
  } else if (bit_depth.type == JXL_BIT_DEPTH_CUSTOM) {
117
0
    return bit_depth.bits_per_sample;
118
0
  } else {
119
0
    return 0;
120
0
  }
121
52
}
enc_frame.cc:unsigned int jxl::(anonymous namespace)::GetBitDepth<jxl::ImageMetadata>(JxlBitDepth, jxl::ImageMetadata const&, JxlPixelFormat)
Line
Count
Source
111
52
                     JxlPixelFormat format) {
112
52
  if (bit_depth.type == JXL_BIT_DEPTH_FROM_PIXEL_FORMAT) {
113
52
    return BitsPerChannel(format.data_type);
114
52
  } else if (bit_depth.type == JXL_BIT_DEPTH_FROM_CODESTREAM) {
115
0
    return metadata.bit_depth.bits_per_sample;
116
0
  } else if (bit_depth.type == JXL_BIT_DEPTH_CUSTOM) {
117
0
    return bit_depth.bits_per_sample;
118
0
  } else {
119
0
    return 0;
120
0
  }
121
52
}
Unexecuted instantiation: enc_frame.cc:unsigned int jxl::(anonymous namespace)::GetBitDepth<jxl::ExtraChannelInfo>(JxlBitDepth, jxl::ExtraChannelInfo const&, JxlPixelFormat)
122
123
Status CopyColorChannels(JxlChunkedFrameInputSource input, Rect rect,
124
                         const FrameInfo& frame_info,
125
                         const ImageMetadata& metadata, ThreadPool* pool,
126
                         Image3F* color, ImageF* alpha,
127
52
                         bool* has_interleaved_alpha) {
128
52
  JxlPixelFormat format = {4, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
129
52
  input.get_color_channels_pixel_format(input.opaque, &format);
130
52
  *has_interleaved_alpha = format.num_channels == 2 || format.num_channels == 4;
131
52
  size_t bits_per_sample =
132
52
      GetBitDepth(frame_info.image_bit_depth, metadata, format);
133
52
  size_t row_offset;
134
52
  auto buffer = GetColorBuffer(input, rect.x0(), rect.y0(), rect.xsize(),
135
52
                               rect.ysize(), &row_offset);
136
52
  if (!buffer) {
137
0
    return JXL_FAILURE("no buffer for color channels given");
138
0
  }
139
52
  size_t color_channels = frame_info.ib_needs_color_transform
140
52
                              ? metadata.color_encoding.Channels()
141
52
                              : 3;
142
52
  if (format.num_channels < color_channels) {
143
0
    return JXL_FAILURE("Expected %" PRIuS
144
0
                       " color channels, received only %u channels",
145
0
                       color_channels, format.num_channels);
146
0
  }
147
52
  const uint8_t* data = reinterpret_cast<const uint8_t*>(buffer.get());
148
208
  for (size_t c = 0; c < color_channels; ++c) {
149
156
    JXL_RETURN_IF_ERROR(ConvertFromExternalNoSizeCheck(
150
156
        data, rect.xsize(), rect.ysize(), row_offset, bits_per_sample, format,
151
156
        c, pool, &color->Plane(c)));
152
156
  }
153
52
  if (color_channels == 1) {
154
0
    CopyImageTo(color->Plane(0), &color->Plane(1));
155
0
    CopyImageTo(color->Plane(0), &color->Plane(2));
156
0
  }
157
52
  if (alpha) {
158
0
    if (*has_interleaved_alpha) {
159
0
      JXL_RETURN_IF_ERROR(ConvertFromExternalNoSizeCheck(
160
0
          data, rect.xsize(), rect.ysize(), row_offset, bits_per_sample, format,
161
0
          format.num_channels - 1, pool, alpha));
162
0
    } else {
163
      // if alpha is not passed, but it is expected, then assume
164
      // it is all-opaque
165
0
      FillImage(1.0f, alpha);
166
0
    }
167
0
  }
168
52
  return true;
169
52
}
170
171
Status CopyExtraChannels(JxlChunkedFrameInputSource input, Rect rect,
172
                         const FrameInfo& frame_info,
173
                         const ImageMetadata& metadata,
174
                         bool has_interleaved_alpha, ThreadPool* pool,
175
52
                         std::vector<ImageF>* extra_channels) {
176
52
  for (size_t ec = 0; ec < metadata.num_extra_channels; ec++) {
177
0
    if (has_interleaved_alpha &&
178
0
        metadata.extra_channel_info[ec].type == ExtraChannel::kAlpha) {
179
      // Skip this alpha channel, but still request additional alpha channels
180
      // if they exist.
181
0
      has_interleaved_alpha = false;
182
0
      continue;
183
0
    }
184
0
    JxlPixelFormat ec_format = {1, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
185
0
    input.get_extra_channel_pixel_format(input.opaque, ec, &ec_format);
186
0
    ec_format.num_channels = 1;
187
0
    size_t row_offset;
188
0
    auto buffer =
189
0
        GetExtraChannelBuffer(input, ec, rect.x0(), rect.y0(), rect.xsize(),
190
0
                              rect.ysize(), &row_offset);
191
0
    if (!buffer) {
192
0
      return JXL_FAILURE("no buffer for extra channel given");
193
0
    }
194
0
    size_t bits_per_sample = GetBitDepth(
195
0
        frame_info.image_bit_depth, metadata.extra_channel_info[ec], ec_format);
196
0
    if (!ConvertFromExternalNoSizeCheck(
197
0
            reinterpret_cast<const uint8_t*>(buffer.get()), rect.xsize(),
198
0
            rect.ysize(), row_offset, bits_per_sample, ec_format, 0, pool,
199
0
            &(*extra_channels)[ec])) {
200
0
      return JXL_FAILURE("Failed to set buffer for extra channel");
201
0
    }
202
0
  }
203
52
  return true;
204
52
}
205
206
void SetProgressiveMode(const CompressParams& cparams,
207
52
                        ProgressiveSplitter* progressive_splitter) {
208
52
  constexpr PassDefinition progressive_passes_dc_vlf_lf_full_ac[] = {
209
52
      {/*num_coefficients=*/2, /*shift=*/0,
210
52
       /*suitable_for_downsampling_of_at_least=*/4},
211
52
      {/*num_coefficients=*/3, /*shift=*/0,
212
52
       /*suitable_for_downsampling_of_at_least=*/2},
213
52
      {/*num_coefficients=*/8, /*shift=*/0,
214
52
       /*suitable_for_downsampling_of_at_least=*/0},
215
52
  };
216
52
  constexpr PassDefinition progressive_passes_dc_quant_ac_full_ac[] = {
217
52
      {/*num_coefficients=*/8, /*shift=*/1,
218
52
       /*suitable_for_downsampling_of_at_least=*/2},
219
52
      {/*num_coefficients=*/8, /*shift=*/0,
220
52
       /*suitable_for_downsampling_of_at_least=*/0},
221
52
  };
222
52
  bool progressive_mode = ApplyOverride(cparams.progressive_mode, false);
223
52
  bool qprogressive_mode = ApplyOverride(cparams.qprogressive_mode, false);
224
52
  if (cparams.custom_progressive_mode) {
225
0
    progressive_splitter->SetProgressiveMode(*cparams.custom_progressive_mode);
226
52
  } else if (qprogressive_mode) {
227
0
    progressive_splitter->SetProgressiveMode(
228
0
        ProgressiveMode{progressive_passes_dc_quant_ac_full_ac});
229
52
  } else if (progressive_mode) {
230
0
    progressive_splitter->SetProgressiveMode(
231
0
        ProgressiveMode{progressive_passes_dc_vlf_lf_full_ac});
232
0
  }
233
52
}
234
235
52
uint64_t FrameFlagsFromParams(const CompressParams& cparams) {
236
52
  uint64_t flags = 0;
237
238
52
  const float dist = cparams.butteraugli_distance;
239
240
  // We don't add noise at low butteraugli distances because the original
241
  // noise is stored within the compressed image and adding noise makes things
242
  // worse.
243
52
  if (ApplyOverride(cparams.noise, dist >= kMinButteraugliForNoise) ||
244
52
      cparams.photon_noise_iso > 0 ||
245
52
      cparams.manual_noise.size() == NoiseParams::kNumNoisePoints) {
246
0
    flags |= FrameHeader::kNoise;
247
0
  }
248
249
52
  if (cparams.progressive_dc > 0 && cparams.modular_mode == false) {
250
0
    flags |= FrameHeader::kUseDcFrame;
251
0
  }
252
253
52
  return flags;
254
52
}
255
256
Status LoopFilterFromParams(const CompressParams& cparams, bool streaming_mode,
257
52
                            FrameHeader* JXL_RESTRICT frame_header) {
258
52
  LoopFilter* loop_filter = &frame_header->loop_filter;
259
260
  // Gaborish defaults to enabled in Hare or slower.
261
52
  loop_filter->gab = ApplyOverride(
262
52
      cparams.gaborish, cparams.speed_tier <= SpeedTier::kHare &&
263
52
                            frame_header->encoding == FrameEncoding::kVarDCT &&
264
52
                            cparams.decoding_speed_tier < 4 &&
265
52
                            !cparams.disable_percepeptual_optimizations);
266
267
52
  if (cparams.epf != -1) {
268
0
    loop_filter->epf_iters = cparams.epf;
269
52
  } else if (cparams.disable_percepeptual_optimizations) {
270
0
    loop_filter->epf_iters = 0;
271
0
    return true;
272
52
  } else {
273
52
    if (frame_header->encoding == FrameEncoding::kModular) {
274
0
      loop_filter->epf_iters = 0;
275
52
    } else {
276
52
      constexpr float kThresholds[3] = {0.7, 1.5, 4.0};
277
52
      loop_filter->epf_iters = 0;
278
52
      if (cparams.decoding_speed_tier < 3) {
279
208
        for (size_t i = cparams.decoding_speed_tier == 2 ? 1 : 0; i < 3; i++) {
280
156
          if (cparams.butteraugli_distance >= kThresholds[i]) {
281
52
            loop_filter->epf_iters++;
282
52
          }
283
156
        }
284
52
      }
285
52
    }
286
52
  }
287
  // Strength of EPF in modular mode.
288
52
  if (frame_header->encoding == FrameEncoding::kModular &&
289
52
      !cparams.IsLossless()) {
290
    // TODO(veluca): this formula is nonsense.
291
0
    loop_filter->epf_sigma_for_modular =
292
0
        std::max(cparams.butteraugli_distance, 1.0f);
293
0
  }
294
52
  if (frame_header->encoding == FrameEncoding::kModular &&
295
52
      cparams.lossy_palette) {
296
0
    loop_filter->epf_sigma_for_modular = 1.0f;
297
0
  }
298
299
52
  return true;
300
52
}
301
302
Status MakeFrameHeader(size_t xsize, size_t ysize,
303
                       const CompressParams& cparams,
304
                       const ProgressiveSplitter& progressive_splitter,
305
                       const FrameInfo& frame_info,
306
                       const jpeg::JPEGData* jpeg_data, bool streaming_mode,
307
52
                       FrameHeader* JXL_RESTRICT frame_header) {
308
52
  frame_header->nonserialized_is_preview = frame_info.is_preview;
309
52
  frame_header->is_last = frame_info.is_last;
310
52
  frame_header->save_before_color_transform =
311
52
      frame_info.save_before_color_transform;
312
52
  frame_header->frame_type = frame_info.frame_type;
313
52
  frame_header->name = frame_info.name;
314
315
52
  progressive_splitter.InitPasses(&frame_header->passes);
316
317
52
  if (cparams.modular_mode) {
318
0
    frame_header->encoding = FrameEncoding::kModular;
319
0
    if (cparams.modular_group_size_shift == -1) {
320
0
      frame_header->group_size_shift = 1;
321
      // no point using groups when only one group is full and the others are
322
      // less than half full: multithreading will not really help much, while
323
      // compression does suffer
324
0
      if (xsize <= 400 && ysize <= 400) {
325
0
        frame_header->group_size_shift = 2;
326
0
      }
327
0
    } else {
328
0
      frame_header->group_size_shift = cparams.modular_group_size_shift;
329
0
    }
330
0
  }
331
332
52
  if (jpeg_data) {
333
    // we are transcoding a JPEG, so we don't get to choose
334
0
    frame_header->encoding = FrameEncoding::kVarDCT;
335
0
    frame_header->x_qm_scale = 2;
336
0
    frame_header->b_qm_scale = 2;
337
0
    JXL_RETURN_IF_ERROR(SetChromaSubsamplingFromJpegData(
338
0
        *jpeg_data, &frame_header->chroma_subsampling));
339
0
    JXL_RETURN_IF_ERROR(SetColorTransformFromJpegData(
340
0
        *jpeg_data, &frame_header->color_transform));
341
52
  } else {
342
52
    frame_header->color_transform = cparams.color_transform;
343
52
    if (!cparams.modular_mode &&
344
52
        (frame_header->chroma_subsampling.MaxHShift() != 0 ||
345
52
         frame_header->chroma_subsampling.MaxVShift() != 0)) {
346
0
      return JXL_FAILURE(
347
0
          "Chroma subsampling is not supported in VarDCT mode when not "
348
0
          "recompressing JPEGs");
349
0
    }
350
52
  }
351
52
  if (frame_header->color_transform != ColorTransform::kYCbCr &&
352
52
      (frame_header->chroma_subsampling.MaxHShift() != 0 ||
353
52
       frame_header->chroma_subsampling.MaxVShift() != 0)) {
354
0
    return JXL_FAILURE(
355
0
        "Chroma subsampling is not supported when color transform is not "
356
0
        "YCbCr");
357
0
  }
358
359
52
  frame_header->flags = FrameFlagsFromParams(cparams);
360
  // Non-photon noise is not supported in the Modular encoder for now.
361
52
  if (frame_header->encoding != FrameEncoding::kVarDCT &&
362
52
      cparams.photon_noise_iso == 0 && cparams.manual_noise.empty()) {
363
0
    frame_header->UpdateFlag(false, FrameHeader::Flags::kNoise);
364
0
  }
365
366
52
  JXL_RETURN_IF_ERROR(
367
52
      LoopFilterFromParams(cparams, streaming_mode, frame_header));
368
369
52
  frame_header->dc_level = frame_info.dc_level;
370
52
  if (frame_header->dc_level > 2) {
371
    // With 3 or more progressive_dc frames, the implementation does not yet
372
    // work, see enc_cache.cc.
373
0
    return JXL_FAILURE("progressive_dc > 2 is not yet supported");
374
0
  }
375
52
  if (cparams.progressive_dc > 0 &&
376
52
      (cparams.ec_resampling != 1 || cparams.resampling != 1)) {
377
0
    return JXL_FAILURE("Resampling not supported with DC frames");
378
0
  }
379
52
  if (cparams.resampling != 1 && cparams.resampling != 2 &&
380
52
      cparams.resampling != 4 && cparams.resampling != 8) {
381
0
    return JXL_FAILURE("Invalid resampling factor");
382
0
  }
383
52
  if (cparams.ec_resampling != 1 && cparams.ec_resampling != 2 &&
384
52
      cparams.ec_resampling != 4 && cparams.ec_resampling != 8) {
385
0
    return JXL_FAILURE("Invalid ec_resampling factor");
386
0
  }
387
  // Resized frames.
388
52
  if (frame_info.frame_type != FrameType::kDCFrame) {
389
52
    frame_header->frame_origin = frame_info.origin;
390
52
    size_t ups = 1;
391
52
    if (cparams.already_downsampled) ups = cparams.resampling;
392
393
    // TODO(lode): this is not correct in case of odd original image sizes in
394
    // combination with cparams.already_downsampled. Likely these values should
395
    // be set to respectively frame_header->default_xsize() and
396
    // frame_header->default_ysize() instead, the original (non downsampled)
397
    // intended decoded image dimensions. But it may be more subtle than that
398
    // if combined with crop. This issue causes custom_size_or_origin to be
399
    // incorrectly set to true in case of already_downsampled with odd output
400
    // image size when no cropping is used.
401
52
    frame_header->frame_size.xsize = xsize * ups;
402
52
    frame_header->frame_size.ysize = ysize * ups;
403
52
    if (frame_info.origin.x0 != 0 || frame_info.origin.y0 != 0 ||
404
52
        frame_header->frame_size.xsize != frame_header->default_xsize() ||
405
52
        frame_header->frame_size.ysize != frame_header->default_ysize()) {
406
0
      frame_header->custom_size_or_origin = true;
407
0
    }
408
52
  }
409
  // Upsampling.
410
52
  frame_header->upsampling = cparams.resampling;
411
52
  const std::vector<ExtraChannelInfo>& extra_channels =
412
52
      frame_header->nonserialized_metadata->m.extra_channel_info;
413
52
  frame_header->extra_channel_upsampling.clear();
414
52
  frame_header->extra_channel_upsampling.resize(extra_channels.size(),
415
52
                                                cparams.ec_resampling);
416
52
  frame_header->save_as_reference = frame_info.save_as_reference;
417
418
  // Set blending-related information.
419
52
  if (frame_info.blend || frame_header->custom_size_or_origin) {
420
    // Set blend_channel to the first alpha channel. These values are only
421
    // encoded in case a blend mode involving alpha is used and there are more
422
    // than one extra channels.
423
0
    size_t index = 0;
424
0
    if (frame_info.alpha_channel == -1) {
425
0
      if (extra_channels.size() > 1) {
426
0
        for (size_t i = 0; i < extra_channels.size(); i++) {
427
0
          if (extra_channels[i].type == ExtraChannel::kAlpha) {
428
0
            index = i;
429
0
            break;
430
0
          }
431
0
        }
432
0
      }
433
0
    } else {
434
0
      index = static_cast<size_t>(frame_info.alpha_channel);
435
0
      JXL_ASSERT(index == 0 || index < extra_channels.size());
436
0
    }
437
0
    frame_header->blending_info.alpha_channel = index;
438
0
    frame_header->blending_info.mode =
439
0
        frame_info.blend ? frame_info.blendmode : BlendMode::kReplace;
440
0
    frame_header->blending_info.source = frame_info.source;
441
0
    frame_header->blending_info.clamp = frame_info.clamp;
442
0
    const auto& extra_channel_info = frame_info.extra_channel_blending_info;
443
0
    for (size_t i = 0; i < extra_channels.size(); i++) {
444
0
      if (i < extra_channel_info.size()) {
445
0
        frame_header->extra_channel_blending_info[i] = extra_channel_info[i];
446
0
      } else {
447
0
        frame_header->extra_channel_blending_info[i].alpha_channel = index;
448
0
        BlendMode default_blend = frame_info.blendmode;
449
0
        if (extra_channels[i].type != ExtraChannel::kBlack && i != index) {
450
          // K needs to be blended, spot colors and other stuff gets added
451
0
          default_blend = BlendMode::kAdd;
452
0
        }
453
0
        frame_header->extra_channel_blending_info[i].mode =
454
0
            frame_info.blend ? default_blend : BlendMode::kReplace;
455
0
        frame_header->extra_channel_blending_info[i].source = 1;
456
0
      }
457
0
    }
458
0
  }
459
460
52
  frame_header->animation_frame.duration = frame_info.duration;
461
52
  frame_header->animation_frame.timecode = frame_info.timecode;
462
463
52
  if (jpeg_data) {
464
0
    frame_header->UpdateFlag(false, FrameHeader::kUseDcFrame);
465
0
    frame_header->UpdateFlag(true, FrameHeader::kSkipAdaptiveDCSmoothing);
466
0
  }
467
468
52
  return true;
469
52
}
470
471
// Invisible (alpha = 0) pixels tend to be a mess in optimized PNGs.
472
// Since they have no visual impact whatsoever, we can replace them with
473
// something that compresses better and reduces artifacts near the edges. This
474
// does some kind of smooth stuff that seems to work.
475
// Replace invisible pixels with a weighted average of the pixel to the left,
476
// the pixel to the topright, and non-invisible neighbours.
477
// Produces downward-blurry smears, with in the upwards direction only a 1px
478
// edge duplication but not more. It would probably be better to smear in all
479
// directions. That requires an alpha-weighed convolution with a large enough
480
// kernel though, which might be overkill...
481
0
void SimplifyInvisible(Image3F* image, const ImageF& alpha, bool lossless) {
482
0
  for (size_t c = 0; c < 3; ++c) {
483
0
    for (size_t y = 0; y < image->ysize(); ++y) {
484
0
      float* JXL_RESTRICT row = image->PlaneRow(c, y);
485
0
      const float* JXL_RESTRICT prow =
486
0
          (y > 0 ? image->PlaneRow(c, y - 1) : nullptr);
487
0
      const float* JXL_RESTRICT nrow =
488
0
          (y + 1 < image->ysize() ? image->PlaneRow(c, y + 1) : nullptr);
489
0
      const float* JXL_RESTRICT a = alpha.Row(y);
490
0
      const float* JXL_RESTRICT pa = (y > 0 ? alpha.Row(y - 1) : nullptr);
491
0
      const float* JXL_RESTRICT na =
492
0
          (y + 1 < image->ysize() ? alpha.Row(y + 1) : nullptr);
493
0
      for (size_t x = 0; x < image->xsize(); ++x) {
494
0
        if (a[x] == 0) {
495
0
          if (lossless) {
496
0
            row[x] = 0;
497
0
            continue;
498
0
          }
499
0
          float d = 0.f;
500
0
          row[x] = 0;
501
0
          if (x > 0) {
502
0
            row[x] += row[x - 1];
503
0
            d++;
504
0
            if (a[x - 1] > 0.f) {
505
0
              row[x] += row[x - 1];
506
0
              d++;
507
0
            }
508
0
          }
509
0
          if (x + 1 < image->xsize()) {
510
0
            if (y > 0) {
511
0
              row[x] += prow[x + 1];
512
0
              d++;
513
0
            }
514
0
            if (a[x + 1] > 0.f) {
515
0
              row[x] += 2.f * row[x + 1];
516
0
              d += 2.f;
517
0
            }
518
0
            if (y > 0 && pa[x + 1] > 0.f) {
519
0
              row[x] += 2.f * prow[x + 1];
520
0
              d += 2.f;
521
0
            }
522
0
            if (y + 1 < image->ysize() && na[x + 1] > 0.f) {
523
0
              row[x] += 2.f * nrow[x + 1];
524
0
              d += 2.f;
525
0
            }
526
0
          }
527
0
          if (y > 0 && pa[x] > 0.f) {
528
0
            row[x] += 2.f * prow[x];
529
0
            d += 2.f;
530
0
          }
531
0
          if (y + 1 < image->ysize() && na[x] > 0.f) {
532
0
            row[x] += 2.f * nrow[x];
533
0
            d += 2.f;
534
0
          }
535
0
          if (d > 1.f) row[x] /= d;
536
0
        }
537
0
      }
538
0
    }
539
0
  }
540
0
}
541
542
struct PixelStatsForChromacityAdjustment {
543
  float dx = 0;
544
  float db = 0;
545
  float exposed_blue = 0;
546
52
  static float CalcPlane(const ImageF* JXL_RESTRICT plane, const Rect& rect) {
547
52
    float xmax = 0;
548
52
    float ymax = 0;
549
416
    for (size_t ty = 1; ty < rect.ysize(); ++ty) {
550
2.91k
      for (size_t tx = 1; tx < rect.xsize(); ++tx) {
551
2.54k
        float cur = rect.Row(plane, ty)[tx];
552
2.54k
        float prev_row = rect.Row(plane, ty - 1)[tx];
553
2.54k
        float prev = rect.Row(plane, ty)[tx - 1];
554
2.54k
        xmax = std::max(xmax, std::abs(cur - prev));
555
2.54k
        ymax = std::max(ymax, std::abs(cur - prev_row));
556
2.54k
      }
557
364
    }
558
52
    return std::max(xmax, ymax);
559
52
  }
560
  void CalcExposedBlue(const ImageF* JXL_RESTRICT plane_y,
561
52
                       const ImageF* JXL_RESTRICT plane_b, const Rect& rect) {
562
52
    float eb = 0;
563
52
    float xmax = 0;
564
52
    float ymax = 0;
565
416
    for (size_t ty = 1; ty < rect.ysize(); ++ty) {
566
2.91k
      for (size_t tx = 1; tx < rect.xsize(); ++tx) {
567
2.54k
        float cur_y = rect.Row(plane_y, ty)[tx];
568
2.54k
        float cur_b = rect.Row(plane_b, ty)[tx];
569
2.54k
        float exposed_b = cur_b - cur_y * 1.2;
570
2.54k
        float diff_b = cur_b - cur_y;
571
2.54k
        float prev_row = rect.Row(plane_b, ty - 1)[tx];
572
2.54k
        float prev = rect.Row(plane_b, ty)[tx - 1];
573
2.54k
        float diff_prev_row = prev_row - rect.Row(plane_y, ty - 1)[tx];
574
2.54k
        float diff_prev = prev - rect.Row(plane_y, ty)[tx - 1];
575
2.54k
        xmax = std::max(xmax, std::abs(diff_b - diff_prev));
576
2.54k
        ymax = std::max(ymax, std::abs(diff_b - diff_prev_row));
577
2.54k
        if (exposed_b >= 0) {
578
0
          exposed_b *= fabs(cur_b - prev) + fabs(cur_b - prev_row);
579
0
          eb = std::max(eb, exposed_b);
580
0
        }
581
2.54k
      }
582
364
    }
583
52
    exposed_blue = eb;
584
52
    db = std::max(xmax, ymax);
585
52
  }
586
52
  void Calc(const Image3F* JXL_RESTRICT opsin, const Rect& rect) {
587
52
    dx = CalcPlane(&opsin->Plane(0), rect);
588
52
    CalcExposedBlue(&opsin->Plane(1), &opsin->Plane(2), rect);
589
52
  }
590
52
  int HowMuchIsXChannelPixelized() const {
591
52
    if (dx >= 0.026) {
592
0
      return 3;
593
0
    }
594
52
    if (dx >= 0.022) {
595
0
      return 2;
596
0
    }
597
52
    if (dx >= 0.015) {
598
0
      return 1;
599
0
    }
600
52
    return 0;
601
52
  }
602
52
  int HowMuchIsBChannelPixelized() const {
603
52
    int add = exposed_blue >= 0.13 ? 1 : 0;
604
52
    if (db > 0.38) {
605
0
      return 2 + add;
606
0
    }
607
52
    if (db > 0.33) {
608
0
      return 1 + add;
609
0
    }
610
52
    if (db > 0.28) {
611
0
      return add;
612
0
    }
613
52
    return 0;
614
52
  }
615
};
616
617
void ComputeChromacityAdjustments(const CompressParams& cparams,
618
                                  const Image3F& opsin, const Rect& rect,
619
52
                                  FrameHeader* frame_header) {
620
52
  if (frame_header->encoding != FrameEncoding::kVarDCT ||
621
52
      cparams.max_error_mode) {
622
0
    return;
623
0
  }
624
  // 1) Distance based approach for chromacity adjustment:
625
52
  float x_qm_scale_steps[3] = {2.5f, 5.5f, 9.5f};
626
52
  frame_header->x_qm_scale = 3;
627
156
  for (float x_qm_scale_step : x_qm_scale_steps) {
628
156
    if (cparams.original_butteraugli_distance > x_qm_scale_step) {
629
0
      frame_header->x_qm_scale++;
630
0
    }
631
156
  }
632
  // 2) Pixel-based approach for chromacity adjustment:
633
  // look at the individual pixels and make a guess how difficult
634
  // the image would be based on the worst case pixel.
635
52
  PixelStatsForChromacityAdjustment pixel_stats;
636
52
  if (cparams.speed_tier <= SpeedTier::kSquirrel) {
637
52
    pixel_stats.Calc(&opsin, rect);
638
52
  }
639
  // For X take the most severe adjustment.
640
52
  frame_header->x_qm_scale = std::max<int>(
641
52
      frame_header->x_qm_scale, 2 + pixel_stats.HowMuchIsXChannelPixelized());
642
  // B only adjusted by pixel-based approach.
643
52
  frame_header->b_qm_scale = 2 + pixel_stats.HowMuchIsBChannelPixelized();
644
52
}
645
646
void ComputeNoiseParams(const CompressParams& cparams, bool streaming_mode,
647
                        bool color_is_jpeg, const Image3F& opsin,
648
                        const FrameDimensions& frame_dim,
649
52
                        FrameHeader* frame_header, NoiseParams* noise_params) {
650
52
  if (cparams.photon_noise_iso > 0) {
651
0
    *noise_params = SimulatePhotonNoise(frame_dim.xsize, frame_dim.ysize,
652
0
                                        cparams.photon_noise_iso);
653
52
  } else if (cparams.manual_noise.size() == NoiseParams::kNumNoisePoints) {
654
0
    for (size_t i = 0; i < NoiseParams::kNumNoisePoints; i++) {
655
0
      noise_params->lut[i] = cparams.manual_noise[i];
656
0
    }
657
52
  } else if (frame_header->encoding == FrameEncoding::kVarDCT &&
658
52
             frame_header->flags & FrameHeader::kNoise && !color_is_jpeg &&
659
52
             !streaming_mode) {
660
    // Don't start at zero amplitude since adding noise is expensive -- it
661
    // significantly slows down decoding, and this is unlikely to
662
    // completely go away even with advanced optimizations. After the
663
    // kNoiseModelingRampUpDistanceRange we have reached the full level,
664
    // i.e. noise is no longer represented by the compressed image, so we
665
    // can add full noise by the noise modeling itself.
666
0
    static const float kNoiseModelingRampUpDistanceRange = 0.6;
667
0
    static const float kNoiseLevelAtStartOfRampUp = 0.25;
668
0
    static const float kNoiseRampupStart = 1.0;
669
    // TODO(user) test and properly select quality_coef with smooth
670
    // filter
671
0
    float quality_coef = 1.0f;
672
0
    const float rampup = (cparams.butteraugli_distance - kNoiseRampupStart) /
673
0
                         kNoiseModelingRampUpDistanceRange;
674
0
    if (rampup < 1.0f) {
675
0
      quality_coef = kNoiseLevelAtStartOfRampUp +
676
0
                     (1.0f - kNoiseLevelAtStartOfRampUp) * rampup;
677
0
    }
678
0
    if (rampup < 0.0f) {
679
0
      quality_coef = kNoiseRampupStart;
680
0
    }
681
0
    if (!GetNoiseParameter(opsin, noise_params, quality_coef)) {
682
0
      frame_header->flags &= ~FrameHeader::kNoise;
683
0
    }
684
0
  }
685
52
}
686
687
Status DownsampleColorChannels(const CompressParams& cparams,
688
                               const FrameHeader& frame_header,
689
52
                               bool color_is_jpeg, Image3F* opsin) {
690
52
  if (color_is_jpeg || frame_header.upsampling == 1 ||
691
52
      cparams.already_downsampled) {
692
52
    return true;
693
52
  }
694
0
  if (frame_header.encoding == FrameEncoding::kVarDCT &&
695
0
      frame_header.upsampling == 2) {
696
    // TODO(lode): use the regular DownsampleImage, or adapt to the custom
697
    // coefficients, if there is are custom upscaling coefficients in
698
    // CustomTransformData
699
0
    if (cparams.speed_tier <= SpeedTier::kSquirrel) {
700
      // TODO(lode): DownsampleImage2_Iterative is currently too slow to
701
      // be used for squirrel, make it faster, and / or enable it only for
702
      // kitten.
703
0
      JXL_RETURN_IF_ERROR(DownsampleImage2_Iterative(opsin));
704
0
    } else {
705
0
      JXL_RETURN_IF_ERROR(DownsampleImage2_Sharper(opsin));
706
0
    }
707
0
  } else {
708
0
    JXL_ASSIGN_OR_RETURN(*opsin,
709
0
                         DownsampleImage(*opsin, frame_header.upsampling));
710
0
  }
711
0
  if (frame_header.encoding == FrameEncoding::kVarDCT) {
712
0
    PadImageToBlockMultipleInPlace(opsin);
713
0
  }
714
0
  return true;
715
0
}
716
717
template <typename V, typename R>
718
0
void FindIndexOfSumMaximum(const V* array, const size_t len, R* idx, V* sum) {
719
0
  JXL_ASSERT(len > 0);
720
0
  V maxval = 0;
721
0
  V val = 0;
722
0
  R maxidx = 0;
723
0
  for (size_t i = 0; i < len; ++i) {
724
0
    val += array[i];
725
0
    if (val > maxval) {
726
0
      maxval = val;
727
0
      maxidx = i;
728
0
    }
729
0
  }
730
0
  *idx = maxidx;
731
0
  *sum = maxval;
732
0
}
733
734
Status ComputeJPEGTranscodingData(const jpeg::JPEGData& jpeg_data,
735
                                  const FrameHeader& frame_header,
736
                                  ThreadPool* pool,
737
                                  ModularFrameEncoder* enc_modular,
738
0
                                  PassesEncoderState* enc_state) {
739
0
  PassesSharedState& shared = enc_state->shared;
740
0
  JxlMemoryManager* memory_manager = enc_state->memory_manager();
741
0
  const FrameDimensions& frame_dim = shared.frame_dim;
742
743
0
  const size_t xsize = frame_dim.xsize_padded;
744
0
  const size_t ysize = frame_dim.ysize_padded;
745
0
  const size_t xsize_blocks = frame_dim.xsize_blocks;
746
0
  const size_t ysize_blocks = frame_dim.ysize_blocks;
747
748
  // no-op chroma from luma
749
0
  JXL_ASSIGN_OR_RETURN(shared.cmap, ColorCorrelationMap::Create(
750
0
                                        memory_manager, xsize, ysize, false));
751
0
  shared.ac_strategy.FillDCT8();
752
0
  FillImage(static_cast<uint8_t>(0), &shared.epf_sharpness);
753
754
0
  enc_state->coeffs.clear();
755
0
  while (enc_state->coeffs.size() < enc_state->passes.size()) {
756
0
    JXL_ASSIGN_OR_RETURN(
757
0
        std::unique_ptr<ACImageT<int32_t>> coeffs,
758
0
        ACImageT<int32_t>::Make(memory_manager, kGroupDim * kGroupDim,
759
0
                                frame_dim.num_groups));
760
0
    enc_state->coeffs.emplace_back(std::move(coeffs));
761
0
  }
762
763
  // convert JPEG quantization table to a Quantizer object
764
0
  float dcquantization[3];
765
0
  std::vector<QuantEncoding> qe(DequantMatrices::kNum,
766
0
                                QuantEncoding::Library(0));
767
768
0
  auto jpeg_c_map =
769
0
      JpegOrder(frame_header.color_transform, jpeg_data.components.size() == 1);
770
771
0
  std::vector<int> qt(192);
772
0
  for (size_t c = 0; c < 3; c++) {
773
0
    size_t jpeg_c = jpeg_c_map[c];
774
0
    const int32_t* quant =
775
0
        jpeg_data.quant[jpeg_data.components[jpeg_c].quant_idx].values.data();
776
777
0
    dcquantization[c] = 255 * 8.0f / quant[0];
778
0
    for (size_t y = 0; y < 8; y++) {
779
0
      for (size_t x = 0; x < 8; x++) {
780
        // JPEG XL transposes the DCT, JPEG doesn't.
781
0
        qt[c * 64 + 8 * x + y] = quant[8 * y + x];
782
0
      }
783
0
    }
784
0
  }
785
0
  DequantMatricesSetCustomDC(memory_manager, &shared.matrices, dcquantization);
786
0
  float dcquantization_r[3] = {1.0f / dcquantization[0],
787
0
                               1.0f / dcquantization[1],
788
0
                               1.0f / dcquantization[2]};
789
790
0
  qe[AcStrategy::Type::DCT] = QuantEncoding::RAW(qt);
791
0
  JXL_RETURN_IF_ERROR(
792
0
      DequantMatricesSetCustom(&shared.matrices, qe, enc_modular));
793
794
  // Ensure that InvGlobalScale() is 1.
795
0
  shared.quantizer = Quantizer(&shared.matrices, 1, kGlobalScaleDenom);
796
  // Recompute MulDC() and InvMulDC().
797
0
  shared.quantizer.RecomputeFromGlobalScale();
798
799
  // Per-block dequant scaling should be 1.
800
0
  FillImage(static_cast<int32_t>(shared.quantizer.InvGlobalScale()),
801
0
            &shared.raw_quant_field);
802
803
0
  std::vector<int32_t> scaled_qtable(192);
804
0
  for (size_t c = 0; c < 3; c++) {
805
0
    for (size_t i = 0; i < 64; i++) {
806
0
      scaled_qtable[64 * c + i] =
807
0
          (1 << kCFLFixedPointPrecision) * qt[64 + i] / qt[64 * c + i];
808
0
    }
809
0
  }
810
811
0
  auto jpeg_row = [&](size_t c, size_t y) {
812
0
    return jpeg_data.components[jpeg_c_map[c]].coeffs.data() +
813
0
           jpeg_data.components[jpeg_c_map[c]].width_in_blocks * kDCTBlockSize *
814
0
               y;
815
0
  };
816
817
0
  bool DCzero = (frame_header.color_transform == ColorTransform::kYCbCr);
818
  // Compute chroma-from-luma for AC (doesn't seem to be useful for DC)
819
0
  if (frame_header.chroma_subsampling.Is444() &&
820
0
      enc_state->cparams.force_cfl_jpeg_recompression &&
821
0
      jpeg_data.components.size() == 3) {
822
0
    for (size_t c : {0, 2}) {
823
0
      ImageSB* map = (c == 0 ? &shared.cmap.ytox_map : &shared.cmap.ytob_map);
824
0
      const float kScale = kDefaultColorFactor;
825
0
      const int kOffset = 127;
826
0
      const float kBase = c == 0 ? shared.cmap.base().YtoXRatio(0)
827
0
                                 : shared.cmap.base().YtoBRatio(0);
828
0
      const float kZeroThresh =
829
0
          kScale * kZeroBiasDefault[c] *
830
0
          0.9999f;  // just epsilon less for better rounding
831
832
0
      auto process_row = [&](const uint32_t task, const size_t thread) {
833
0
        size_t ty = task;
834
0
        int8_t* JXL_RESTRICT row_out = map->Row(ty);
835
0
        for (size_t tx = 0; tx < map->xsize(); ++tx) {
836
0
          const size_t y0 = ty * kColorTileDimInBlocks;
837
0
          const size_t x0 = tx * kColorTileDimInBlocks;
838
0
          const size_t y1 = std::min(frame_dim.ysize_blocks,
839
0
                                     (ty + 1) * kColorTileDimInBlocks);
840
0
          const size_t x1 = std::min(frame_dim.xsize_blocks,
841
0
                                     (tx + 1) * kColorTileDimInBlocks);
842
0
          int32_t d_num_zeros[257] = {0};
843
          // TODO(veluca): this needs SIMD + fixed point adaptation, and/or
844
          // conversion to the new CfL algorithm.
845
0
          for (size_t y = y0; y < y1; ++y) {
846
0
            const int16_t* JXL_RESTRICT row_m = jpeg_row(1, y);
847
0
            const int16_t* JXL_RESTRICT row_s = jpeg_row(c, y);
848
0
            for (size_t x = x0; x < x1; ++x) {
849
0
              for (size_t coeffpos = 1; coeffpos < kDCTBlockSize; coeffpos++) {
850
0
                const float scaled_m = row_m[x * kDCTBlockSize + coeffpos] *
851
0
                                       scaled_qtable[64 * c + coeffpos] *
852
0
                                       (1.0f / (1 << kCFLFixedPointPrecision));
853
0
                const float scaled_s =
854
0
                    kScale * row_s[x * kDCTBlockSize + coeffpos] +
855
0
                    (kOffset - kBase * kScale) * scaled_m;
856
0
                if (std::abs(scaled_m) > 1e-8f) {
857
0
                  float from;
858
0
                  float to;
859
0
                  if (scaled_m > 0) {
860
0
                    from = (scaled_s - kZeroThresh) / scaled_m;
861
0
                    to = (scaled_s + kZeroThresh) / scaled_m;
862
0
                  } else {
863
0
                    from = (scaled_s + kZeroThresh) / scaled_m;
864
0
                    to = (scaled_s - kZeroThresh) / scaled_m;
865
0
                  }
866
0
                  if (from < 0.0f) {
867
0
                    from = 0.0f;
868
0
                  }
869
0
                  if (to > 255.0f) {
870
0
                    to = 255.0f;
871
0
                  }
872
                  // Instead of clamping the both values
873
                  // we just check that range is sane.
874
0
                  if (from <= to) {
875
0
                    d_num_zeros[static_cast<int>(std::ceil(from))]++;
876
0
                    d_num_zeros[static_cast<int>(std::floor(to + 1))]--;
877
0
                  }
878
0
                }
879
0
              }
880
0
            }
881
0
          }
882
0
          int best = 0;
883
0
          int32_t best_sum = 0;
884
0
          FindIndexOfSumMaximum(d_num_zeros, 256, &best, &best_sum);
885
0
          int32_t offset_sum = 0;
886
0
          for (int i = 0; i < 256; ++i) {
887
0
            if (i <= kOffset) {
888
0
              offset_sum += d_num_zeros[i];
889
0
            }
890
0
          }
891
0
          row_out[tx] = 0;
892
0
          if (best_sum > offset_sum + 1) {
893
0
            row_out[tx] = best - kOffset;
894
0
          }
895
0
        }
896
0
      };
897
898
0
      JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, map->ysize(), ThreadPool::NoInit,
899
0
                                    process_row, "FindCorrelation"));
900
0
    }
901
0
  }
902
903
0
  JXL_ASSIGN_OR_RETURN(
904
0
      Image3F dc, Image3F::Create(memory_manager, xsize_blocks, ysize_blocks));
905
0
  if (!frame_header.chroma_subsampling.Is444()) {
906
0
    ZeroFillImage(&dc);
907
0
    for (auto& coeff : enc_state->coeffs) {
908
0
      coeff->ZeroFill();
909
0
    }
910
0
  }
911
  // JPEG DC is from -1024 to 1023.
912
0
  std::vector<size_t> dc_counts[3] = {};
913
0
  dc_counts[0].resize(2048);
914
0
  dc_counts[1].resize(2048);
915
0
  dc_counts[2].resize(2048);
916
0
  size_t total_dc[3] = {};
917
0
  for (size_t c : {1, 0, 2}) {
918
0
    if (jpeg_data.components.size() == 1 && c != 1) {
919
0
      for (auto& coeff : enc_state->coeffs) {
920
0
        coeff->ZeroFillPlane(c);
921
0
      }
922
0
      ZeroFillImage(&dc.Plane(c));
923
      // Ensure no division by 0.
924
0
      dc_counts[c][1024] = 1;
925
0
      total_dc[c] = 1;
926
0
      continue;
927
0
    }
928
0
    size_t hshift = frame_header.chroma_subsampling.HShift(c);
929
0
    size_t vshift = frame_header.chroma_subsampling.VShift(c);
930
0
    ImageSB& map = (c == 0 ? shared.cmap.ytox_map : shared.cmap.ytob_map);
931
0
    for (size_t group_index = 0; group_index < frame_dim.num_groups;
932
0
         group_index++) {
933
0
      const size_t gx = group_index % frame_dim.xsize_groups;
934
0
      const size_t gy = group_index / frame_dim.xsize_groups;
935
0
      int32_t* coeffs[kMaxNumPasses];
936
0
      for (size_t i = 0; i < enc_state->coeffs.size(); i++) {
937
0
        coeffs[i] = enc_state->coeffs[i]->PlaneRow(c, group_index, 0).ptr32;
938
0
      }
939
0
      int32_t block[64];
940
0
      for (size_t by = gy * kGroupDimInBlocks;
941
0
           by < ysize_blocks && by < (gy + 1) * kGroupDimInBlocks; ++by) {
942
0
        if ((by >> vshift) << vshift != by) continue;
943
0
        const int16_t* JXL_RESTRICT inputjpeg = jpeg_row(c, by >> vshift);
944
0
        const int16_t* JXL_RESTRICT inputjpegY = jpeg_row(1, by);
945
0
        float* JXL_RESTRICT fdc = dc.PlaneRow(c, by >> vshift);
946
0
        const int8_t* JXL_RESTRICT cm =
947
0
            map.ConstRow(by / kColorTileDimInBlocks);
948
0
        for (size_t bx = gx * kGroupDimInBlocks;
949
0
             bx < xsize_blocks && bx < (gx + 1) * kGroupDimInBlocks; ++bx) {
950
0
          if ((bx >> hshift) << hshift != bx) continue;
951
0
          size_t base = (bx >> hshift) * kDCTBlockSize;
952
0
          int idc;
953
0
          if (DCzero) {
954
0
            idc = inputjpeg[base];
955
0
          } else {
956
0
            idc = inputjpeg[base] + 1024 / qt[c * 64];
957
0
          }
958
0
          dc_counts[c][std::min(static_cast<uint32_t>(idc + 1024),
959
0
                                static_cast<uint32_t>(2047))]++;
960
0
          total_dc[c]++;
961
0
          fdc[bx >> hshift] = idc * dcquantization_r[c];
962
0
          if (c == 1 || !enc_state->cparams.force_cfl_jpeg_recompression ||
963
0
              !frame_header.chroma_subsampling.Is444()) {
964
0
            for (size_t y = 0; y < 8; y++) {
965
0
              for (size_t x = 0; x < 8; x++) {
966
0
                block[y * 8 + x] = inputjpeg[base + x * 8 + y];
967
0
              }
968
0
            }
969
0
          } else {
970
0
            const int32_t scale =
971
0
                ColorCorrelation::RatioJPEG(cm[bx / kColorTileDimInBlocks]);
972
973
0
            for (size_t y = 0; y < 8; y++) {
974
0
              for (size_t x = 0; x < 8; x++) {
975
0
                int Y = inputjpegY[kDCTBlockSize * bx + x * 8 + y];
976
0
                int QChroma = inputjpeg[kDCTBlockSize * bx + x * 8 + y];
977
                // Fixed-point multiply of CfL scale with quant table ratio
978
                // first, and Y value second.
979
0
                int coeff_scale = (scale * scaled_qtable[64 * c + y * 8 + x] +
980
0
                                   (1 << (kCFLFixedPointPrecision - 1))) >>
981
0
                                  kCFLFixedPointPrecision;
982
0
                int cfl_factor =
983
0
                    (Y * coeff_scale + (1 << (kCFLFixedPointPrecision - 1))) >>
984
0
                    kCFLFixedPointPrecision;
985
0
                int QCR = QChroma - cfl_factor;
986
0
                block[y * 8 + x] = QCR;
987
0
              }
988
0
            }
989
0
          }
990
0
          enc_state->progressive_splitter.SplitACCoefficients(
991
0
              block, AcStrategy::FromRawStrategy(AcStrategy::Type::DCT), bx, by,
992
0
              coeffs);
993
0
          for (size_t i = 0; i < enc_state->coeffs.size(); i++) {
994
0
            coeffs[i] += kDCTBlockSize;
995
0
          }
996
0
        }
997
0
      }
998
0
    }
999
0
  }
1000
1001
0
  auto& dct = enc_state->shared.block_ctx_map.dc_thresholds;
1002
0
  auto& num_dc_ctxs = enc_state->shared.block_ctx_map.num_dc_ctxs;
1003
0
  num_dc_ctxs = 1;
1004
0
  for (size_t i = 0; i < 3; i++) {
1005
0
    dct[i].clear();
1006
0
    int num_thresholds = (CeilLog2Nonzero(total_dc[i]) - 12) / 2;
1007
    // up to 3 buckets per channel:
1008
    // dark/medium/bright, yellow/unsat/blue, green/unsat/red
1009
0
    num_thresholds = std::min(std::max(num_thresholds, 0), 2);
1010
0
    size_t cumsum = 0;
1011
0
    size_t cut = total_dc[i] / (num_thresholds + 1);
1012
0
    for (int j = 0; j < 2048; j++) {
1013
0
      cumsum += dc_counts[i][j];
1014
0
      if (cumsum > cut) {
1015
0
        dct[i].push_back(j - 1025);
1016
0
        cut = total_dc[i] * (dct[i].size() + 1) / (num_thresholds + 1);
1017
0
      }
1018
0
    }
1019
0
    num_dc_ctxs *= dct[i].size() + 1;
1020
0
  }
1021
1022
0
  auto& ctx_map = enc_state->shared.block_ctx_map.ctx_map;
1023
0
  ctx_map.clear();
1024
0
  ctx_map.resize(3 * kNumOrders * num_dc_ctxs, 0);
1025
1026
0
  int lbuckets = (dct[1].size() + 1);
1027
0
  for (size_t i = 0; i < num_dc_ctxs; i++) {
1028
    // up to 9 contexts for luma
1029
0
    ctx_map[i] = i / lbuckets;
1030
    // up to 3 contexts for chroma
1031
0
    ctx_map[kNumOrders * num_dc_ctxs + i] =
1032
0
        ctx_map[2 * kNumOrders * num_dc_ctxs + i] =
1033
0
            num_dc_ctxs / lbuckets + (i % lbuckets);
1034
0
  }
1035
0
  enc_state->shared.block_ctx_map.num_ctxs =
1036
0
      *std::max_element(ctx_map.begin(), ctx_map.end()) + 1;
1037
1038
  // disable DC frame for now
1039
0
  std::atomic<bool> has_error{false};
1040
0
  auto compute_dc_coeffs = [&](const uint32_t group_index,
1041
0
                               size_t /* thread */) {
1042
0
    if (has_error) return;
1043
0
    const Rect r = enc_state->shared.frame_dim.DCGroupRect(group_index);
1044
0
    if (!enc_modular->AddVarDCTDC(frame_header, dc, r, group_index,
1045
0
                                  /*nl_dc=*/false, enc_state,
1046
0
                                  /*jpeg_transcode=*/true)) {
1047
0
      has_error = true;
1048
0
      return;
1049
0
    }
1050
0
    if (!enc_modular->AddACMetadata(r, group_index, /*jpeg_transcode=*/true,
1051
0
                                    enc_state)) {
1052
0
      has_error = true;
1053
0
      return;
1054
0
    }
1055
0
  };
1056
0
  JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, shared.frame_dim.num_dc_groups,
1057
0
                                ThreadPool::NoInit, compute_dc_coeffs,
1058
0
                                "Compute DC coeffs"));
1059
0
  if (has_error) return JXL_FAILURE("Compute DC coeffs failed");
1060
1061
0
  return true;
1062
0
}
1063
1064
Status ComputeVarDCTEncodingData(const FrameHeader& frame_header,
1065
                                 const Image3F* linear,
1066
                                 Image3F* JXL_RESTRICT opsin, const Rect& rect,
1067
                                 const JxlCmsInterface& cms, ThreadPool* pool,
1068
                                 ModularFrameEncoder* enc_modular,
1069
                                 PassesEncoderState* enc_state,
1070
52
                                 AuxOut* aux_out) {
1071
52
  JXL_ASSERT((rect.xsize() % kBlockDim) == 0 &&
1072
52
             (rect.ysize() % kBlockDim) == 0);
1073
52
  JxlMemoryManager* memory_manager = enc_state->memory_manager();
1074
  // Save pre-Gaborish opsin for AR control field heuristics computation.
1075
52
  Image3F orig_opsin;
1076
52
  JXL_ASSIGN_OR_RETURN(
1077
52
      orig_opsin, Image3F::Create(memory_manager, rect.xsize(), rect.ysize()));
1078
52
  CopyImageTo(rect, *opsin, Rect(orig_opsin), &orig_opsin);
1079
52
  orig_opsin.ShrinkTo(enc_state->shared.frame_dim.xsize,
1080
52
                      enc_state->shared.frame_dim.ysize);
1081
1082
52
  JXL_RETURN_IF_ERROR(LossyFrameHeuristics(frame_header, enc_state, enc_modular,
1083
52
                                           linear, opsin, rect, cms, pool,
1084
52
                                           aux_out));
1085
1086
52
  JXL_RETURN_IF_ERROR(InitializePassesEncoder(
1087
52
      frame_header, *opsin, rect, cms, pool, enc_state, enc_modular, aux_out));
1088
1089
52
  JXL_RETURN_IF_ERROR(
1090
52
      ComputeARHeuristics(frame_header, enc_state, orig_opsin, rect, pool));
1091
1092
52
  JXL_RETURN_IF_ERROR(ComputeACMetadata(pool, enc_state, enc_modular));
1093
1094
52
  return true;
1095
52
}
1096
1097
void ComputeAllCoeffOrders(PassesEncoderState& enc_state,
1098
52
                           const FrameDimensions& frame_dim) {
1099
52
  auto used_orders_info = ComputeUsedOrders(
1100
52
      enc_state.cparams.speed_tier, enc_state.shared.ac_strategy,
1101
52
      Rect(enc_state.shared.raw_quant_field));
1102
52
  enc_state.used_orders.resize(enc_state.progressive_splitter.GetNumPasses());
1103
104
  for (size_t i = 0; i < enc_state.progressive_splitter.GetNumPasses(); i++) {
1104
52
    ComputeCoeffOrder(
1105
52
        enc_state.cparams.speed_tier, *enc_state.coeffs[i],
1106
52
        enc_state.shared.ac_strategy, frame_dim, enc_state.used_orders[i],
1107
52
        enc_state.used_acs, used_orders_info.first, used_orders_info.second,
1108
52
        &enc_state.shared.coeff_orders[i * enc_state.shared.coeff_order_size]);
1109
52
  }
1110
52
  enc_state.used_acs |= used_orders_info.first;
1111
52
}
1112
1113
// Working area for TokenizeCoefficients (per-group!)
1114
struct EncCache {
1115
  // Allocates memory when first called.
1116
52
  Status InitOnce(JxlMemoryManager* memory_manager) {
1117
52
    if (num_nzeroes.xsize() == 0) {
1118
52
      JXL_ASSIGN_OR_RETURN(num_nzeroes,
1119
52
                           Image3I::Create(memory_manager, kGroupDimInBlocks,
1120
52
                                           kGroupDimInBlocks));
1121
52
    }
1122
52
    return true;
1123
52
  }
1124
  // TokenizeCoefficients
1125
  Image3I num_nzeroes;
1126
};
1127
1128
Status TokenizeAllCoefficients(const FrameHeader& frame_header,
1129
                               ThreadPool* pool,
1130
52
                               PassesEncoderState* enc_state) {
1131
52
  PassesSharedState& shared = enc_state->shared;
1132
52
  std::vector<EncCache> group_caches;
1133
52
  JxlMemoryManager* memory_manager = enc_state->memory_manager();
1134
52
  const auto tokenize_group_init = [&](const size_t num_threads) {
1135
52
    group_caches.resize(num_threads);
1136
52
    return true;
1137
52
  };
1138
52
  std::atomic<bool> has_error{false};
1139
52
  const auto tokenize_group = [&](const uint32_t group_index,
1140
52
                                  const size_t thread) {
1141
52
    if (has_error) return;
1142
    // Tokenize coefficients.
1143
52
    const Rect rect = shared.frame_dim.BlockGroupRect(group_index);
1144
104
    for (size_t idx_pass = 0; idx_pass < enc_state->passes.size(); idx_pass++) {
1145
52
      JXL_ASSERT(enc_state->coeffs[idx_pass]->Type() == ACType::k32);
1146
52
      const int32_t* JXL_RESTRICT ac_rows[3] = {
1147
52
          enc_state->coeffs[idx_pass]->PlaneRow(0, group_index, 0).ptr32,
1148
52
          enc_state->coeffs[idx_pass]->PlaneRow(1, group_index, 0).ptr32,
1149
52
          enc_state->coeffs[idx_pass]->PlaneRow(2, group_index, 0).ptr32,
1150
52
      };
1151
      // Ensure group cache is initialized.
1152
52
      if (!group_caches[thread].InitOnce(memory_manager)) {
1153
0
        has_error = true;
1154
0
        return;
1155
0
      }
1156
52
      TokenizeCoefficients(
1157
52
          &shared.coeff_orders[idx_pass * shared.coeff_order_size], rect,
1158
52
          ac_rows, shared.ac_strategy, frame_header.chroma_subsampling,
1159
52
          &group_caches[thread].num_nzeroes,
1160
52
          &enc_state->passes[idx_pass].ac_tokens[group_index], shared.quant_dc,
1161
52
          shared.raw_quant_field, shared.block_ctx_map);
1162
52
    }
1163
52
  };
1164
52
  JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, shared.frame_dim.num_groups,
1165
52
                                tokenize_group_init, tokenize_group,
1166
52
                                "TokenizeGroup"));
1167
52
  if (has_error) return JXL_FAILURE("TokenizeGroup failed");
1168
52
  return true;
1169
52
}
1170
1171
Status EncodeGlobalDCInfo(const PassesSharedState& shared, BitWriter* writer,
1172
52
                          AuxOut* aux_out) {
1173
  // Encode quantizer DC and global scale.
1174
52
  QuantizerParams params = shared.quantizer.GetParams();
1175
52
  JXL_RETURN_IF_ERROR(
1176
52
      WriteQuantizerParams(params, writer, kLayerQuant, aux_out));
1177
52
  EncodeBlockCtxMap(shared.block_ctx_map, writer, aux_out);
1178
52
  ColorCorrelationEncodeDC(shared.cmap.base(), writer, kLayerDC, aux_out);
1179
52
  return true;
1180
52
}
1181
1182
// In streaming mode, this function only performs the histogram clustering and
1183
// saves the histogram bitstreams in enc_state, the actual AC global bitstream
1184
// is written in OutputAcGlobal() function after all the groups are processed.
1185
Status EncodeGlobalACInfo(PassesEncoderState* enc_state, BitWriter* writer,
1186
52
                          ModularFrameEncoder* enc_modular, AuxOut* aux_out) {
1187
52
  PassesSharedState& shared = enc_state->shared;
1188
52
  JxlMemoryManager* memory_manager = enc_state->memory_manager();
1189
52
  JXL_RETURN_IF_ERROR(DequantMatricesEncode(memory_manager, shared.matrices,
1190
52
                                            writer, kLayerQuant, aux_out,
1191
52
                                            enc_modular));
1192
52
  size_t num_histo_bits = CeilLog2Nonzero(shared.frame_dim.num_groups);
1193
52
  if (!enc_state->streaming_mode && num_histo_bits != 0) {
1194
0
    BitWriter::Allotment allotment(writer, num_histo_bits);
1195
0
    writer->Write(num_histo_bits, shared.num_histograms - 1);
1196
0
    allotment.ReclaimAndCharge(writer, kLayerAC, aux_out);
1197
0
  }
1198
1199
104
  for (size_t i = 0; i < enc_state->progressive_splitter.GetNumPasses(); i++) {
1200
    // Encode coefficient orders.
1201
52
    if (!enc_state->streaming_mode) {
1202
52
      size_t order_bits = 0;
1203
52
      JXL_RETURN_IF_ERROR(U32Coder::CanEncode(
1204
52
          kOrderEnc, enc_state->used_orders[i], &order_bits));
1205
52
      BitWriter::Allotment allotment(writer, order_bits);
1206
52
      JXL_CHECK(U32Coder::Write(kOrderEnc, enc_state->used_orders[i], writer));
1207
52
      allotment.ReclaimAndCharge(writer, kLayerOrder, aux_out);
1208
52
      EncodeCoeffOrders(enc_state->used_orders[i],
1209
52
                        &shared.coeff_orders[i * shared.coeff_order_size],
1210
52
                        writer, kLayerOrder, aux_out);
1211
52
    }
1212
1213
    // Encode histograms.
1214
52
    HistogramParams hist_params(enc_state->cparams.speed_tier,
1215
52
                                shared.block_ctx_map.NumACContexts());
1216
52
    if (enc_state->cparams.speed_tier > SpeedTier::kTortoise) {
1217
52
      hist_params.lz77_method = HistogramParams::LZ77Method::kNone;
1218
52
    }
1219
52
    if (enc_state->cparams.decoding_speed_tier >= 1) {
1220
0
      hist_params.max_histograms = 6;
1221
0
    }
1222
52
    size_t num_histogram_groups = shared.num_histograms;
1223
52
    if (enc_state->streaming_mode) {
1224
0
      size_t prev_num_histograms =
1225
0
          enc_state->passes[i].codes.encoding_info.size();
1226
0
      if (enc_state->initialize_global_state) {
1227
0
        prev_num_histograms += kNumFixedHistograms;
1228
0
        hist_params.add_fixed_histograms = true;
1229
0
      }
1230
0
      size_t remaining_histograms = kClustersLimit - prev_num_histograms;
1231
      // Heuristic to assign budget of new histograms to DC groups.
1232
      // TODO(szabadka) Tune this together with the DC group ordering.
1233
0
      size_t max_histograms = remaining_histograms < 20
1234
0
                                  ? std::min<size_t>(remaining_histograms, 4)
1235
0
                                  : remaining_histograms / 4;
1236
0
      hist_params.max_histograms =
1237
0
          std::min(max_histograms, hist_params.max_histograms);
1238
0
      num_histogram_groups = 1;
1239
0
    }
1240
52
    hist_params.streaming_mode = enc_state->streaming_mode;
1241
52
    hist_params.initialize_global_state = enc_state->initialize_global_state;
1242
52
    BuildAndEncodeHistograms(
1243
52
        memory_manager, hist_params,
1244
52
        num_histogram_groups * shared.block_ctx_map.NumACContexts(),
1245
52
        enc_state->passes[i].ac_tokens, &enc_state->passes[i].codes,
1246
52
        &enc_state->passes[i].context_map, writer, kLayerAC, aux_out);
1247
52
  }
1248
1249
52
  return true;
1250
52
}
1251
1252
Status EncodeGroups(const FrameHeader& frame_header,
1253
                    PassesEncoderState* enc_state,
1254
                    ModularFrameEncoder* enc_modular, ThreadPool* pool,
1255
                    std::vector<std::unique_ptr<BitWriter>>* group_codes,
1256
52
                    AuxOut* aux_out) {
1257
52
  const PassesSharedState& shared = enc_state->shared;
1258
52
  JxlMemoryManager* memory_manager = shared.memory_manager;
1259
52
  const FrameDimensions& frame_dim = shared.frame_dim;
1260
52
  const size_t num_groups = frame_dim.num_groups;
1261
52
  const size_t num_passes = enc_state->progressive_splitter.GetNumPasses();
1262
52
  const size_t global_ac_index = frame_dim.num_dc_groups + 1;
1263
52
  const bool is_small_image =
1264
52
      !enc_state->streaming_mode && num_groups == 1 && num_passes == 1;
1265
52
  const size_t num_toc_entries =
1266
52
      is_small_image ? 1
1267
52
                     : AcGroupIndex(0, 0, num_groups, frame_dim.num_dc_groups) +
1268
0
                           num_groups * num_passes;
1269
52
  JXL_ASSERT(group_codes->empty());
1270
52
  group_codes->reserve(num_toc_entries);
1271
104
  for (size_t i = 0; i < num_toc_entries; ++i) {
1272
52
    group_codes->emplace_back(jxl::make_unique<BitWriter>(memory_manager));
1273
52
  }
1274
1275
416
  const auto get_output = [&](const size_t index) -> BitWriter* {
1276
416
    return (*group_codes)[is_small_image ? 0 : index].get();
1277
416
  };
1278
104
  auto ac_group_code = [&](size_t pass, size_t group) {
1279
104
    return get_output(AcGroupIndex(pass, group, frame_dim.num_groups,
1280
104
                                   frame_dim.num_dc_groups));
1281
104
  };
1282
1283
52
  if (enc_state->initialize_global_state) {
1284
52
    if (frame_header.flags & FrameHeader::kPatches) {
1285
0
      PatchDictionaryEncoder::Encode(shared.image_features.patches,
1286
0
                                     get_output(0), kLayerDictionary, aux_out);
1287
0
    }
1288
52
    if (frame_header.flags & FrameHeader::kSplines) {
1289
0
      EncodeSplines(shared.image_features.splines, get_output(0), kLayerSplines,
1290
0
                    HistogramParams(), aux_out);
1291
0
    }
1292
52
    if (frame_header.flags & FrameHeader::kNoise) {
1293
0
      EncodeNoise(shared.image_features.noise_params, get_output(0),
1294
0
                  kLayerNoise, aux_out);
1295
0
    }
1296
1297
52
    JXL_RETURN_IF_ERROR(DequantMatricesEncodeDC(shared.matrices, get_output(0),
1298
52
                                                kLayerQuant, aux_out));
1299
52
    if (frame_header.encoding == FrameEncoding::kVarDCT) {
1300
52
      JXL_RETURN_IF_ERROR(EncodeGlobalDCInfo(shared, get_output(0), aux_out));
1301
52
    }
1302
52
    JXL_RETURN_IF_ERROR(enc_modular->EncodeGlobalInfo(enc_state->streaming_mode,
1303
52
                                                      get_output(0), aux_out));
1304
52
    JXL_RETURN_IF_ERROR(enc_modular->EncodeStream(get_output(0), aux_out,
1305
52
                                                  kLayerModularGlobal,
1306
52
                                                  ModularStreamId::Global()));
1307
52
  }
1308
1309
52
  std::vector<std::unique_ptr<AuxOut>> aux_outs;
1310
52
  auto resize_aux_outs = [&aux_outs,
1311
156
                          aux_out](const size_t num_threads) -> Status {
1312
156
    if (aux_out == nullptr) {
1313
156
      aux_outs.resize(num_threads);
1314
156
    } else {
1315
0
      while (aux_outs.size() > num_threads) {
1316
0
        aux_out->Assimilate(*aux_outs.back());
1317
0
        aux_outs.pop_back();
1318
0
      }
1319
0
      while (num_threads > aux_outs.size()) {
1320
0
        aux_outs.emplace_back(jxl::make_unique<AuxOut>());
1321
0
      }
1322
0
    }
1323
156
    return true;
1324
156
  };
1325
1326
52
  const auto process_dc_group = [&](const uint32_t group_index,
1327
52
                                    const size_t thread) {
1328
52
    AuxOut* my_aux_out = aux_outs[thread].get();
1329
52
    BitWriter* output = get_output(group_index + 1);
1330
52
    int modular_group_index = group_index;
1331
52
    if (enc_state->streaming_mode) {
1332
0
      JXL_ASSERT(group_index == 0);
1333
0
      modular_group_index = enc_state->dc_group_index;
1334
0
    }
1335
52
    if (frame_header.encoding == FrameEncoding::kVarDCT &&
1336
52
        !(frame_header.flags & FrameHeader::kUseDcFrame)) {
1337
52
      BitWriter::Allotment allotment(output, 2);
1338
52
      output->Write(2, enc_modular->extra_dc_precision[modular_group_index]);
1339
52
      allotment.ReclaimAndCharge(output, kLayerDC, my_aux_out);
1340
52
      JXL_CHECK(enc_modular->EncodeStream(
1341
52
          output, my_aux_out, kLayerDC,
1342
52
          ModularStreamId::VarDCTDC(modular_group_index)));
1343
52
    }
1344
52
    JXL_CHECK(enc_modular->EncodeStream(
1345
52
        output, my_aux_out, kLayerModularDcGroup,
1346
52
        ModularStreamId::ModularDC(modular_group_index)));
1347
52
    if (frame_header.encoding == FrameEncoding::kVarDCT) {
1348
52
      const Rect& rect = enc_state->shared.frame_dim.DCGroupRect(group_index);
1349
52
      size_t nb_bits = CeilLog2Nonzero(rect.xsize() * rect.ysize());
1350
52
      if (nb_bits != 0) {
1351
0
        BitWriter::Allotment allotment(output, nb_bits);
1352
0
        output->Write(nb_bits,
1353
0
                      enc_modular->ac_metadata_size[modular_group_index] - 1);
1354
0
        allotment.ReclaimAndCharge(output, kLayerControlFields, my_aux_out);
1355
0
      }
1356
52
      JXL_CHECK(enc_modular->EncodeStream(
1357
52
          output, my_aux_out, kLayerControlFields,
1358
52
          ModularStreamId::ACMetadata(modular_group_index)));
1359
52
    }
1360
52
  };
1361
52
  JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, frame_dim.num_dc_groups,
1362
52
                                resize_aux_outs, process_dc_group,
1363
52
                                "EncodeDCGroup"));
1364
1365
52
  if (frame_header.encoding == FrameEncoding::kVarDCT) {
1366
52
    JXL_RETURN_IF_ERROR(EncodeGlobalACInfo(
1367
52
        enc_state, get_output(global_ac_index), enc_modular, aux_out));
1368
52
  }
1369
1370
52
  std::atomic<bool> has_error{false};
1371
52
  const auto process_group = [&](const uint32_t group_index,
1372
52
                                 const size_t thread) {
1373
52
    if (has_error) return;
1374
52
    AuxOut* my_aux_out = aux_outs[thread].get();
1375
1376
52
    size_t ac_group_id =
1377
52
        enc_state->streaming_mode
1378
52
            ? enc_modular->ComputeStreamingAbsoluteAcGroupId(
1379
0
                  enc_state->dc_group_index, group_index, shared.frame_dim)
1380
52
            : group_index;
1381
1382
104
    for (size_t i = 0; i < num_passes; i++) {
1383
52
      JXL_DEBUG_V(2, "Encoding AC group %u [abs %" PRIuS "] pass %" PRIuS,
1384
52
                  group_index, ac_group_id, i);
1385
52
      if (frame_header.encoding == FrameEncoding::kVarDCT) {
1386
52
        if (!EncodeGroupTokenizedCoefficients(
1387
52
                group_index, i, enc_state->histogram_idx[group_index],
1388
52
                *enc_state, ac_group_code(i, group_index), my_aux_out)) {
1389
0
          has_error = true;
1390
0
          return;
1391
0
        }
1392
52
      }
1393
      // Write all modular encoded data (color?, alpha, depth, extra channels)
1394
52
      if (!enc_modular->EncodeStream(
1395
52
              ac_group_code(i, group_index), my_aux_out, kLayerModularAcGroup,
1396
52
              ModularStreamId::ModularAC(ac_group_id, i))) {
1397
0
        has_error = true;
1398
0
        return;
1399
0
      }
1400
52
      JXL_DEBUG_V(2,
1401
52
                  "AC group %u [abs %" PRIuS "] pass %" PRIuS
1402
52
                  " encoded size is %" PRIuS " bits",
1403
52
                  group_index, ac_group_id, i,
1404
52
                  ac_group_code(i, group_index)->BitsWritten());
1405
52
    }
1406
52
  };
1407
52
  JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, num_groups, resize_aux_outs,
1408
52
                                process_group, "EncodeGroupCoefficients"));
1409
52
  if (has_error) return JXL_FAILURE("EncodeGroupCoefficients failed");
1410
  // Resizing aux_outs to 0 also Assimilates the array.
1411
52
  static_cast<void>(resize_aux_outs(0));
1412
1413
52
  for (std::unique_ptr<BitWriter>& bw : *group_codes) {
1414
52
    BitWriter::Allotment allotment(bw.get(), 8);
1415
52
    bw->ZeroPadToByte();  // end of group.
1416
52
    allotment.ReclaimAndCharge(bw.get(), kLayerAC, aux_out);
1417
52
  }
1418
52
  return true;
1419
52
}
1420
1421
Status ComputeEncodingData(
1422
    const CompressParams& cparams, const FrameInfo& frame_info,
1423
    const CodecMetadata* metadata, JxlEncoderChunkedFrameAdapter& frame_data,
1424
    const jpeg::JPEGData* jpeg_data, size_t x0, size_t y0, size_t xsize,
1425
    size_t ysize, const JxlCmsInterface& cms, ThreadPool* pool,
1426
    FrameHeader& mutable_frame_header, ModularFrameEncoder& enc_modular,
1427
    PassesEncoderState& enc_state,
1428
52
    std::vector<std::unique_ptr<BitWriter>>* group_codes, AuxOut* aux_out) {
1429
52
  JXL_ASSERT(x0 + xsize <= frame_data.xsize);
1430
52
  JXL_ASSERT(y0 + ysize <= frame_data.ysize);
1431
52
  JxlMemoryManager* memory_manager = enc_state.memory_manager();
1432
52
  const FrameHeader& frame_header = mutable_frame_header;
1433
52
  PassesSharedState& shared = enc_state.shared;
1434
52
  shared.metadata = metadata;
1435
52
  if (enc_state.streaming_mode) {
1436
0
    shared.frame_dim.Set(
1437
0
        xsize, ysize, frame_header.group_size_shift,
1438
0
        /*max_hshift=*/0, /*max_vshift=*/0,
1439
0
        mutable_frame_header.encoding == FrameEncoding::kModular,
1440
0
        /*upsampling=*/1);
1441
52
  } else {
1442
52
    shared.frame_dim = frame_header.ToFrameDimensions();
1443
52
  }
1444
1445
52
  shared.image_features.patches.SetPassesSharedState(&shared);
1446
52
  const FrameDimensions& frame_dim = shared.frame_dim;
1447
52
  JXL_ASSIGN_OR_RETURN(
1448
52
      shared.ac_strategy,
1449
52
      AcStrategyImage::Create(memory_manager, frame_dim.xsize_blocks,
1450
52
                              frame_dim.ysize_blocks));
1451
52
  JXL_ASSIGN_OR_RETURN(shared.raw_quant_field,
1452
52
                       ImageI::Create(memory_manager, frame_dim.xsize_blocks,
1453
52
                                      frame_dim.ysize_blocks));
1454
52
  JXL_ASSIGN_OR_RETURN(shared.epf_sharpness,
1455
52
                       ImageB::Create(memory_manager, frame_dim.xsize_blocks,
1456
52
                                      frame_dim.ysize_blocks));
1457
52
  JXL_ASSIGN_OR_RETURN(
1458
52
      shared.cmap, ColorCorrelationMap::Create(memory_manager, frame_dim.xsize,
1459
52
                                               frame_dim.ysize));
1460
52
  shared.coeff_order_size = kCoeffOrderMaxSize;
1461
52
  if (frame_header.encoding == FrameEncoding::kVarDCT) {
1462
52
    shared.coeff_orders.resize(frame_header.passes.num_passes *
1463
52
                               kCoeffOrderMaxSize);
1464
52
  }
1465
1466
52
  JXL_ASSIGN_OR_RETURN(shared.quant_dc,
1467
52
                       ImageB::Create(memory_manager, frame_dim.xsize_blocks,
1468
52
                                      frame_dim.ysize_blocks));
1469
52
  JXL_ASSIGN_OR_RETURN(shared.dc_storage,
1470
52
                       Image3F::Create(memory_manager, frame_dim.xsize_blocks,
1471
52
                                       frame_dim.ysize_blocks));
1472
52
  shared.dc = &shared.dc_storage;
1473
1474
52
  const size_t num_extra_channels = metadata->m.num_extra_channels;
1475
52
  const ExtraChannelInfo* alpha_eci = metadata->m.Find(ExtraChannel::kAlpha);
1476
52
  const ExtraChannelInfo* black_eci = metadata->m.Find(ExtraChannel::kBlack);
1477
52
  const size_t alpha_idx = alpha_eci - metadata->m.extra_channel_info.data();
1478
52
  const size_t black_idx = black_eci - metadata->m.extra_channel_info.data();
1479
52
  const ColorEncoding c_enc = metadata->m.color_encoding;
1480
1481
  // Make the image patch bigger than the currently processed group in streaming
1482
  // mode so that we can take into account border pixels around the group when
1483
  // computing inverse Gaborish and adaptive quantization map.
1484
52
  int max_border = enc_state.streaming_mode ? kBlockDim : 0;
1485
52
  Rect frame_rect(0, 0, frame_data.xsize, frame_data.ysize);
1486
52
  Rect frame_area_rect = Rect(x0, y0, xsize, ysize);
1487
52
  Rect patch_rect = frame_area_rect.Extend(max_border, frame_rect);
1488
52
  JXL_ASSERT(patch_rect.IsInside(frame_rect));
1489
1490
  // Allocating a large enough image avoids a copy when padding.
1491
104
  JXL_ASSIGN_OR_RETURN(
1492
104
      Image3F color,
1493
104
      Image3F::Create(memory_manager, RoundUpToBlockDim(patch_rect.xsize()),
1494
104
                      RoundUpToBlockDim(patch_rect.ysize())));
1495
104
  color.ShrinkTo(patch_rect.xsize(), patch_rect.ysize());
1496
104
  std::vector<ImageF> extra_channels(num_extra_channels);
1497
104
  for (auto& extra_channel : extra_channels) {
1498
0
    JXL_ASSIGN_OR_RETURN(
1499
0
        extra_channel,
1500
0
        ImageF::Create(memory_manager, patch_rect.xsize(), patch_rect.ysize()));
1501
0
  }
1502
52
  ImageF* alpha = alpha_eci ? &extra_channels[alpha_idx] : nullptr;
1503
52
  ImageF* black = black_eci ? &extra_channels[black_idx] : nullptr;
1504
52
  bool has_interleaved_alpha = false;
1505
52
  JxlChunkedFrameInputSource input = frame_data.GetInputSource();
1506
52
  if (!frame_data.IsJPEG()) {
1507
52
    JXL_RETURN_IF_ERROR(CopyColorChannels(input, patch_rect, frame_info,
1508
52
                                          metadata->m, pool, &color, alpha,
1509
52
                                          &has_interleaved_alpha));
1510
52
  }
1511
52
  JXL_RETURN_IF_ERROR(CopyExtraChannels(input, patch_rect, frame_info,
1512
52
                                        metadata->m, has_interleaved_alpha,
1513
52
                                        pool, &extra_channels));
1514
1515
52
  shared.image_features.patches.SetPassesSharedState(&shared);
1516
52
  enc_state.cparams = cparams;
1517
1518
52
  Image3F linear_storage;
1519
52
  Image3F* linear = nullptr;
1520
1521
52
  if (!jpeg_data) {
1522
52
    if (frame_header.color_transform == ColorTransform::kXYB &&
1523
52
        frame_info.ib_needs_color_transform) {
1524
52
      if (frame_header.encoding == FrameEncoding::kVarDCT &&
1525
52
          cparams.speed_tier <= SpeedTier::kKitten) {
1526
0
        JXL_ASSIGN_OR_RETURN(linear_storage,
1527
0
                             Image3F::Create(memory_manager, patch_rect.xsize(),
1528
0
                                             patch_rect.ysize()));
1529
0
        linear = &linear_storage;
1530
0
      }
1531
52
      ToXYB(c_enc, metadata->m.IntensityTarget(), black, pool, &color, cms,
1532
52
            linear);
1533
52
    } else {
1534
      // Nothing to do.
1535
      // RGB or YCbCr: forward YCbCr is not implemented, this is only used when
1536
      // the input is already in YCbCr
1537
      // If encoding a special DC or reference frame: input is already in XYB.
1538
0
    }
1539
52
    bool lossless = cparams.IsLossless();
1540
52
    if (alpha && !alpha_eci->alpha_associated &&
1541
52
        frame_header.frame_type == FrameType::kRegularFrame &&
1542
52
        !ApplyOverride(cparams.keep_invisible, cparams.IsLossless()) &&
1543
52
        cparams.ec_resampling == cparams.resampling &&
1544
52
        !cparams.disable_percepeptual_optimizations) {
1545
      // simplify invisible pixels
1546
0
      SimplifyInvisible(&color, *alpha, lossless);
1547
0
      if (linear) {
1548
0
        SimplifyInvisible(linear, *alpha, lossless);
1549
0
      }
1550
0
    }
1551
52
    PadImageToBlockMultipleInPlace(&color);
1552
52
  }
1553
1554
  // Rectangle within color that corresponds to the currently processed group in
1555
  // streaming mode.
1556
52
  Rect group_rect(x0 - patch_rect.x0(), y0 - patch_rect.y0(),
1557
52
                  RoundUpToBlockDim(xsize), RoundUpToBlockDim(ysize));
1558
1559
52
  if (enc_state.initialize_global_state && !jpeg_data) {
1560
52
    ComputeChromacityAdjustments(cparams, color, group_rect,
1561
52
                                 &mutable_frame_header);
1562
52
  }
1563
1564
52
  bool has_jpeg_data = (jpeg_data != nullptr);
1565
52
  ComputeNoiseParams(cparams, enc_state.streaming_mode, has_jpeg_data, color,
1566
52
                     frame_dim, &mutable_frame_header,
1567
52
                     &shared.image_features.noise_params);
1568
1569
52
  JXL_RETURN_IF_ERROR(
1570
52
      DownsampleColorChannels(cparams, frame_header, has_jpeg_data, &color));
1571
1572
52
  if (cparams.ec_resampling != 1 && !cparams.already_downsampled) {
1573
0
    for (ImageF& ec : extra_channels) {
1574
0
      JXL_ASSIGN_OR_RETURN(ec, DownsampleImage(ec, cparams.ec_resampling));
1575
0
    }
1576
0
  }
1577
1578
52
  if (!enc_state.streaming_mode) {
1579
52
    group_rect = Rect(color);
1580
52
  }
1581
1582
52
  if (frame_header.encoding == FrameEncoding::kVarDCT) {
1583
52
    enc_state.passes.resize(enc_state.progressive_splitter.GetNumPasses());
1584
52
    for (PassesEncoderState::PassData& pass : enc_state.passes) {
1585
52
      pass.ac_tokens.resize(shared.frame_dim.num_groups);
1586
52
    }
1587
52
    if (jpeg_data) {
1588
0
      JXL_RETURN_IF_ERROR(ComputeJPEGTranscodingData(
1589
0
          *jpeg_data, frame_header, pool, &enc_modular, &enc_state));
1590
52
    } else {
1591
52
      JXL_RETURN_IF_ERROR(ComputeVarDCTEncodingData(
1592
52
          frame_header, linear, &color, group_rect, cms, pool, &enc_modular,
1593
52
          &enc_state, aux_out));
1594
52
    }
1595
52
    ComputeAllCoeffOrders(enc_state, frame_dim);
1596
52
    if (!enc_state.streaming_mode) {
1597
52
      shared.num_histograms = 1;
1598
52
      enc_state.histogram_idx.resize(frame_dim.num_groups);
1599
52
    }
1600
52
    JXL_RETURN_IF_ERROR(
1601
52
        TokenizeAllCoefficients(frame_header, pool, &enc_state));
1602
52
  }
1603
1604
52
  if (cparams.modular_mode || !extra_channels.empty()) {
1605
0
    JXL_RETURN_IF_ERROR(enc_modular.ComputeEncodingData(
1606
0
        frame_header, metadata->m, &color, extra_channels, group_rect,
1607
0
        frame_dim, frame_area_rect, &enc_state, cms, pool, aux_out,
1608
0
        /*do_color=*/cparams.modular_mode));
1609
0
  }
1610
1611
52
  if (!enc_state.streaming_mode) {
1612
52
    if (cparams.speed_tier < SpeedTier::kTortoise ||
1613
52
        !cparams.ModularPartIsLossless() || cparams.responsive ||
1614
52
        !cparams.custom_fixed_tree.empty()) {
1615
      // Use local trees if doing lossless modular, unless at very slow speeds.
1616
52
      JXL_RETURN_IF_ERROR(enc_modular.ComputeTree(pool));
1617
52
      JXL_RETURN_IF_ERROR(enc_modular.ComputeTokens(pool));
1618
52
    }
1619
52
    mutable_frame_header.UpdateFlag(shared.image_features.patches.HasAny(),
1620
52
                                    FrameHeader::kPatches);
1621
52
    mutable_frame_header.UpdateFlag(shared.image_features.splines.HasAny(),
1622
52
                                    FrameHeader::kSplines);
1623
52
  }
1624
1625
52
  JXL_RETURN_IF_ERROR(EncodeGroups(frame_header, &enc_state, &enc_modular, pool,
1626
52
                                   group_codes, aux_out));
1627
52
  if (enc_state.streaming_mode) {
1628
0
    const size_t group_index = enc_state.dc_group_index;
1629
0
    enc_modular.ClearStreamData(ModularStreamId::VarDCTDC(group_index));
1630
0
    enc_modular.ClearStreamData(ModularStreamId::ACMetadata(group_index));
1631
0
    enc_modular.ClearModularStreamData();
1632
0
  }
1633
52
  return true;
1634
52
}
1635
1636
Status PermuteGroups(const CompressParams& cparams,
1637
                     const FrameDimensions& frame_dim, size_t num_passes,
1638
                     std::vector<coeff_order_t>* permutation,
1639
52
                     std::vector<std::unique_ptr<BitWriter>>* group_codes) {
1640
52
  const size_t num_groups = frame_dim.num_groups;
1641
52
  if (!cparams.centerfirst || (num_passes == 1 && num_groups == 1)) {
1642
52
    return true;
1643
52
  }
1644
  // Don't permute global DC/AC or DC.
1645
0
  permutation->resize(frame_dim.num_dc_groups + 2);
1646
0
  std::iota(permutation->begin(), permutation->end(), 0);
1647
0
  std::vector<coeff_order_t> ac_group_order(num_groups);
1648
0
  std::iota(ac_group_order.begin(), ac_group_order.end(), 0);
1649
0
  size_t group_dim = frame_dim.group_dim;
1650
1651
  // The center of the image is either given by parameters or chosen
1652
  // to be the middle of the image by default if center_x, center_y resp.
1653
  // are not provided.
1654
1655
0
  int64_t imag_cx;
1656
0
  if (cparams.center_x != static_cast<size_t>(-1)) {
1657
0
    JXL_RETURN_IF_ERROR(cparams.center_x < frame_dim.xsize);
1658
0
    imag_cx = cparams.center_x;
1659
0
  } else {
1660
0
    imag_cx = frame_dim.xsize / 2;
1661
0
  }
1662
1663
0
  int64_t imag_cy;
1664
0
  if (cparams.center_y != static_cast<size_t>(-1)) {
1665
0
    JXL_RETURN_IF_ERROR(cparams.center_y < frame_dim.ysize);
1666
0
    imag_cy = cparams.center_y;
1667
0
  } else {
1668
0
    imag_cy = frame_dim.ysize / 2;
1669
0
  }
1670
1671
  // The center of the group containing the center of the image.
1672
0
  int64_t cx = (imag_cx / group_dim) * group_dim + group_dim / 2;
1673
0
  int64_t cy = (imag_cy / group_dim) * group_dim + group_dim / 2;
1674
  // This identifies in what area of the central group the center of the image
1675
  // lies in.
1676
0
  double direction = -std::atan2(imag_cy - cy, imag_cx - cx);
1677
  // This identifies the side of the central group the center of the image
1678
  // lies closest to. This can take values 0, 1, 2, 3 corresponding to left,
1679
  // bottom, right, top.
1680
0
  int64_t side = std::fmod((direction + 5 * kPi / 4), 2 * kPi) * 2 / kPi;
1681
0
  auto get_distance_from_center = [&](size_t gid) {
1682
0
    Rect r = frame_dim.GroupRect(gid);
1683
0
    int64_t gcx = r.x0() + group_dim / 2;
1684
0
    int64_t gcy = r.y0() + group_dim / 2;
1685
0
    int64_t dx = gcx - cx;
1686
0
    int64_t dy = gcy - cy;
1687
    // The angle is determined by taking atan2 and adding an appropriate
1688
    // starting point depending on the side we want to start on.
1689
0
    double angle = std::remainder(
1690
0
        std::atan2(dy, dx) + kPi / 4 + side * (kPi / 2), 2 * kPi);
1691
    // Concentric squares in clockwise order.
1692
0
    return std::make_pair(std::max(std::abs(dx), std::abs(dy)), angle);
1693
0
  };
1694
0
  std::sort(ac_group_order.begin(), ac_group_order.end(),
1695
0
            [&](coeff_order_t a, coeff_order_t b) {
1696
0
              return get_distance_from_center(a) < get_distance_from_center(b);
1697
0
            });
1698
0
  std::vector<coeff_order_t> inv_ac_group_order(ac_group_order.size(), 0);
1699
0
  for (size_t i = 0; i < ac_group_order.size(); i++) {
1700
0
    inv_ac_group_order[ac_group_order[i]] = i;
1701
0
  }
1702
0
  for (size_t i = 0; i < num_passes; i++) {
1703
0
    size_t pass_start = permutation->size();
1704
0
    for (coeff_order_t v : inv_ac_group_order) {
1705
0
      permutation->push_back(pass_start + v);
1706
0
    }
1707
0
  }
1708
0
  std::vector<std::unique_ptr<BitWriter>> new_group_codes(group_codes->size());
1709
0
  for (size_t i = 0; i < permutation->size(); i++) {
1710
0
    new_group_codes[(*permutation)[i]] = std::move((*group_codes)[i]);
1711
0
  }
1712
0
  group_codes->swap(new_group_codes);
1713
0
  return true;
1714
0
}
1715
1716
bool CanDoStreamingEncoding(const CompressParams& cparams,
1717
                            const FrameInfo& frame_info,
1718
                            const CodecMetadata& metadata,
1719
52
                            const JxlEncoderChunkedFrameAdapter& frame_data) {
1720
52
  if (cparams.buffering == 0) {
1721
0
    return false;
1722
0
  }
1723
52
  if (cparams.buffering == -1) {
1724
52
    if (cparams.speed_tier < SpeedTier::kTortoise) return false;
1725
52
    if (cparams.speed_tier < SpeedTier::kSquirrel &&
1726
52
        cparams.butteraugli_distance > 0.5f) {
1727
0
      return false;
1728
0
    }
1729
52
    if (cparams.speed_tier == SpeedTier::kSquirrel &&
1730
52
        cparams.butteraugli_distance >= 3.f) {
1731
0
      return false;
1732
0
    }
1733
52
  }
1734
1735
  // TODO(veluca): handle different values of `buffering`.
1736
52
  if (frame_data.xsize <= 2048 && frame_data.ysize <= 2048) {
1737
52
    return false;
1738
52
  }
1739
0
  if (frame_data.IsJPEG()) {
1740
0
    return false;
1741
0
  }
1742
0
  if (cparams.noise == Override::kOn || cparams.patches == Override::kOn) {
1743
0
    return false;
1744
0
  }
1745
0
  if (cparams.progressive_dc != 0 || frame_info.dc_level != 0) {
1746
0
    return false;
1747
0
  }
1748
0
  if (cparams.resampling != 1 || cparams.ec_resampling != 1) {
1749
0
    return false;
1750
0
  }
1751
0
  if (cparams.max_error_mode) {
1752
0
    return false;
1753
0
  }
1754
0
  if (!cparams.ModularPartIsLossless() || cparams.responsive > 0) {
1755
0
    if (metadata.m.num_extra_channels > 0 || cparams.modular_mode) {
1756
0
      return false;
1757
0
    }
1758
0
  }
1759
0
  ColorTransform ok_color_transform =
1760
0
      cparams.modular_mode ? ColorTransform::kNone : ColorTransform::kXYB;
1761
0
  if (cparams.color_transform != ok_color_transform) {
1762
0
    return false;
1763
0
  }
1764
0
  return true;
1765
0
}
1766
1767
void ComputePermutationForStreaming(size_t xsize, size_t ysize,
1768
                                    size_t group_size, size_t num_passes,
1769
                                    std::vector<coeff_order_t>& permutation,
1770
0
                                    std::vector<size_t>& dc_group_order) {
1771
  // This is only valid in VarDCT mode, otherwise there can be group shift.
1772
0
  const size_t dc_group_size = group_size * kBlockDim;
1773
0
  const size_t group_xsize = DivCeil(xsize, group_size);
1774
0
  const size_t group_ysize = DivCeil(ysize, group_size);
1775
0
  const size_t dc_group_xsize = DivCeil(xsize, dc_group_size);
1776
0
  const size_t dc_group_ysize = DivCeil(ysize, dc_group_size);
1777
0
  const size_t num_groups = group_xsize * group_ysize;
1778
0
  const size_t num_dc_groups = dc_group_xsize * dc_group_ysize;
1779
0
  const size_t num_sections = 2 + num_dc_groups + num_passes * num_groups;
1780
0
  permutation.resize(num_sections);
1781
0
  size_t new_ix = 0;
1782
  // DC Global is first
1783
0
  permutation[0] = new_ix++;
1784
  // TODO(szabadka) Change the dc group order to center-first.
1785
0
  for (size_t dc_y = 0; dc_y < dc_group_ysize; ++dc_y) {
1786
0
    for (size_t dc_x = 0; dc_x < dc_group_xsize; ++dc_x) {
1787
0
      size_t dc_ix = dc_y * dc_group_xsize + dc_x;
1788
0
      dc_group_order.push_back(dc_ix);
1789
0
      permutation[1 + dc_ix] = new_ix++;
1790
0
      size_t ac_y0 = dc_y * kBlockDim;
1791
0
      size_t ac_x0 = dc_x * kBlockDim;
1792
0
      size_t ac_y1 = std::min<size_t>(group_ysize, ac_y0 + kBlockDim);
1793
0
      size_t ac_x1 = std::min<size_t>(group_xsize, ac_x0 + kBlockDim);
1794
0
      for (size_t pass = 0; pass < num_passes; ++pass) {
1795
0
        for (size_t ac_y = ac_y0; ac_y < ac_y1; ++ac_y) {
1796
0
          for (size_t ac_x = ac_x0; ac_x < ac_x1; ++ac_x) {
1797
0
            size_t group_ix = ac_y * group_xsize + ac_x;
1798
0
            size_t old_ix =
1799
0
                AcGroupIndex(pass, group_ix, num_groups, num_dc_groups);
1800
0
            permutation[old_ix] = new_ix++;
1801
0
          }
1802
0
        }
1803
0
      }
1804
0
    }
1805
0
  }
1806
  // AC Global is last
1807
0
  permutation[1 + num_dc_groups] = new_ix++;
1808
0
  JXL_ASSERT(new_ix == num_sections);
1809
0
}
1810
1811
constexpr size_t kGroupSizeOffset[4] = {
1812
    static_cast<size_t>(0),
1813
    static_cast<size_t>(1024),
1814
    static_cast<size_t>(17408),
1815
    static_cast<size_t>(4211712),
1816
};
1817
constexpr size_t kTOCBits[4] = {12, 16, 24, 32};
1818
1819
0
size_t TOCBucket(size_t group_size) {
1820
0
  size_t bucket = 0;
1821
0
  while (bucket < 3 && group_size >= kGroupSizeOffset[bucket + 1]) ++bucket;
1822
0
  return bucket;
1823
0
}
1824
1825
0
size_t TOCSize(const std::vector<size_t>& group_sizes) {
1826
0
  size_t toc_bits = 0;
1827
0
  for (size_t group_size : group_sizes) {
1828
0
    toc_bits += kTOCBits[TOCBucket(group_size)];
1829
0
  }
1830
0
  return (toc_bits + 7) / 8;
1831
0
}
1832
1833
PaddedBytes EncodeTOC(JxlMemoryManager* memory_manager,
1834
0
                      const std::vector<size_t>& group_sizes, AuxOut* aux_out) {
1835
0
  BitWriter writer{memory_manager};
1836
0
  BitWriter::Allotment allotment(&writer, 32 * group_sizes.size());
1837
0
  for (size_t group_size : group_sizes) {
1838
0
    JXL_CHECK(U32Coder::Write(kTocDist, group_size, &writer));
1839
0
  }
1840
0
  writer.ZeroPadToByte();  // before first group
1841
0
  allotment.ReclaimAndCharge(&writer, kLayerTOC, aux_out);
1842
0
  return std::move(writer).TakeBytes();
1843
0
}
1844
1845
void ComputeGroupDataOffset(size_t frame_header_size, size_t dc_global_size,
1846
                            size_t num_sections, size_t& min_dc_global_size,
1847
0
                            size_t& group_offset) {
1848
0
  size_t max_toc_bits = (num_sections - 1) * 32;
1849
0
  size_t min_toc_bits = (num_sections - 1) * 12;
1850
0
  size_t max_padding = (max_toc_bits - min_toc_bits + 7) / 8;
1851
0
  min_dc_global_size = dc_global_size;
1852
0
  size_t dc_global_bucket = TOCBucket(min_dc_global_size);
1853
0
  while (TOCBucket(min_dc_global_size + max_padding) > dc_global_bucket) {
1854
0
    dc_global_bucket = TOCBucket(min_dc_global_size + max_padding);
1855
0
    min_dc_global_size = kGroupSizeOffset[dc_global_bucket];
1856
0
  }
1857
0
  JXL_ASSERT(TOCBucket(min_dc_global_size) == dc_global_bucket);
1858
0
  JXL_ASSERT(TOCBucket(min_dc_global_size + max_padding) == dc_global_bucket);
1859
0
  max_toc_bits += kTOCBits[dc_global_bucket];
1860
0
  size_t max_toc_size = (max_toc_bits + 7) / 8;
1861
0
  group_offset = frame_header_size + max_toc_size + min_dc_global_size;
1862
0
}
1863
1864
size_t ComputeDcGlobalPadding(const std::vector<size_t>& group_sizes,
1865
                              size_t frame_header_size,
1866
                              size_t group_data_offset,
1867
0
                              size_t min_dc_global_size) {
1868
0
  std::vector<size_t> new_group_sizes = group_sizes;
1869
0
  new_group_sizes[0] = min_dc_global_size;
1870
0
  size_t toc_size = TOCSize(new_group_sizes);
1871
0
  size_t actual_offset = frame_header_size + toc_size + group_sizes[0];
1872
0
  return group_data_offset - actual_offset;
1873
0
}
1874
1875
Status OutputGroups(std::vector<std::unique_ptr<BitWriter>>&& group_codes,
1876
                    std::vector<size_t>* group_sizes,
1877
0
                    JxlEncoderOutputProcessorWrapper* output_processor) {
1878
0
  JXL_ASSERT(group_codes.size() >= 4);
1879
0
  {
1880
0
    PaddedBytes dc_group = std::move(*group_codes[1]).TakeBytes();
1881
0
    group_sizes->push_back(dc_group.size());
1882
0
    JXL_RETURN_IF_ERROR(AppendData(*output_processor, dc_group));
1883
0
  }
1884
0
  for (size_t i = 3; i < group_codes.size(); ++i) {
1885
0
    PaddedBytes ac_group = std::move(*group_codes[i]).TakeBytes();
1886
0
    group_sizes->push_back(ac_group.size());
1887
0
    JXL_RETURN_IF_ERROR(AppendData(*output_processor, ac_group));
1888
0
  }
1889
0
  return true;
1890
0
}
1891
1892
void RemoveUnusedHistograms(std::vector<uint8_t>& context_map,
1893
0
                            EntropyEncodingData& codes) {
1894
0
  std::vector<int> remap(256, -1);
1895
0
  std::vector<uint8_t> inv_remap;
1896
0
  for (uint8_t& context : context_map) {
1897
0
    const uint8_t histo_ix = context;
1898
0
    if (remap[histo_ix] == -1) {
1899
0
      remap[histo_ix] = inv_remap.size();
1900
0
      inv_remap.push_back(histo_ix);
1901
0
    }
1902
0
    context = remap[histo_ix];
1903
0
  }
1904
0
  EntropyEncodingData new_codes;
1905
0
  new_codes.use_prefix_code = codes.use_prefix_code;
1906
0
  new_codes.lz77 = codes.lz77;
1907
0
  for (uint8_t histo_idx : inv_remap) {
1908
0
    new_codes.encoding_info.emplace_back(
1909
0
        std::move(codes.encoding_info[histo_idx]));
1910
0
    new_codes.uint_config.emplace_back(codes.uint_config[histo_idx]);
1911
0
    new_codes.encoded_histograms.emplace_back(
1912
0
        std::move(codes.encoded_histograms[histo_idx]));
1913
0
  }
1914
0
  codes = std::move(new_codes);
1915
0
}
1916
1917
Status OutputAcGlobal(PassesEncoderState& enc_state,
1918
                      const FrameDimensions& frame_dim,
1919
                      std::vector<size_t>* group_sizes,
1920
                      JxlEncoderOutputProcessorWrapper* output_processor,
1921
0
                      AuxOut* aux_out) {
1922
0
  JXL_ASSERT(frame_dim.num_groups > 1);
1923
0
  JxlMemoryManager* memory_manager = enc_state.memory_manager();
1924
0
  BitWriter writer{memory_manager};
1925
0
  {
1926
0
    size_t num_histo_bits = CeilLog2Nonzero(frame_dim.num_groups);
1927
0
    BitWriter::Allotment allotment(&writer, num_histo_bits + 1);
1928
0
    writer.Write(1, 1);  // default dequant matrices
1929
0
    writer.Write(num_histo_bits, frame_dim.num_dc_groups - 1);
1930
0
    allotment.ReclaimAndCharge(&writer, kLayerAC, aux_out);
1931
0
  }
1932
0
  const PassesSharedState& shared = enc_state.shared;
1933
0
  for (size_t i = 0; i < enc_state.progressive_splitter.GetNumPasses(); i++) {
1934
    // Encode coefficient orders.
1935
0
    size_t order_bits = 0;
1936
0
    JXL_RETURN_IF_ERROR(
1937
0
        U32Coder::CanEncode(kOrderEnc, enc_state.used_orders[i], &order_bits));
1938
0
    BitWriter::Allotment allotment(&writer, order_bits);
1939
0
    JXL_CHECK(U32Coder::Write(kOrderEnc, enc_state.used_orders[i], &writer));
1940
0
    allotment.ReclaimAndCharge(&writer, kLayerOrder, aux_out);
1941
0
    EncodeCoeffOrders(enc_state.used_orders[i],
1942
0
                      &shared.coeff_orders[i * shared.coeff_order_size],
1943
0
                      &writer, kLayerOrder, aux_out);
1944
    // Fix up context map and entropy codes to remove any fix histograms that
1945
    // were not selected by clustering.
1946
0
    RemoveUnusedHistograms(enc_state.passes[i].context_map,
1947
0
                           enc_state.passes[i].codes);
1948
0
    EncodeHistograms(enc_state.passes[i].context_map, enc_state.passes[i].codes,
1949
0
                     &writer, kLayerAC, aux_out);
1950
0
  }
1951
0
  {
1952
0
    BitWriter::Allotment allotment(&writer, 8);
1953
0
    writer.ZeroPadToByte();  // end of group.
1954
0
    allotment.ReclaimAndCharge(&writer, kLayerAC, aux_out);
1955
0
  }
1956
0
  PaddedBytes ac_global = std::move(writer).TakeBytes();
1957
0
  group_sizes->push_back(ac_global.size());
1958
0
  JXL_RETURN_IF_ERROR(AppendData(*output_processor, ac_global));
1959
0
  return true;
1960
0
}
1961
1962
Status EncodeFrameStreaming(JxlMemoryManager* memory_manager,
1963
                            const CompressParams& cparams,
1964
                            const FrameInfo& frame_info,
1965
                            const CodecMetadata* metadata,
1966
                            JxlEncoderChunkedFrameAdapter& frame_data,
1967
                            const JxlCmsInterface& cms, ThreadPool* pool,
1968
                            JxlEncoderOutputProcessorWrapper* output_processor,
1969
0
                            AuxOut* aux_out) {
1970
0
  PassesEncoderState enc_state{memory_manager};
1971
0
  SetProgressiveMode(cparams, &enc_state.progressive_splitter);
1972
0
  FrameHeader frame_header(metadata);
1973
0
  std::unique_ptr<jpeg::JPEGData> jpeg_data;
1974
0
  if (frame_data.IsJPEG()) {
1975
0
    jpeg_data = make_unique<jpeg::JPEGData>(frame_data.TakeJPEGData());
1976
0
  }
1977
0
  JXL_RETURN_IF_ERROR(MakeFrameHeader(frame_data.xsize, frame_data.ysize,
1978
0
                                      cparams, enc_state.progressive_splitter,
1979
0
                                      frame_info, jpeg_data.get(), true,
1980
0
                                      &frame_header));
1981
0
  const size_t num_passes = enc_state.progressive_splitter.GetNumPasses();
1982
0
  ModularFrameEncoder enc_modular(memory_manager, frame_header, cparams, true);
1983
0
  std::vector<coeff_order_t> permutation;
1984
0
  std::vector<size_t> dc_group_order;
1985
0
  size_t group_size = frame_header.ToFrameDimensions().group_dim;
1986
0
  ComputePermutationForStreaming(frame_data.xsize, frame_data.ysize, group_size,
1987
0
                                 num_passes, permutation, dc_group_order);
1988
0
  enc_state.shared.num_histograms = dc_group_order.size();
1989
0
  size_t dc_group_size = group_size * kBlockDim;
1990
0
  size_t dc_group_xsize = DivCeil(frame_data.xsize, dc_group_size);
1991
0
  size_t min_dc_global_size = 0;
1992
0
  size_t group_data_offset = 0;
1993
0
  PaddedBytes frame_header_bytes{memory_manager};
1994
0
  PaddedBytes dc_global_bytes{memory_manager};
1995
0
  std::vector<size_t> group_sizes;
1996
0
  size_t start_pos = output_processor->CurrentPosition();
1997
0
  for (size_t i = 0; i < dc_group_order.size(); ++i) {
1998
0
    size_t dc_ix = dc_group_order[i];
1999
0
    size_t dc_y = dc_ix / dc_group_xsize;
2000
0
    size_t dc_x = dc_ix % dc_group_xsize;
2001
0
    size_t y0 = dc_y * dc_group_size;
2002
0
    size_t x0 = dc_x * dc_group_size;
2003
0
    size_t ysize = std::min<size_t>(dc_group_size, frame_data.ysize - y0);
2004
0
    size_t xsize = std::min<size_t>(dc_group_size, frame_data.xsize - x0);
2005
0
    size_t group_xsize = DivCeil(xsize, group_size);
2006
0
    size_t group_ysize = DivCeil(ysize, group_size);
2007
0
    JXL_DEBUG_V(2,
2008
0
                "Encoding DC group #%" PRIuS " dc_y = %" PRIuS " dc_x = %" PRIuS
2009
0
                " (x0, y0) = (%" PRIuS ", %" PRIuS ") (xsize, ysize) = (%" PRIuS
2010
0
                ", %" PRIuS ")",
2011
0
                dc_ix, dc_y, dc_x, x0, y0, xsize, ysize);
2012
0
    enc_state.streaming_mode = true;
2013
0
    enc_state.initialize_global_state = (i == 0);
2014
0
    enc_state.dc_group_index = dc_ix;
2015
0
    enc_state.histogram_idx = std::vector<size_t>(group_xsize * group_ysize, i);
2016
0
    std::vector<std::unique_ptr<BitWriter>> group_codes;
2017
0
    JXL_RETURN_IF_ERROR(ComputeEncodingData(
2018
0
        cparams, frame_info, metadata, frame_data, jpeg_data.get(), x0, y0,
2019
0
        xsize, ysize, cms, pool, frame_header, enc_modular, enc_state,
2020
0
        &group_codes, aux_out));
2021
0
    JXL_ASSERT(enc_state.special_frames.empty());
2022
0
    if (i == 0) {
2023
0
      BitWriter writer{memory_manager};
2024
0
      JXL_RETURN_IF_ERROR(WriteFrameHeader(frame_header, &writer, aux_out));
2025
0
      BitWriter::Allotment allotment(&writer, 8);
2026
0
      writer.Write(1, 1);  // write permutation
2027
0
      EncodePermutation(permutation.data(), /*skip=*/0, permutation.size(),
2028
0
                        &writer, kLayerHeader, aux_out);
2029
0
      writer.ZeroPadToByte();
2030
0
      allotment.ReclaimAndCharge(&writer, kLayerHeader, aux_out);
2031
0
      frame_header_bytes = std::move(writer).TakeBytes();
2032
0
      dc_global_bytes = std::move(*group_codes[0]).TakeBytes();
2033
0
      ComputeGroupDataOffset(frame_header_bytes.size(), dc_global_bytes.size(),
2034
0
                             permutation.size(), min_dc_global_size,
2035
0
                             group_data_offset);
2036
0
      JXL_DEBUG_V(2, "Frame header size: %" PRIuS, frame_header_bytes.size());
2037
0
      JXL_DEBUG_V(2, "DC global size: %" PRIuS ", min size for TOC: %" PRIuS,
2038
0
                  dc_global_bytes.size(), min_dc_global_size);
2039
0
      JXL_DEBUG_V(2, "Num groups: %" PRIuS " group data offset: %" PRIuS,
2040
0
                  permutation.size(), group_data_offset);
2041
0
      group_sizes.push_back(dc_global_bytes.size());
2042
0
      output_processor->Seek(start_pos + group_data_offset);
2043
0
    }
2044
0
    JXL_RETURN_IF_ERROR(
2045
0
        OutputGroups(std::move(group_codes), &group_sizes, output_processor));
2046
0
  }
2047
0
  if (frame_header.encoding == FrameEncoding::kVarDCT) {
2048
0
    JXL_RETURN_IF_ERROR(
2049
0
        OutputAcGlobal(enc_state, frame_header.ToFrameDimensions(),
2050
0
                       &group_sizes, output_processor, aux_out));
2051
0
  } else {
2052
0
    group_sizes.push_back(0);
2053
0
  }
2054
0
  JXL_ASSERT(group_sizes.size() == permutation.size());
2055
0
  size_t end_pos = output_processor->CurrentPosition();
2056
0
  output_processor->Seek(start_pos);
2057
0
  size_t padding_size =
2058
0
      ComputeDcGlobalPadding(group_sizes, frame_header_bytes.size(),
2059
0
                             group_data_offset, min_dc_global_size);
2060
0
  group_sizes[0] += padding_size;
2061
0
  PaddedBytes toc_bytes = EncodeTOC(memory_manager, group_sizes, aux_out);
2062
0
  std::vector<uint8_t> padding_bytes(padding_size);
2063
0
  JXL_RETURN_IF_ERROR(AppendData(*output_processor, frame_header_bytes));
2064
0
  JXL_RETURN_IF_ERROR(AppendData(*output_processor, toc_bytes));
2065
0
  JXL_RETURN_IF_ERROR(AppendData(*output_processor, dc_global_bytes));
2066
0
  JXL_RETURN_IF_ERROR(AppendData(*output_processor, padding_bytes));
2067
0
  JXL_DEBUG_V(2, "TOC size: %" PRIuS " padding bytes after DC global: %" PRIuS,
2068
0
              toc_bytes.size(), padding_size);
2069
0
  JXL_ASSERT(output_processor->CurrentPosition() ==
2070
0
             start_pos + group_data_offset);
2071
0
  output_processor->Seek(end_pos);
2072
0
  return true;
2073
0
}
2074
2075
Status EncodeFrameOneShot(JxlMemoryManager* memory_manager,
2076
                          const CompressParams& cparams,
2077
                          const FrameInfo& frame_info,
2078
                          const CodecMetadata* metadata,
2079
                          JxlEncoderChunkedFrameAdapter& frame_data,
2080
                          const JxlCmsInterface& cms, ThreadPool* pool,
2081
                          JxlEncoderOutputProcessorWrapper* output_processor,
2082
52
                          AuxOut* aux_out) {
2083
52
  PassesEncoderState enc_state{memory_manager};
2084
52
  SetProgressiveMode(cparams, &enc_state.progressive_splitter);
2085
52
  FrameHeader frame_header(metadata);
2086
52
  std::unique_ptr<jpeg::JPEGData> jpeg_data;
2087
52
  if (frame_data.IsJPEG()) {
2088
0
    jpeg_data = make_unique<jpeg::JPEGData>(frame_data.TakeJPEGData());
2089
0
  }
2090
52
  JXL_RETURN_IF_ERROR(MakeFrameHeader(frame_data.xsize, frame_data.ysize,
2091
52
                                      cparams, enc_state.progressive_splitter,
2092
52
                                      frame_info, jpeg_data.get(), false,
2093
52
                                      &frame_header));
2094
52
  const size_t num_passes = enc_state.progressive_splitter.GetNumPasses();
2095
52
  ModularFrameEncoder enc_modular(memory_manager, frame_header, cparams, false);
2096
52
  std::vector<std::unique_ptr<BitWriter>> group_codes;
2097
52
  JXL_RETURN_IF_ERROR(ComputeEncodingData(
2098
52
      cparams, frame_info, metadata, frame_data, jpeg_data.get(), 0, 0,
2099
52
      frame_data.xsize, frame_data.ysize, cms, pool, frame_header, enc_modular,
2100
52
      enc_state, &group_codes, aux_out));
2101
2102
52
  BitWriter writer{memory_manager};
2103
52
  writer.AppendByteAligned(enc_state.special_frames);
2104
52
  JXL_RETURN_IF_ERROR(WriteFrameHeader(frame_header, &writer, aux_out));
2105
2106
52
  std::vector<coeff_order_t> permutation;
2107
52
  JXL_RETURN_IF_ERROR(PermuteGroups(cparams, enc_state.shared.frame_dim,
2108
52
                                    num_passes, &permutation, &group_codes));
2109
2110
52
  JXL_RETURN_IF_ERROR(
2111
52
      WriteGroupOffsets(group_codes, permutation, &writer, aux_out));
2112
2113
52
  writer.AppendByteAligned(group_codes);
2114
52
  PaddedBytes frame_bytes = std::move(writer).TakeBytes();
2115
52
  JXL_RETURN_IF_ERROR(AppendData(*output_processor, frame_bytes));
2116
2117
52
  return true;
2118
52
}
2119
2120
}  // namespace
2121
2122
Status EncodeFrame(JxlMemoryManager* memory_manager,
2123
                   const CompressParams& cparams_orig,
2124
                   const FrameInfo& frame_info, const CodecMetadata* metadata,
2125
                   JxlEncoderChunkedFrameAdapter& frame_data,
2126
                   const JxlCmsInterface& cms, ThreadPool* pool,
2127
                   JxlEncoderOutputProcessorWrapper* output_processor,
2128
52
                   AuxOut* aux_out) {
2129
52
  CompressParams cparams = cparams_orig;
2130
52
  if (cparams.speed_tier == SpeedTier::kTectonicPlate &&
2131
52
      !cparams.IsLossless()) {
2132
0
    cparams.speed_tier = SpeedTier::kGlacier;
2133
0
  }
2134
  // Lightning mode is handled externally, so switch to Thunder mode to handle
2135
  // potentially weird cases.
2136
52
  if (cparams.speed_tier == SpeedTier::kLightning) {
2137
0
    cparams.speed_tier = SpeedTier::kThunder;
2138
0
  }
2139
52
  if (cparams.speed_tier == SpeedTier::kTectonicPlate) {
2140
0
    std::vector<CompressParams> all_params;
2141
0
    std::vector<size_t> size;
2142
2143
0
    CompressParams cparams_attempt = cparams_orig;
2144
0
    cparams_attempt.speed_tier = SpeedTier::kGlacier;
2145
0
    cparams_attempt.options.max_properties = 4;
2146
2147
0
    for (float x : {0.0f, 80.f}) {
2148
0
      cparams_attempt.channel_colors_percent = x;
2149
0
      for (float y : {0.0f, 95.0f}) {
2150
0
        cparams_attempt.channel_colors_pre_transform_percent = y;
2151
        // 70000 ensures that the number of palette colors is representable in
2152
        // modular headers.
2153
0
        for (int K : {0, 1 << 10, 70000}) {
2154
0
          cparams_attempt.palette_colors = K;
2155
0
          for (int tree_mode :
2156
0
               {-1, static_cast<int>(ModularOptions::TreeMode::kNoWP),
2157
0
                static_cast<int>(ModularOptions::TreeMode::kDefault)}) {
2158
0
            if (tree_mode == -1) {
2159
              // LZ77 only
2160
0
              cparams_attempt.options.nb_repeats = 0;
2161
0
            } else {
2162
0
              cparams_attempt.options.nb_repeats = 1;
2163
0
              cparams_attempt.options.wp_tree_mode =
2164
0
                  static_cast<ModularOptions::TreeMode>(tree_mode);
2165
0
            }
2166
0
            for (Predictor pred : {Predictor::Zero, Predictor::Variable}) {
2167
0
              cparams_attempt.options.predictor = pred;
2168
0
              for (int g : {0, -1, 3}) {
2169
0
                cparams_attempt.modular_group_size_shift = g;
2170
0
                for (Override patches : {Override::kDefault, Override::kOff}) {
2171
0
                  cparams_attempt.patches = patches;
2172
0
                  all_params.push_back(cparams_attempt);
2173
0
                }
2174
0
              }
2175
0
            }
2176
0
          }
2177
0
        }
2178
0
      }
2179
0
    }
2180
2181
0
    size.resize(all_params.size());
2182
2183
0
    std::atomic<bool> has_error{false};
2184
2185
0
    JXL_RETURN_IF_ERROR(RunOnPool(
2186
0
        pool, 0, all_params.size(), ThreadPool::NoInit,
2187
0
        [&](size_t task, size_t) {
2188
0
          if (has_error) return;
2189
0
          std::vector<uint8_t> output(64);
2190
0
          uint8_t* next_out = output.data();
2191
0
          size_t avail_out = output.size();
2192
0
          JxlEncoderOutputProcessorWrapper local_output(memory_manager);
2193
0
          local_output.SetAvailOut(&next_out, &avail_out);
2194
0
          if (!EncodeFrame(memory_manager, all_params[task], frame_info,
2195
0
                           metadata, frame_data, cms, nullptr, &local_output,
2196
0
                           aux_out)) {
2197
0
            has_error = true;
2198
0
            return;
2199
0
          }
2200
0
          size[task] = local_output.CurrentPosition();
2201
0
        },
2202
0
        "Compress kTectonicPlate"));
2203
0
    if (has_error) return JXL_FAILURE("Compress kTectonicPlate failed");
2204
2205
0
    size_t best_idx = 0;
2206
0
    for (size_t i = 1; i < all_params.size(); i++) {
2207
0
      if (size[best_idx] > size[i]) {
2208
0
        best_idx = i;
2209
0
      }
2210
0
    }
2211
0
    cparams = all_params[best_idx];
2212
0
  }
2213
2214
52
  JXL_RETURN_IF_ERROR(ParamsPostInit(&cparams));
2215
2216
52
  if (cparams.butteraugli_distance < 0) {
2217
0
    return JXL_FAILURE("Expected non-negative distance");
2218
0
  }
2219
2220
52
  if (cparams.progressive_dc < 0) {
2221
52
    if (cparams.progressive_dc != -1) {
2222
0
      return JXL_FAILURE("Invalid progressive DC setting value (%d)",
2223
0
                         cparams.progressive_dc);
2224
0
    }
2225
52
    cparams.progressive_dc = 0;
2226
52
  }
2227
52
  if (cparams.ec_resampling < cparams.resampling) {
2228
0
    cparams.ec_resampling = cparams.resampling;
2229
0
  }
2230
52
  if (cparams.resampling > 1 || frame_info.is_preview) {
2231
0
    cparams.progressive_dc = 0;
2232
0
  }
2233
2234
52
  if (frame_info.dc_level + cparams.progressive_dc > 4) {
2235
0
    return JXL_FAILURE("Too many levels of progressive DC");
2236
0
  }
2237
2238
52
  if (cparams.butteraugli_distance != 0 &&
2239
52
      cparams.butteraugli_distance < kMinButteraugliDistance) {
2240
0
    return JXL_FAILURE("Butteraugli distance is too low (%f)",
2241
0
                       cparams.butteraugli_distance);
2242
0
  }
2243
2244
52
  if (frame_data.IsJPEG()) {
2245
0
    cparams.gaborish = Override::kOff;
2246
0
    cparams.epf = 0;
2247
0
    cparams.modular_mode = false;
2248
0
  }
2249
2250
52
  if (frame_data.xsize == 0 || frame_data.ysize == 0) {
2251
0
    return JXL_FAILURE("Empty image");
2252
0
  }
2253
2254
  // Assert that this metadata is correctly set up for the compression params,
2255
  // this should have been done by enc_file.cc
2256
52
  JXL_ASSERT(metadata->m.xyb_encoded ==
2257
52
             (cparams.color_transform == ColorTransform::kXYB));
2258
2259
52
  if (frame_data.IsJPEG() && cparams.color_transform == ColorTransform::kXYB) {
2260
0
    return JXL_FAILURE("Can't add JPEG frame to XYB codestream");
2261
0
  }
2262
2263
52
  if (CanDoStreamingEncoding(cparams, frame_info, *metadata, frame_data)) {
2264
0
    return EncodeFrameStreaming(memory_manager, cparams, frame_info, metadata,
2265
0
                                frame_data, cms, pool, output_processor,
2266
0
                                aux_out);
2267
52
  } else {
2268
52
    return EncodeFrameOneShot(memory_manager, cparams, frame_info, metadata,
2269
52
                              frame_data, cms, pool, output_processor, aux_out);
2270
52
  }
2271
52
}
2272
2273
Status EncodeFrame(JxlMemoryManager* memory_manager,
2274
                   const CompressParams& cparams_orig,
2275
                   const FrameInfo& frame_info, const CodecMetadata* metadata,
2276
                   const ImageBundle& ib, const JxlCmsInterface& cms,
2277
0
                   ThreadPool* pool, BitWriter* writer, AuxOut* aux_out) {
2278
0
  JxlEncoderChunkedFrameAdapter frame_data(ib.xsize(), ib.ysize(),
2279
0
                                           ib.extra_channels().size());
2280
0
  std::vector<uint8_t> color;
2281
0
  if (ib.IsJPEG()) {
2282
0
    frame_data.SetJPEGData(*ib.jpeg_data);
2283
0
  } else {
2284
0
    uint32_t num_channels =
2285
0
        ib.IsGray() && frame_info.ib_needs_color_transform ? 1 : 3;
2286
0
    size_t stride = ib.xsize() * num_channels * 4;
2287
0
    color.resize(ib.ysize() * stride);
2288
0
    JXL_RETURN_IF_ERROR(ConvertToExternal(
2289
0
        ib, /*bits_per_sample=*/32, /*float_out=*/true, num_channels,
2290
0
        JXL_NATIVE_ENDIAN, stride, pool, color.data(), color.size(),
2291
0
        /*out_callback=*/{}, Orientation::kIdentity));
2292
0
    JxlPixelFormat format{num_channels, JXL_TYPE_FLOAT, JXL_NATIVE_ENDIAN, 0};
2293
0
    frame_data.SetFromBuffer(0, color.data(), color.size(), format);
2294
0
  }
2295
0
  for (size_t ec = 0; ec < ib.extra_channels().size(); ++ec) {
2296
0
    JxlPixelFormat ec_format{1, JXL_TYPE_FLOAT, JXL_NATIVE_ENDIAN, 0};
2297
0
    size_t ec_stride = ib.xsize() * 4;
2298
0
    std::vector<uint8_t> ec_data(ib.ysize() * ec_stride);
2299
0
    const ImageF* channel = &ib.extra_channels()[ec];
2300
0
    JXL_RETURN_IF_ERROR(ConvertChannelsToExternal(
2301
0
        &channel, 1,
2302
0
        /*bits_per_sample=*/32,
2303
0
        /*float_out=*/true, JXL_NATIVE_ENDIAN, ec_stride, pool, ec_data.data(),
2304
0
        ec_data.size(), /*out_callback=*/{}, Orientation::kIdentity));
2305
0
    frame_data.SetFromBuffer(1 + ec, ec_data.data(), ec_data.size(), ec_format);
2306
0
  }
2307
0
  FrameInfo fi = frame_info;
2308
0
  fi.origin = ib.origin;
2309
0
  fi.blend = ib.blend;
2310
0
  fi.blendmode = ib.blendmode;
2311
0
  fi.duration = ib.duration;
2312
0
  fi.timecode = ib.timecode;
2313
0
  fi.name = ib.name;
2314
0
  std::vector<uint8_t> output(64);
2315
0
  uint8_t* next_out = output.data();
2316
0
  size_t avail_out = output.size();
2317
0
  JxlEncoderOutputProcessorWrapper output_processor(memory_manager);
2318
0
  output_processor.SetAvailOut(&next_out, &avail_out);
2319
0
  JXL_RETURN_IF_ERROR(EncodeFrame(memory_manager, cparams_orig, fi, metadata,
2320
0
                                  frame_data, cms, pool, &output_processor,
2321
0
                                  aux_out));
2322
0
  output_processor.SetFinalizedPosition();
2323
0
  output_processor.CopyOutput(output, next_out, avail_out);
2324
0
  writer->AppendByteAligned(Bytes(output));
2325
0
  return true;
2326
0
}
2327
2328
}  // namespace jxl