upsample stage (2x) (#53)

jxl/src/headers/transform_data.rs

@@ -7,7 +7,7 @@
 
 use crate::{bit_reader::BitReader, error::Error, headers::encodings::*};
 use jxl_macros::UnconditionalCoder;
-
+#[derive(Default)]
 pub struct CustomTransformDataNonserialized {
     pub xyb_encoded: bool,
 }
@@ -332,11 +332,11 @@ pub struct CustomTransformData {
     custom_weight_mask: u32,
     #[condition((custom_weight_mask & 1) != 0)]
     #[default(DEFAULT_KERN_2)]
-    weights2: [f32; 15],
+    pub weights2: [f32; 15],
     #[condition((custom_weight_mask & 2) != 0)]
     #[default(DEFAULT_KERN_4)]
-    weights4: [f32; 55],
+    pub weights4: [f32; 55],
     #[condition((custom_weight_mask & 4) != 0)]
     #[default(DEFAULT_KERN_8)]
-    weights8: [f32; 210],
+    pub weights8: [f32; 210],
 }

jxl/src/image.rs

@@ -184,6 +184,13 @@ impl<T: ImageDataType> Image<T> {
         Ok(img)
     }
 
+    #[cfg(test)]
+    pub fn new_constant(size: (usize, usize), val: T) -> Result<Image<T>> {
+        let mut img = Self::new(size)?;
+        img.data.iter_mut().for_each(|x| *x = val);
+        Ok(img)
+    }
+
     pub fn size(&self) -> (usize, usize) {
         self.size
     }

jxl/src/render/mod.rs

@@ -88,6 +88,7 @@ pub trait RenderPipelineStage: Any + std::fmt::Display {
         &mut self,
         position: (usize, usize),
         xsize: usize,
+        // one for each channel
         row: &mut [<Self::Type as RenderPipelineStageInfo>::RowType<'_>],
     );
 
@@ -111,7 +112,6 @@ pub trait RenderPipelineBuilder: Sized {
     fn build(self) -> Result<Self::RenderPipeline>;
 }
 
-#[allow(dead_code)]
 pub struct GroupFillInfo<F> {
     group_id: usize,
     num_filled_passes: usize,

jxl/src/render/stages/chroma_upsample.rs

@@ -0,0 +1,149 @@
+// 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::render::{RenderPipelineInOutStage, RenderPipelineStage};
+
+pub struct HorizontalChromaUpsample {
+    channel: usize,
+}
+
+impl HorizontalChromaUpsample {
+    pub fn new(channel: usize) -> HorizontalChromaUpsample {
+        HorizontalChromaUpsample { channel }
+    }
+}
+
+impl std::fmt::Display for HorizontalChromaUpsample {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "chroma upsample of channel {}, horizontally",
+            self.channel
+        )
+    }
+}
+
+impl RenderPipelineStage for HorizontalChromaUpsample {
+    type Type = RenderPipelineInOutStage<f32, f32, 1, 0, 1, 0>;
+
+    fn uses_channel(&self, c: usize) -> bool {
+        c == self.channel
+    }
+
+    fn process_row_chunk(
+        &mut self,
+        _position: (usize, usize),
+        xsize: usize,
+        row: &mut [(&[&[f32]], &mut [&mut [f32]])],
+    ) {
+        let (input, output) = &mut row[0];
+        for i in 0..xsize {
+            let scaled_cur = input[0][i + 1] * 0.75;
+            let prev = input[0][i];
+            let next = input[0][i + 2];
+            let left = 0.25 * prev + scaled_cur;
+            let right = 0.25 * next + scaled_cur;
+            output[0][2 * i] = left;
+            output[0][2 * i + 1] = right;
+        }
+    }
+}
+
+pub struct VerticalChromaUpsample {
+    channel: usize,
+}
+
+impl VerticalChromaUpsample {
+    pub fn new(channel: usize) -> VerticalChromaUpsample {
+        VerticalChromaUpsample { channel }
+    }
+}
+
+impl std::fmt::Display for VerticalChromaUpsample {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "chroma upsample of channel {}, vertically", self.channel)
+    }
+}
+
+impl RenderPipelineStage for VerticalChromaUpsample {
+    type Type = RenderPipelineInOutStage<f32, f32, 0, 1, 0, 1>;
+
+    fn uses_channel(&self, c: usize) -> bool {
+        c == self.channel
+    }
+
+    fn process_row_chunk(
+        &mut self,
+        _position: (usize, usize),
+        xsize: usize,
+        row: &mut [(&[&[f32]], &mut [&mut [f32]])],
+    ) {
+        let (input, output) = &mut row[0];
+        for i in 0..xsize {
+            let scaled_cur = input[1][i] * 0.75;
+            let prev = input[0][i];
+            let next = input[2][i];
+            let up = 0.25 * prev + scaled_cur;
+            let down = 0.25 * next + scaled_cur;
+            output[0][i] = up;
+            output[1][i] = down;
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::{error::Result, image::Image, render::test::make_and_run_simple_pipeline};
+    use test_log::test;
+
+    #[test]
+    fn hchr_consistency() -> Result<()> {
+        crate::render::test::test_stage_consistency::<_, f32, f32>(
+            HorizontalChromaUpsample::new(0),
+            (500, 500),
+            1,
+        )
+    }
+
+    #[test]
+    fn test_hchr() -> Result<()> {
+        let mut input = Image::new((3, 1))?;
+        input
+            .as_rect_mut()
+            .row(0)
+            .copy_from_slice(&[1.0f32, 2.0, 4.0]);
+        let stage = HorizontalChromaUpsample::new(0);
+        let output: Vec<Image<f32>> = make_and_run_simple_pipeline(stage, &[input], (6, 1), 256)?.1;
+        assert_eq!(output[0].as_rect().row(0), [1.0, 1.25, 1.75, 2.5, 3.5, 4.0]);
+        Ok(())
+    }
+
+    #[test]
+    fn vchr_consistency() -> Result<()> {
+        crate::render::test::test_stage_consistency::<_, f32, f32>(
+            VerticalChromaUpsample::new(0),
+            (500, 500),
+            1,
+        )
+    }
+
+    #[test]
+    fn test_vchr() -> Result<()> {
+        let mut input = Image::new((1, 3))?;
+        input.as_rect_mut().row(0)[0] = 1.0f32;
+        input.as_rect_mut().row(1)[0] = 2.0f32;
+        input.as_rect_mut().row(2)[0] = 4.0f32;
+        let stage = VerticalChromaUpsample::new(0);
+        let output: Vec<Image<f32>> = make_and_run_simple_pipeline(stage, &[input], (1, 6), 256)?.1;
+        assert_eq!(output[0].as_rect().row(0)[0], 1.0);
+        assert_eq!(output[0].as_rect().row(1)[0], 1.25);
+        assert_eq!(output[0].as_rect().row(2)[0], 1.75);
+        assert_eq!(output[0].as_rect().row(3)[0], 2.5);
+        assert_eq!(output[0].as_rect().row(4)[0], 3.5);
+        assert_eq!(output[0].as_rect().row(5)[0], 4.0);
+        Ok(())
+    }
+}

jxl/src/render/stages/mod.rs

@@ -3,10 +3,12 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+mod chroma_upsample;
 mod convert;
+mod nearest_neighbor;
 mod save;
 mod upsample;
 
+pub use chroma_upsample::*;
 pub use convert::*;
 pub use save::*;
-pub use upsample::*;

jxl/src/render/stages/nearest_neighbor.rs

@@ -0,0 +1,91 @@
+// 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::render::{RenderPipelineInOutStage, RenderPipelineStage};
+pub struct NearestNeighbourUpsample {
+    channel: usize,
+}
+
+impl NearestNeighbourUpsample {
+    #[allow(dead_code)]
+    pub fn new(channel: usize) -> NearestNeighbourUpsample {
+        NearestNeighbourUpsample { channel }
+    }
+}
+
+impl std::fmt::Display for NearestNeighbourUpsample {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "2x2 nearest neighbour upsample of channel {}",
+            self.channel
+        )
+    }
+}
+
+impl RenderPipelineStage for NearestNeighbourUpsample {
+    type Type = RenderPipelineInOutStage<u8, u8, 0, 0, 1, 1>;
+
+    fn uses_channel(&self, c: usize) -> bool {
+        c == self.channel
+    }
+
+    fn process_row_chunk(
+        &mut self,
+        _position: (usize, usize),
+        xsize: usize,
+        row: &mut [(&[&[u8]], &mut [&mut [u8]])],
+    ) {
+        let (input, output) = &mut row[0];
+        for i in 0..xsize {
+            output[0][i * 2] = input[0][i];
+            output[0][i * 2 + 1] = input[0][i];
+            output[1][i * 2] = input[0][i];
+            output[1][i * 2 + 1] = input[0][i];
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::{error::Result, image::Image, render::test::make_and_run_simple_pipeline};
+    use rand::SeedableRng;
+    use test_log::test;
+
+    #[test]
+    fn nn_consistency() -> Result<()> {
+        crate::render::test::test_stage_consistency::<_, u8, u8>(
+            NearestNeighbourUpsample::new(0),
+            (500, 500),
+            1,
+        )
+    }
+
+    #[test]
+    fn test_nn() -> Result<()> {
+        let image_size = (500, 400);
+        let input_size = (image_size.0 / 2, image_size.1 / 2);
+        let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(0);
+        let input = vec![Image::<u8>::new_random(input_size, &mut rng)?];
+        let stage = NearestNeighbourUpsample::new(0);
+        let output: Vec<Image<u8>> =
+            make_and_run_simple_pipeline(stage, &input, image_size, 256)?.1;
+        assert_eq!(image_size, output[0].size());
+        for y in 0..image_size.1 {
+            for x in 0..image_size.0 {
+                let ix = x / 2;
+                let iy = y / 2;
+                let i = input[0].as_rect().row(iy)[ix];
+                let o = output[0].as_rect().row(y)[x];
+                if i != o {
+                    panic!("Mismatch at output position {x}x{y}: {i} vs output {o}");
+                }
+            }
+        }
+
+        Ok(())
+    }
+}