[RFC] Add API for setting adaptive quantization parameters

lib/jpegli/encode.cc

@@ -480,11 +480,16 @@ void AllocateBuffers(j_compress_ptr cinfo) {
     m->quant_field.FillRow(0, 0, m->xsize_blocks);
   }
   for (int c = 0; c < cinfo->num_components; ++c) {
-    m->zero_bias_offset[c] =
-        Allocate<float>(cinfo, DCTSIZE2, JPOOL_IMAGE_ALIGNED);
-    m->zero_bias_mul[c] = Allocate<float>(cinfo, DCTSIZE2, JPOOL_IMAGE_ALIGNED);
-    memset(m->zero_bias_mul[c], 0, DCTSIZE2 * sizeof(float));
-    memset(m->zero_bias_offset[c], 0, DCTSIZE2 * sizeof(float));
+    if (m->zero_bias_offset[c] == nullptr) {
+      m->zero_bias_offset[c] =
+          Allocate<float>(cinfo, DCTSIZE2, JPOOL_IMAGE_ALIGNED);
+      memset(m->zero_bias_offset[c], 0, DCTSIZE2 * sizeof(float));
+    }
+    if (m->zero_bias_offset[c] == nullptr) {
+      m->zero_bias_mul[c] =
+          Allocate<float>(cinfo, DCTSIZE2, JPOOL_IMAGE_ALIGNED);
+      memset(m->zero_bias_mul[c], 0, DCTSIZE2 * sizeof(float));
+    }
   }
 }
 
@@ -689,6 +694,11 @@ void jpegli_CreateCompress(j_compress_ptr cinfo, int version,
   cinfo->master->cicp_transfer_function = 2;  // unknown transfer function code
   cinfo->master->use_std_tables = false;
   cinfo->master->use_adaptive_quantization = true;
+  for (int c = 0; c < jpegli::kMaxComponents; ++c) {
+    cinfo->master->zero_bias_offset[c] = nullptr;
+    cinfo->master->zero_bias_mul[c] = nullptr;
+  }
+  cinfo->master->zero_bias_params_set = false;
   cinfo->master->progressive_level = jpegli::kDefaultProgressiveLevel;
   cinfo->master->data_type = JPEGLI_TYPE_UINT8;
   cinfo->master->endianness = JPEGLI_NATIVE_ENDIAN;
@@ -923,6 +933,34 @@ void jpegli_enable_adaptive_quantization(j_compress_ptr cinfo, boolean value) {
   cinfo->master->use_adaptive_quantization = FROM_JXL_BOOL(value);
 }
 
+void jpegli_set_adaptive_quantization_settings(j_compress_ptr cinfo,
+                                               int which_component,
+                                               const float* mul,
+                                               const float* offset) {
+  CheckState(cinfo, jpegli::kEncStart);
+  if (which_component < 0 || which_component >= cinfo->num_components) {
+    JPEGLI_ERROR("Invalid component index %d", which_component);
+  }
+  if (mul == nullptr || offset == nullptr) {
+    JPEGLI_ERROR("Invalid adaptive quantization parameters");
+  }
+  auto* master = cinfo->master;
+  // These are normally allocated in AllocateBuffers but we need to
+  // allocate them early to be abe to set override values.
+  if (master->zero_bias_mul[which_component] == nullptr) {
+    master->zero_bias_mul[which_component] =
+        jpegli::Allocate<float>(cinfo, DCTSIZE2, JPOOL_IMAGE_ALIGNED);
+  }
+  if (master->zero_bias_offset[which_component] == nullptr) {
+    master->zero_bias_offset[which_component] =
+        jpegli::Allocate<float>(cinfo, DCTSIZE2, JPOOL_IMAGE_ALIGNED);
+  }
+  memcpy(master->zero_bias_mul[which_component], mul, DCTSIZE2 * sizeof(float));
+  memcpy(master->zero_bias_offset[which_component], offset,
+         DCTSIZE2 * sizeof(float));
+  master->zero_bias_params_set = true;
+}
+
 void jpegli_simple_progression(j_compress_ptr cinfo) {
   CheckState(cinfo, jpegli::kEncStart);
   jpegli_set_progressive_level(cinfo, 2);

lib/jpegli/encode.h

@@ -145,6 +145,17 @@ void jpegli_set_input_format(j_compress_ptr cinfo, JpegliDataType data_type,
 // Enabled by default.
 void jpegli_enable_adaptive_quantization(j_compress_ptr cinfo, boolean value);
 
+// Sets the adaptive quantization settings, allowing for control of the
+// conditions for when coefficients are zero'd out. Each of the mul and
+// offset parameters consists of an array of 64 values, one for each
+// coefficient.
+// Note that when used, this function must be called once for each
+// component of the image.
+void jpegli_set_adaptive_quantization_settings(j_compress_ptr cinfo,
+                                               int which_component,
+                                               const float* mul,
+                                               const float offset);
+
 // Sets the default progression parameters, where level 0 is sequential, and
 // greater level value means more progression steps. Default is 2.
 void jpegli_set_progressive_level(j_compress_ptr cinfo, int level);

lib/jpegli/encode_internal.h

@@ -90,6 +90,7 @@ struct jpeg_comp_master {
   float* quant_mul[jpegli::kMaxComponents];
   float* zero_bias_offset[jpegli::kMaxComponents];
   float* zero_bias_mul[jpegli::kMaxComponents];
+  bool zero_bias_params_set;
   int h_factor[jpegli::kMaxComponents];
   int v_factor[jpegli::kMaxComponents];
   // Array of Huffman tables that will be encoded in one or more DHT segments.

lib/jpegli/quant.cc

@@ -732,7 +732,7 @@ void InitQuantizer(j_compress_ptr cinfo, QuantPass pass) {
       }
     }
   }
-  if (m->use_adaptive_quantization) {
+  if (m->use_adaptive_quantization && !m->zero_bias_params_set) {
     for (int c = 0; c < cinfo->num_components; ++c) {
       for (int k = 0; k < DCTSIZE2; ++k) {
         m->zero_bias_mul[c][k] = k == 0 ? 0.0f : 0.5f;