From 722295d763b3828a21033c8af1bb8f07a7f61322 Mon Sep 17 00:00:00 2001 From: Luca Versari Date: Sat, 5 Oct 2024 12:40:06 +0200 Subject: [PATCH] Add chroma upsampling stages. --- jxl/src/image.rs | 28 +++++- jxl/src/render/stages/upsample.rs | 140 +++++++++++++++++++++++++++++- 2 files changed, 163 insertions(+), 5 deletions(-) diff --git a/jxl/src/image.rs b/jxl/src/image.rs index 0aa3ea9..51f01ac 100644 --- a/jxl/src/image.rs +++ b/jxl/src/image.rs @@ -37,8 +37,28 @@ pub trait ImageDataType: private::Sealed + Copy + Default + 'static + Debug + Pa fn random(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(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(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) } } diff --git a/jxl/src/render/stages/upsample.rs b/jxl/src/render/stages/upsample.rs index fa40772..d94a62d 100644 --- a/jxl/src/render/stages/upsample.rs +++ b/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; + + 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; + + 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::::new_random(input_size, &mut rng)?]; let stage = NearestNeighbourUpsample::new(0); let output: Vec> = - 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> = 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> = 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(()) + } }