Skip to content

Commit

Permalink
Add chroma upsampling stages.
Browse files Browse the repository at this point in the history
  • Loading branch information
veluca93 committed Oct 5, 2024
1 parent 7229eb7 commit 722295d
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 5 deletions.
28 changes: 25 additions & 3 deletions jxl/src/image.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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)
}
}
};
Expand All @@ -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)
}
}

Expand Down
140 changes: 138 additions & 2 deletions jxl/src/render/stages/upsample.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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 {
Expand All @@ -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(())
}
}

0 comments on commit 722295d

Please sign in to comment.