add way to set distance for extra channels (in particular, alpha) (li…

…bjxl#2117)

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 
 ### Added
+ - encoder API: add `JxlEncoderSetExtraChannelDistance` to adjust the quality
+   of extra channels (like alpha) separately.
 
 ### Removed
 

lib/extras/enc/jxl.cc

@@ -115,6 +115,12 @@ bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
       fprintf(stderr, "JxlEncoderSetFrameBitDepth() failed.\n");
       return false;
     }
+    if (num_alpha_channels != 0 &&
+        JXL_ENC_SUCCESS != JxlEncoderSetExtraChannelDistance(
+                               settings, 0, params.alpha_distance)) {
+      fprintf(stderr, "Setting alpha distance failed.\n");
+      return false;
+    }
     if (lossless &&
         JXL_ENC_SUCCESS != JxlEncoderSetFrameLossless(settings, JXL_TRUE)) {
       fprintf(stderr, "JxlEncoderSetFrameLossless() failed.\n");

lib/extras/enc/jxl.h

@@ -38,6 +38,7 @@ struct JXLCompressParams {
   std::vector<JXLOption> options;
   // Target butteraugli distance, 0.0 means lossless.
   float distance = 1.0f;
+  float alpha_distance = 1.0f;
   // If set to true, forces container mode.
   bool use_container = false;
   // Whether to enable/disable byte-exact jpeg reconstruction for jpeg inputs.

lib/include/jxl/encode.h

@@ -1137,6 +1137,22 @@ JXL_EXPORT JxlEncoderStatus JxlEncoderSetFrameDistance(
 JXL_DEPRECATED JXL_EXPORT JxlEncoderStatus
 JxlEncoderOptionsSetDistance(JxlEncoderFrameSettings*, float);
 
+/**
+ * Sets the distance level for lossy compression of extra channels.
+ * The distance is as in JxlEncoderSetFrameDistance (lower = higher quality).
+ * If not set, or if set to the special value -1, the distance that was set with
+ * JxlEncoderSetFrameDistance will be used.
+ *
+ * @param frame_settings set of options and metadata for this frame. Also
+ * includes reference to the encoder object.
+ * @param index index of the extra channel to set a distance value for.
+ * @param distance the distance value to set.
+ * @return JXL_ENC_SUCCESS if the operation was successful, JXL_ENC_ERROR
+ * otherwise.
+ */
+JXL_EXPORT JxlEncoderStatus JxlEncoderSetExtraChannelDistance(
+    JxlEncoderFrameSettings* frame_settings, size_t index, float distance);
+
 /**
  * Create a new set of encoder options, with all values initially copied from
  * the @p source options, or set to default if @p source is NULL.

lib/jxl/enc_modular.cc

@@ -306,7 +306,7 @@ ModularFrameEncoder::ModularFrameEncoder(const FrameHeader& frame_header,
     : frame_dim_(frame_header.ToFrameDimensions()), cparams_(cparams_orig) {
   size_t num_streams =
       ModularStreamId::Num(frame_dim_, frame_header.passes.num_passes);
-  if (cparams_.IsLossless()) {
+  if (cparams_.ModularPartIsLossless()) {
     switch (cparams_.decoding_speed_tier) {
       case 0:
         break;
@@ -331,7 +331,7 @@ ModularFrameEncoder::ModularFrameEncoder(const FrameHeader& frame_header,
     }
   }
   if (cparams_.decoding_speed_tier >= 1 && cparams_.responsive &&
-      cparams_.IsLossless()) {
+      cparams_.ModularPartIsLossless()) {
     cparams_.options.tree_kind =
         ModularOptions::TreeKind::kTrivialTreeNoPredictor;
     cparams_.options.nb_repeats = 0;
@@ -341,7 +341,7 @@ ModularFrameEncoder::ModularFrameEncoder(const FrameHeader& frame_header,
   // use a sensible default if nothing explicit is specified:
   // Squeeze for lossy, no squeeze for lossless
   if (cparams_.responsive < 0) {
-    if (cparams_.IsLossless()) {
+    if (cparams_.ModularPartIsLossless()) {
       cparams_.responsive = 0;
     } else {
       cparams_.responsive = 1;
@@ -428,7 +428,7 @@ ModularFrameEncoder::ModularFrameEncoder(const FrameHeader& frame_header,
     delta_pred_ = cparams_.options.predictor;
     if (cparams_.lossy_palette) cparams_.options.predictor = Predictor::Zero;
   }
-  if (!cparams_.IsLossless()) {
+  if (!cparams_.ModularPartIsLossless()) {
     if (cparams_.options.predictor == Predictor::Weighted ||
         cparams_.options.predictor == Predictor::Variable ||
         cparams_.options.predictor == Predictor::Best)
@@ -637,8 +637,7 @@ Status ModularFrameEncoder::ComputeEncodingData(
         cparams_.level, max_bitdepth, level_max_bitdepth);
 
   // Set options and apply transformations
-
-  if (cparams_.butteraugli_distance > 0) {
+  if (!cparams_.ModularPartIsLossless()) {
     if (cparams_.palette_colors != 0) {
       JXL_DEBUG_V(3, "Lossy encode, not doing palette transforms");
     }
@@ -752,8 +751,8 @@ Status ModularFrameEncoder::ComputeEncodingData(
   if (cparams_.color_transform == ColorTransform::kNone && do_color &&
       gi.channel.size() - gi.nb_meta_channels >= 3 &&
       max_bitdepth + 1 < level_max_bitdepth) {
-    if (cparams_.colorspace < 0 &&
-        (!cparams_.IsLossless() || cparams_.speed_tier > SpeedTier::kHare)) {
+    if (cparams_.colorspace < 0 && (!cparams_.ModularPartIsLossless() ||
+                                    cparams_.speed_tier > SpeedTier::kHare)) {
       Transform ycocg{TransformId::kRCT};
       ycocg.rct_type = 6;
       ycocg.begin_c = gi.nb_meta_channels;
@@ -785,20 +784,32 @@ Status ModularFrameEncoder::ComputeEncodingData(
 
   std::vector<uint32_t> quants;
 
-  if (cparams_.butteraugli_distance > 0) {
+  if (!cparams_.ModularPartIsLossless()) {
     quants.resize(gi.channel.size(), 1);
-    float quality = 0.25f * cparams_.butteraugli_distance;
-    JXL_DEBUG_V(2,
-                "Adding quantization constants corresponding to distance %.3f ",
-                quality);
+    float quantizer = 0.25f;
     if (!cparams_.responsive) {
       JXL_DEBUG_V(1,
                   "Warning: lossy compression without Squeeze "
                   "transform is just color quantization.");
-      quality *= 0.1f;
+      quantizer *= 0.1f;
     }
+    float bitdepth_correction = 1.f;
     if (cparams_.color_transform != ColorTransform::kXYB) {
-      quality *= maxval / 255.f;
+      bitdepth_correction = maxval / 255.f;
+    }
+    std::vector<float> quantizers;
+    float dist = cparams_.butteraugli_distance;
+    for (size_t i = 0; i < 3; i++) {
+      quantizers.push_back(quantizer * dist * bitdepth_correction);
+    }
+    for (size_t i = 0; i < extra_channels.size(); i++) {
+      int ec_bitdepth =
+          metadata.extra_channel_info[i].bit_depth.bits_per_sample;
+      pixel_type ec_maxval = ec_bitdepth < 32 ? (1u << ec_bitdepth) - 1 : 0;
+      bitdepth_correction = ec_maxval / 255.f;
+      if (i < cparams_.ec_distance.size()) dist = cparams_.ec_distance[i];
+      if (dist < 0) dist = cparams_.butteraugli_distance;
+      quantizers.push_back(quantizer * dist * bitdepth_correction);
     }
     if (cparams_.options.nb_repeats == 0) {
       return JXL_FAILURE("nb_repeats = 0 not supported with modular lossy!");
@@ -817,14 +828,15 @@ Status ModularFrameEncoder::ComputeEncodingData(
         component = 1;
       }
       if (cparams_.color_transform == ColorTransform::kXYB && component < 3) {
-        q = quality * squeeze_quality_factor_xyb *
+        q = quantizers[component] * squeeze_quality_factor_xyb *
             squeeze_xyb_qtable[component][shift];
       } else {
         if (cparams_.colorspace != 0 && component > 0 && component < 3) {
-          q = quality * squeeze_quality_factor * squeeze_chroma_qtable[shift];
+          q = quantizers[component] * squeeze_quality_factor *
+              squeeze_chroma_qtable[shift];
         } else {
-          q = quality * squeeze_quality_factor * squeeze_luma_factor *
-              squeeze_luma_qtable[shift];
+          q = quantizers[component] * squeeze_quality_factor *
+              squeeze_luma_factor * squeeze_luma_qtable[shift];
         }
       }
       if (q < 1) q = 1;

lib/jxl/enc_params.h

@@ -54,6 +54,10 @@ enum class SpeedTier {
 // NOLINTNEXTLINE(clang-analyzer-optin.performance.Padding)
 struct CompressParams {
   float butteraugli_distance = 1.0f;
+
+  // explicit distances for extra channels (defaults to butteraugli_distance
+  // when not set; value of -1 can be used to represent 'default')
+  std::vector<float> ec_distance;
   size_t target_size = 0;
   float target_bitrate = 0.0f;
 
@@ -73,7 +77,6 @@ struct CompressParams {
   // 0 = default.
   // 1 = slightly worse quality.
   // 4 = fastest speed, lowest quality
-  // TODO(veluca): hook this up to the C API.
   size_t decoding_speed_tier = 0;
 
   int max_butteraugli_iters = 4;
@@ -150,17 +153,32 @@ struct CompressParams {
   bool lossy_palette = false;
 
   // Returns whether these params are lossless as defined by SetLossless();
-  bool IsLossless() const {
-    // YCbCr is also considered lossless here since it's intended for
-    // source material that is already YCbCr (we don't do the fwd transform)
-    return modular_mode && butteraugli_distance == 0.0f &&
-           color_transform != jxl::ColorTransform::kXYB;
+  bool IsLossless() const { return modular_mode && ModularPartIsLossless(); }
+
+  bool ModularPartIsLossless() const {
+    if (modular_mode) {
+      // YCbCr is also considered lossless here since it's intended for
+      // source material that is already YCbCr (we don't do the fwd transform)
+      if (butteraugli_distance != 0 ||
+          color_transform == jxl::ColorTransform::kXYB)
+        return false;
+    }
+    for (float f : ec_distance) {
+      if (f > 0) return false;
+      if (f < 0 && butteraugli_distance != 0) return false;
+    }
+    // if no explicit ec_distance given, and using vardct, then the modular part
+    // is empty or not lossless
+    if (!modular_mode && !ec_distance.size()) return false;
+    // all modular channels are encoded at distance 0
+    return true;
   }
 
   // Sets the parameters required to make the codec lossless.
   void SetLossless() {
     modular_mode = true;
     butteraugli_distance = 0.0f;
+    for (float &f : ec_distance) f = 0.0f;
     color_transform = jxl::ColorTransform::kNone;
   }
 

lib/jxl/encode.cc

@@ -982,6 +982,9 @@ JxlEncoderFrameSettings* JxlEncoderFrameSettingsCreate(
     opts->values.lossless = false;
   }
   opts->values.cparams.level = enc->codestream_level;
+  opts->values.cparams.ec_distance.resize(enc->metadata.m.num_extra_channels,
+                                          -1);
+
   JxlEncoderFrameSettings* ret = opts.get();
   enc->encoder_options.emplace_back(std::move(opts));
   return ret;
@@ -1031,6 +1034,33 @@ JxlEncoderStatus JxlEncoderSetFrameDistance(
   return JXL_ENC_SUCCESS;
 }
 
+JxlEncoderStatus JxlEncoderSetExtraChannelDistance(
+    JxlEncoderFrameSettings* frame_settings, size_t index, float distance) {
+  if (index >= frame_settings->enc->metadata.m.num_extra_channels) {
+    return JXL_API_ERROR(frame_settings->enc, JXL_ENC_ERR_API_USAGE,
+                         "Invalid value for the index of extra channel");
+  }
+  if (distance != -1.f && (distance < 0.f || distance > 25.f)) {
+    return JXL_API_ERROR(
+        frame_settings->enc, JXL_ENC_ERR_API_USAGE,
+        "Distance has to be -1 or in [0.0..25.0] (corresponding to "
+        "quality in [0.0..100.0])");
+  }
+  if (distance > 0.f && distance < 0.01f) {
+    distance = 0.01f;
+  }
+
+  if (index >= frame_settings->values.cparams.ec_distance.size()) {
+    // This can only happen if JxlEncoderFrameSettingsCreate() was called before
+    // JxlEncoderSetBasicInfo().
+    frame_settings->values.cparams.ec_distance.resize(
+        frame_settings->enc->metadata.m.num_extra_channels, -1);
+  }
+
+  frame_settings->values.cparams.ec_distance[index] = distance;
+  return JXL_ENC_SUCCESS;
+}
+
 JxlEncoderStatus JxlEncoderOptionsSetDistance(
     JxlEncoderFrameSettings* frame_settings, float distance) {
   // Deprecated function name, call the non-deprecated function

lib/jxl/jxl_test.cc

@@ -882,7 +882,7 @@ TEST(JxlTest, RoundtripAlpha16) {
 
   PackedPixelFile ppf_out;
   // TODO(szabadka) Investigate big size difference on i686
-  EXPECT_NEAR(Roundtrip(t.ppf(), cparams, {}, &pool, &ppf_out), 4107, 200);
+  EXPECT_NEAR(Roundtrip(t.ppf(), cparams, {}, &pool, &ppf_out), 3620, 50);
   EXPECT_THAT(ButteraugliDistance(t.ppf(), ppf_out), IsSlightlyBelow(0.7));
 }
 
@@ -1177,7 +1177,7 @@ TEST(JxlTest, RoundtripAnimationPatches) {
   PackedPixelFile ppf_out;
   // 40k with no patches, 27k with patch frames encoded multiple times.
   EXPECT_THAT(Roundtrip(t.ppf(), cparams, dparams, pool, &ppf_out),
-              IsSlightlyBelow(14400));
+              IsSlightlyBelow(16000));
   EXPECT_EQ(ppf_out.frames.size(), t.ppf().frames.size());
   // >10 with broken patches
   EXPECT_THAT(ButteraugliDistance(t.ppf(), ppf_out), IsSlightlyBelow(1.2));

tools/benchmark/benchmark_codec_jxl.cc

@@ -177,6 +177,9 @@ class JxlCodec : public ImageCodec {
         return JXL_FAILURE("failed to parse uniform quant parameter %s",
                            param.c_str());
       }
+    } else if (param[0] == 'D') {
+      cparams_.ec_distance.clear();
+      cparams_.ec_distance.push_back(strtof(param.substr(1).c_str(), nullptr));
     } else if (param.substr(0, kMaxPassesPrefix.size()) == kMaxPassesPrefix) {
       std::istringstream parser(param.substr(kMaxPassesPrefix.size()));
       parser >> dparams_.max_passes;

tools/cjxl_main.cc

@@ -127,6 +127,15 @@ struct CompressArgs {
         "    Mutually exclusive with --quality.",
         &distance, &ParseFloat);
 
+    opt_alpha_distance_id = cmdline->AddOptionValue(
+        'a', "alpha_distance", "maxError",
+        "Max. butteraugli distance for the alpha channel, lower = higher "
+        "quality.\n"
+        "    0.0 = mathematically lossless. 1.0 = visually lossless.\n"
+        "    Default is to use the same value as for the color image.\n"
+        "    Recommended range: 0.5 .. 3.0. Allowed range: 0.0 ... 25.0.",
+        &alpha_distance, &ParseFloat);
+
     // High-level options
     opt_quality_id = cmdline->AddOptionValue(
         'q', "quality", "QUALITY",
@@ -488,6 +497,7 @@ struct CompressArgs {
   int64_t codestream_level = -1;
   int64_t responsive = -1;
   float distance = 1.0;
+  float alpha_distance = 1.0;
   size_t effort = 7;
   size_t brotli_effort = 9;
   std::string frame_indexing;
@@ -501,6 +511,7 @@ struct CompressArgs {
   CommandLineParser::OptionId opt_lossless_jpeg_id = -1;
   CommandLineParser::OptionId opt_responsive_id = -1;
   CommandLineParser::OptionId opt_distance_id = -1;
+  CommandLineParser::OptionId opt_alpha_distance_id = -1;
   CommandLineParser::OptionId opt_quality_id = -1;
   CommandLineParser::OptionId opt_modular_group_size_id = -1;
 };
@@ -647,6 +658,8 @@ void SetDistanceFromFlags(CommandLineParser* cmdline, CompressArgs* args,
                           jxl::extras::JXLCompressParams* params,
                           const jxl::extras::Codec& codec) {
   bool distance_set = cmdline->GetOption(args->opt_distance_id)->matched();
+  bool alpha_distance_set =
+      cmdline->GetOption(args->opt_alpha_distance_id)->matched();
   bool quality_set = cmdline->GetOption(args->opt_quality_id)->matched();
   if (((distance_set && (args->distance != 0.0)) ||
        (quality_set && (args->quality != 100))) &&
@@ -677,6 +690,8 @@ void SetDistanceFromFlags(CommandLineParser* cmdline, CompressArgs* args,
     args->lossless_jpeg = 0;
   }
   params->distance = args->distance;
+  params->alpha_distance =
+      alpha_distance_set ? args->alpha_distance : params->distance;
 }
 
 void ProcessFlags(const jxl::extras::Codec codec,