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

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/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,21 @@ 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.