Add chroma upsampling stages.

jxl/src/image.rs

@@ -37,8 +37,28 @@ pub trait ImageDataType: private::Sealed + Copy + Default + 'static + Debug + Pa
  fn random<R: rand::Rng>(rng: &mut R) -> Self;
 }
 
+#[cfg(test)]
+macro_rules! type_min {
+ (f32) => {
+ 0.0f32
+ };
+ ($ty: ty) => {
+ <$ty>::MIN
+ };
+}
+
+#[cfg(test)]
+macro_rules! type_max {
+ (f32) => {
+ 1.0f32
+ };
+ ($ty: ty) => {
+ <$ty>::MAX
+ };
+}
+
 macro_rules! impl_image_data_type {
- ($ty: ty, $id: ident) => {
+ ($ty: ident, $id: ident) => {
  impl private::Sealed for $ty {}
  impl ImageDataType for $ty {
  const DATA_TYPE_ID: DataTypeTag = DataTypeTag::$id;
@@ -51,7 +71,9 @@ macro_rules! impl_image_data_type {
  #[cfg(test)]
  fn random<R: rand::Rng>(rng: &mut R) -> Self {
  use rand::distributions::{Distribution, Uniform};
- Uniform::new(<$ty>::MIN, <$ty>::MAX).sample(rng)
+ let min = type_min!($ty);
+ let max = type_max!($ty);
+ Uniform::new_inclusive(min, max).sample(rng)
  }
  }
  };
@@ -77,7 +99,7 @@ impl ImageDataType for half::f16 {
  #[cfg(test)]
  fn random<R: rand::Rng>(rng: &mut R) -> Self {
  use rand::distributions::{Distribution, Uniform};
- Self::from_f64(Uniform::new(Self::MIN.to_f64(), Self::MAX.to_f64()).sample(rng))
+ Self::from_f64(Uniform::new(0.0f32, 1.0f32).sample(rng) as f64)
  }
 }
 

jxl/src/render/stages/upsample.rs

@@ -48,10 +48,98 @@ impl RenderPipelineStage for NearestNeighbourUpsample {
  }
 }
 
+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};
+ use crate::{error::Result, image::Image, render::test::make_and_run_simple_pipeline};
  use rand::SeedableRng;
  use test_log::test;
 
@@ -72,7 +160,7 @@ mod test {
  let input = vec![Image::<u8>::new_random(input_size, &mut rng)?];
  let stage = NearestNeighbourUpsample::new(0);
  let output: Vec<Image<u8>> =
- crate::render::test::make_and_run_simple_pipeline(stage, &input, image_size, 256)?.1;
+ 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 {
@@ -88,4 +176,52 @@ mod test {
 
  Ok(())
  }
+
+ #[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(())
+ }
 }