XYB to linear sRGB stage

jxl/src/headers/transform_data.rs

@@ -21,9 +21,9 @@ pub struct OpsinInverseMatrix {
     #[default([11.031566901960783, -9.866943921568629, -0.16462299647058826,
                -3.254147380392157,  4.418770392156863,  -0.16462299647058826,
                -3.6588512862745097, 2.7129230470588235, 1.9459282392156863])]
-    inverse_matrix: [f32; 9],
-    #[default([0.0037930732552754493, 0.0037930732552754493, 0.0037930732552754493])]
-    opsin_biases: [f32; 3],
+    pub inverse_matrix: [f32; 9],
+    #[default([-0.0037930732552754493; 3])]
+    pub opsin_biases: [f32; 3],
     #[default([1.0 - 0.05465007330715401, 1.0 - 0.07005449891748593, 1.0 - 0.049935103337343655, 0.145])]
     quant_biases: [f32; 4],
 }
@@ -326,7 +326,7 @@ pub struct CustomTransformData {
     all_default: bool,
     #[condition(nonserialized.xyb_encoded)]
     #[default(OpsinInverseMatrix::default())]
-    opsin_inverse_matrix: OpsinInverseMatrix,
+    pub opsin_inverse_matrix: OpsinInverseMatrix,
     #[default(0)]
     #[coder(Bits(3))]
     custom_weight_mask: u32,

jxl/src/render/stages/mod.rs

@@ -9,6 +9,7 @@ mod nearest_neighbor;
 mod noise;
 mod save;
 mod upsample;
+mod xyb;
 
 pub use chroma_upsample::*;
 pub use convert::*;

jxl/src/render/stages/xyb.rs

@@ -0,0 +1,145 @@
+// Copyright (c) the JPEG XL Project Authors. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+use crate::headers::OpsinInverseMatrix;
+use crate::render::{RenderPipelineInPlaceStage, RenderPipelineStage};
+
+/// Convert XYB to linear sRGB, where 1.0 corresponds to `intensity_target` nits.
+pub struct XybToLinearSrgbStage {
+    first_channel: usize,
+    opsin: OpsinInverseMatrix,
+    intensity_target: f32,
+}
+
+impl XybToLinearSrgbStage {
+    // TODO(tirr-c): remove once we use this!
+    #[allow(unused)]
+    pub fn new(first_channel: usize, opsin: OpsinInverseMatrix, intensity_target: f32) -> Self {
+        Self {
+            first_channel,
+            opsin,
+            intensity_target,
+        }
+    }
+}
+
+impl std::fmt::Display for XybToLinearSrgbStage {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let channel = self.first_channel;
+        write!(
+            f,
+            "XYB to linear sRGB for channel [{},{},{}]",
+            channel,
+            channel + 1,
+            channel + 2
+        )
+    }
+}
+
+impl RenderPipelineStage for XybToLinearSrgbStage {
+    type Type = RenderPipelineInPlaceStage<f32>;
+
+    fn uses_channel(&self, c: usize) -> bool {
+        (self.first_channel..self.first_channel + 3).contains(&c)
+    }
+
+    fn process_row_chunk(
+        &mut self,
+        _position: (usize, usize),
+        xsize: usize,
+        row: &mut [&mut [f32]],
+    ) {
+        let [row_x, row_y, row_b] = row else {
+            panic!(
+                "incorrect number of channels; expected 3, found {}",
+                row.len()
+            );
+        };
+
+        let OpsinInverseMatrix {
+            inverse_matrix: mat,
+            opsin_biases: bias,
+            ..
+        } = self.opsin;
+        let bias_cbrt = bias.map(|x| x.cbrt());
+        let intensity_scale = 255.0 / self.intensity_target;
+
+        for idx in 0..xsize {
+            let x = row_x[idx];
+            let y = row_y[idx];
+            let b = row_b[idx];
+
+            // Mix and apply bias
+            let l = y + x - bias_cbrt[0];
+            let m = y - x - bias_cbrt[1];
+            let s = b - bias_cbrt[2];
+
+            // Apply biased inverse gamma and scale (1.0 corresponds to `intensity_target` nits)
+            let l = (l * l * l + bias[0]) * intensity_scale;
+            let m = (m * m * m + bias[1]) * intensity_scale;
+            let s = (s * s * s + bias[2]) * intensity_scale;
+
+            // Apply opsin inverse matrix (linear LMS to linear sRGB)
+            let [r, g, b] = crate::util::matmul3_vec(mat, [l, m, s]);
+            row_x[idx] = r;
+            row_y[idx] = g;
+            row_b[idx] = b;
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use test_log::test;
+
+    use super::*;
+    use crate::error::Result;
+    use crate::image::Image;
+    use crate::render::test::make_and_run_simple_pipeline;
+    use crate::util::test::assert_all_almost_eq;
+
+    #[test]
+    fn consistency() -> Result<()> {
+        crate::render::test::test_stage_consistency::<_, f32, f32>(
+            XybToLinearSrgbStage::new(0, OpsinInverseMatrix::default(), 255.0),
+            (500, 500),
+            3,
+        )
+    }
+
+    #[test]
+    fn srgb_primaries() -> Result<()> {
+        let mut input_x = Image::new((3, 1))?;
+        let mut input_y = Image::new((3, 1))?;
+        let mut input_b = Image::new((3, 1))?;
+        input_x
+            .as_rect_mut()
+            .row(0)
+            .copy_from_slice(&[0.028100073, -0.015386105, 0.0]);
+        input_y
+            .as_rect_mut()
+            .row(0)
+            .copy_from_slice(&[0.4881882, 0.71478134, 0.2781282]);
+        input_b
+            .as_rect_mut()
+            .row(0)
+            .copy_from_slice(&[0.471659, 0.43707693, 0.66613984]);
+
+        let stage = XybToLinearSrgbStage::new(0, OpsinInverseMatrix::default(), 255.0);
+        let output = make_and_run_simple_pipeline::<_, f32, f32>(
+            stage,
+            &[input_x, input_y, input_b],
+            (3, 1),
+            256,
+        )?
+        .1;
+
+        assert_all_almost_eq!(output[0].as_rect().row(0), &[1.0, 0.0, 0.0], 1e-6);
+        assert_all_almost_eq!(output[1].as_rect().row(0), &[0.0, 1.0, 0.0], 1e-6);
+        assert_all_almost_eq!(output[2].as_rect().row(0), &[0.0, 0.0, 1.0], 1e-6);
+
+        Ok(())
+    }
+}

jxl/src/util.rs

@@ -8,13 +8,15 @@ pub mod test;
 
 mod bits;
 mod concat_slice;
+mod linalg;
 mod log2;
 mod shift_right_ceil;
 pub mod tracing_wrappers;
 mod vec_helpers;
 
 pub use bits::*;
 pub use concat_slice::*;
+pub use linalg::*;
 pub use log2::*;
 pub use shift_right_ceil::*;
 pub use vec_helpers::*;

jxl/src/util/linalg.rs

@@ -0,0 +1,12 @@
+// Copyright (c) the JPEG XL Project Authors. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+pub fn matmul3_vec(m: [f32; 9], v: [f32; 3]) -> [f32; 3] {
+    [
+        v[0] * m[0] + v[1] * m[1] + v[2] * m[2],
+        v[0] * m[3] + v[1] * m[4] + v[2] * m[5],
+        v[0] * m[6] + v[1] * m[7] + v[2] * m[8],
+    ]
+}