make kXYB images work, and make color_space=XYB do something better

lib/include/jxl/color_encoding.h

@@ -28,7 +28,10 @@ typedef enum {
  * value implies that num_color_channels in JxlBasicInfo is 1, any other value
  * implies num_color_channels is 3. */
  JXL_COLOR_SPACE_GRAY,
- /** XYB (opsin) color space */
+ /** XYB (opsin) color space. When used with floating point sample value types,
+ * this is unscaled XYB as it is used internally in JPEG XL. When used with
+ * uint sample value types, this is XYB scaled to suitable ranges for that.
+ * The ICC profile generated for this color space applies to scaled XYB. */
  JXL_COLOR_SPACE_XYB,
  /** None of the other table entries describe the color space appropriately */
  JXL_COLOR_SPACE_UNKNOWN,

lib/jxl/color_encoding_internal.cc

@@ -545,6 +545,7 @@ Status ColorEncoding::VisitFields(Visitor* JXL_RESTRICT visitor) {
  }
  }
 
+ tf.nonserialized_color_space = color_space_;
  JXL_QUIET_RETURN_IF_ERROR(visitor->VisitNested(&tf));
 
  JXL_QUIET_RETURN_IF_ERROR(

lib/jxl/color_management.cc

@@ -401,6 +401,14 @@ Status CreateICCCurvParaTag(std::vector<float> params, size_t curve_type,
  return true;
 }
 
+// TODO(jon): This produces an ICC profile corresponding to scaled XYB
+// That is OK if the output will also be scaled XYB, e.g. when outputting
+// to PNG. When outputting to PFM it will be wrong but it doesn't matter
+// since PFM cannot store the ICC profile anyway. But in general the
+// ICC profile returned here should correspond to the pixel data that
+// is produced, so if it's unscaled XYB, this should be a different
+// ICC profile. (I'm not sure if such an ICC profile can even be made,
+// there might be an implicit assumption that samples are in 0..1 range)
 Status CreateICCLutAtoBTagForXYB(PaddedBytes* JXL_RESTRICT tags) {
  WriteICCTag("mAB ", tags->size(), tags);
  // 4 reserved bytes set to 0

lib/jxl/dec_cache.cc

@@ -154,7 +154,9 @@ Status PassesDecoderState::PreparePipeline(ImageBundle* decoded,
  if (frame_header.color_transform == ColorTransform::kYCbCr) {
  builder.AddStage(GetYCbCrStage());
  } else if (frame_header.color_transform == ColorTransform::kXYB) {
- builder.AddStage(GetXYBStage(output_encoding_info));
+ JxlDataType ptype = main_output.format.data_type;
+ bool scaled_xyb = (ptype == JXL_TYPE_UINT8 || ptype == JXL_TYPE_UINT16);
+ builder.AddStage(GetXYBStage(output_encoding_info, scaled_xyb));
  if (output_encoding_info.color_encoding.GetColorSpace() !=
  ColorSpace::kXYB) {
  linear = true;

lib/jxl/decode_test.cc

@@ -1713,22 +1713,28 @@ double ButteraugliDistance(size_t xsize, size_t ysize,
  float intensity_out) {
  jxl::CodecInOut in;
  in.metadata.m.color_encoding = color_in;
+ if (color_in.GetColorSpace() == jxl::ColorSpace::kXYB) {
+ in.metadata.m.color_encoding.SetColorSpace(jxl::ColorSpace::kRGB);
+ }
  in.metadata.m.SetIntensityTarget(intensity_in);
  JxlPixelFormat format_in = {static_cast<uint32_t>(color_in.Channels()),
  JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
  EXPECT_TRUE(jxl::ConvertFromExternal(
  jxl::Span<const uint8_t>(pixels_in.data(), pixels_in.size()), xsize,
- ysize, color_in,
+ ysize, in.metadata.m.color_encoding,
  /*bits_per_sample=*/16, format_in,
  /*pool=*/nullptr, &in.Main()));
  jxl::CodecInOut out;
  out.metadata.m.color_encoding = color_out;
+ if (color_out.GetColorSpace() == jxl::ColorSpace::kXYB) {
+ out.metadata.m.color_encoding.SetColorSpace(jxl::ColorSpace::kRGB);
+ }
  out.metadata.m.SetIntensityTarget(intensity_out);
  JxlPixelFormat format_out = {static_cast<uint32_t>(color_out.Channels()),
  JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
  EXPECT_TRUE(jxl::ConvertFromExternal(
  jxl::Span<const uint8_t>(pixels_out.data(), pixels_out.size()), xsize,
- ysize, color_out,
+ ysize, out.metadata.m.color_encoding,
  /*bits_per_sample=*/16, format_out,
  /*pool=*/nullptr, &out.Main()));
  return ButteraugliDistance(in.frames, out.frames, jxl::ButteraugliParams(),

lib/jxl/enc_image_bundle.cc

@@ -37,6 +37,16 @@ Status CopyToT(const ImageMetadata* metadata, const ImageBundle* ib,
  } else {
  out->ShrinkTo(rect.xsize(), rect.ysize());
  }
+
+ const ColorEncoding& c_linear_srgb = ColorEncoding::LinearSRGB(ib->IsGray());
+ if (ib->c_current().GetColorSpace() == ColorSpace::kXYB &&
+ c_desired.SameColorEncoding(c_linear_srgb)) {
+ OpsinParams opsin_params;
+ opsin_params.Init(ib->metadata()->IntensityTarget());
+ OpsinToLinear(ib->color(), rect, pool, out, opsin_params);
+ return true;
+ }
+
  std::atomic<bool> ok{true};
  JXL_RETURN_IF_ERROR(RunOnPool(
  pool, 0, rect.ysize(),

lib/jxl/enc_xyb.cc

@@ -20,6 +20,7 @@
 #include "lib/jxl/base/status.h"
 #include "lib/jxl/color_encoding_internal.h"
 #include "lib/jxl/color_management.h"
+#include "lib/jxl/dec_xyb.h"
 #include "lib/jxl/enc_bit_writer.h"
 #include "lib/jxl/enc_color_management.h"
 #include "lib/jxl/enc_image_bundle.h"
@@ -323,8 +324,22 @@ const ImageBundle* ToXYB(const ImageBundle& in, ThreadPool* pool,
  ComputePremulAbsorb(in.metadata()->IntensityTarget(), premul_absorb);
 
  const bool want_linear = linear != nullptr;
-
  const ColorEncoding& c_linear_srgb = ColorEncoding::LinearSRGB(in.IsGray());
+
+ // If the input is already XYB, just copy it
+ if (in.c_current().GetColorSpace() == ColorSpace::kXYB) {
+ *xyb = CopyImage(in.color());
+ if (want_linear) {
+ // and convert it to linear sRGB for the slow encode settings
+ OpsinParams opsin_params;
+ opsin_params.Init(in.metadata()->IntensityTarget());
+ linear->SetFromImage(Image3F(xsize, ysize), c_linear_srgb);
+ OpsinToLinear(in.color(), Rect(in.color()), pool, linear->color(),
+ opsin_params);
+ return linear;
+ }
+ return ∈
+ }
  // Linear sRGB inputs are rare but can be useful for the fastest encoders, for
  // which undoing the sRGB transfer function would be a large part of the cost.
  if (c_linear_srgb.SameColorEncoding(in.c_current())) {

lib/jxl/render_pipeline/stage_xyb.cc

@@ -20,11 +20,13 @@ namespace HWY_NAMESPACE {
 
 class XYBStage : public RenderPipelineStage {
  public:
- explicit XYBStage(const OutputEncodingInfo& output_encoding_info)
+ explicit XYBStage(const OutputEncodingInfo& output_encoding_info,
+ bool scaled_xyb)
  : RenderPipelineStage(RenderPipelineStage::Settings()),
  opsin_params_(output_encoding_info.opsin_params),
  output_is_xyb_(output_encoding_info.color_encoding.GetColorSpace() ==
- ColorSpace::kXYB) {}
+ ColorSpace::kXYB),
+ scaled_xyb_(scaled_xyb) {}
 
  void ProcessRow(const RowInfo& input_rows, const RowInfo& output_rows,
  size_t xextra, size_t xsize, size_t xpos, size_t ypos,
@@ -45,7 +47,8 @@ class XYBStage : public RenderPipelineStage {
  msan::UnpoisonMemory(row2 + xsize, sizeof(float) * (xsize_v - xsize));
  // TODO(eustas): when using frame origin, addresses might be unaligned;
  // making them aligned will void performance penalty.
- if (output_is_xyb_) {
+ if (output_is_xyb_ && scaled_xyb_) {
+ // scale XYB to fit in uint types (nominal range of 0..1)
  const auto scale_x = Set(d, kScaledXYBScale[0]);
  const auto scale_y = Set(d, kScaledXYBScale[1]);
  const auto scale_bmy = Set(d, kScaledXYBScale[2]);
@@ -63,6 +66,8 @@ class XYBStage : public RenderPipelineStage {
  StoreU(out_y, d, row1 + x);
  StoreU(out_b, d, row2 + x);
  }
+ } else if (output_is_xyb_) {
+ // no need to do anything, just keep XYB as it is
  } else {
  for (ssize_t x = -xextra; x < (ssize_t)(xsize + xextra); x += Lanes(d)) {
  const auto in_opsin_x = LoadU(d, row0 + x);
@@ -93,11 +98,12 @@ class XYBStage : public RenderPipelineStage {
  private:
  const OpsinParams opsin_params_;
  const bool output_is_xyb_;
+ const bool scaled_xyb_;
 };
 
 std::unique_ptr<RenderPipelineStage> GetXYBStage(
- const OutputEncodingInfo& output_encoding_info) {
- return jxl::make_unique<XYBStage>(output_encoding_info);
+ const OutputEncodingInfo& output_encoding_info, bool scaled_xyb) {
+ return jxl::make_unique<XYBStage>(output_encoding_info, scaled_xyb);
 }
 
 // NOLINTNEXTLINE(google-readability-namespace-comments)
@@ -111,8 +117,8 @@ namespace jxl {
 HWY_EXPORT(GetXYBStage);
 
 std::unique_ptr<RenderPipelineStage> GetXYBStage(
- const OutputEncodingInfo& output_encoding_info) {
- return HWY_DYNAMIC_DISPATCH(GetXYBStage)(output_encoding_info);
+ const OutputEncodingInfo& output_encoding_info, bool scaled_xyb) {
+ return HWY_DYNAMIC_DISPATCH(GetXYBStage)(output_encoding_info, scaled_xyb);
 }
 
 namespace {

lib/jxl/render_pipeline/stage_xyb.h

@@ -14,7 +14,7 @@ namespace jxl {
 
 // Converts the color channels from XYB to linear with appropriate primaries.
 std::unique_ptr<RenderPipelineStage> GetXYBStage(
- const OutputEncodingInfo& output_encoding_info);
+ const OutputEncodingInfo& output_encoding_info, bool scaled_xyb);
 
 // Gets a stage to convert with fixed point arithmetic from XYB to sRGB8 and
 // write to a uint8 buffer.