Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
mo271 committed Dec 4, 2024
1 parent 3b90808 commit a762336
Show file tree
Hide file tree
Showing 5 changed files with 175 additions and 18 deletions.
8 changes: 8 additions & 0 deletions jxl/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,14 @@ pub enum Error {
InvalidPredictor(u32),
#[error("Invalid modular mode property: {0}")]
InvalidProperty(u32),
#[error("Invalid alpha channel for blending: {0}, limit is {1}")]
PatchesInvalidAlphaChannel(u32, u32),
#[error("Invalid patch blend mode: {0}, limit is {1}")]
PatchesInvalidBlendMode(u8, u8),
#[error("Invalid Patch: negative {0}-coordinate: {1} base {0}, {2} delta {0}")]
PatchesInvalidDelta(String, usize, i32),
#[error("Invalid Patch {0}: at {1} + {2} > {3}")]
PatchesOutOfBounds(String, usize, usize, u32),
#[error("Too many patches: {0}, limit is {1}")]
PatchesTooMany(u32, u32),
#[error("Reference too large: {0}, limit is {1}")]
Expand Down
2 changes: 1 addition & 1 deletion jxl/src/features/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,6 @@
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

pub mod patches;
pub mod noise;
pub mod patches;
pub mod spline;
163 changes: 150 additions & 13 deletions jxl/src/features/patches.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@ use crate::{
bit_reader::BitReader,
entropy_coding::decode::Histograms,
error::{Error, Result},
image::Image,
util::tracing_wrappers::*,
};

// TODO(firsching): move to some common place?
const MAX_NUM_REFERENCE_FRAMES : u32 = 4;
const MAX_NUM_REFERENCE_FRAMES: u32 = 4;

// Context numbers as specified in Section C.4.5, Listing C.2:
const NUM_REF_PATCH_CONTEXT: usize = 0;
Expand Down Expand Up @@ -90,17 +91,23 @@ pub struct PatchPosition {

#[derive(Debug, Default)]
pub struct PatchesDictionary {
pub reference_frames: [Option<Box<Image<u8>>>; 4],
pub positions: Vec<PatchPosition>,
pub ref_positions: Vec<PatchReferencePosition>,
ref_positions: Vec<PatchReferencePosition>,
blendings: Vec<PatchBlending>,
num_patches: Vec<usize>,
sorted_patches_y0: Vec<(usize, usize)>,
sorted_patches_y1: Vec<(usize, usize)>,
blendings_stride: u32,
}

impl PatchesDictionary {
#[instrument(level = "debug", skip(br), ret, err)]
pub fn read(br: &mut BitReader, xsize: u32, ysize: u32) -> Result<PatchesDictionary> {
pub fn read(
br: &mut BitReader,
xsize: u32,
ysize: u32,
num_extra_channels: u32,
reference_frames: [Option<Box<Image<u8>>>; 4],
) -> Result<PatchesDictionary> {
let blendings_stride = num_extra_channels + 1;
trace!(pos = br.total_bits_read());
let patches_histograms = Histograms::decode(NUM_PATCH_DICTIONARY_CONTEXTS, br, true)?;
let mut patches_reader = patches_histograms.make_reader(br)?;
Expand All @@ -112,19 +119,149 @@ impl PatchesDictionary {
if num_ref_patch > max_ref_patches {
return Err(Error::PatchesTooMany(num_ref_patch, max_ref_patches));
}

let mut ref_positions: Vec<PatchReferencePosition> = Vec::with_capacity(num_ref_patch as usize);
let mut total_patches = 0;
let mut next_size = 1;
let mut positions: Vec<PatchPosition> = Vec::new();
let mut blendings = Vec::new();
let mut ref_positions: Vec<PatchReferencePosition> =
Vec::with_capacity(num_ref_patch as usize);
for _ in 0..num_ref_patch {
let reference = patches_reader.read(br, REFERENCE_FRAME_CONTEXT)?;
if reference >= MAX_NUM_REFERENCE_FRAMES {
return Err(Error::PatchesRefTooLarge(reference, MAX_NUM_REFERENCE_FRAMES));
return Err(Error::PatchesRefTooLarge(
reference,
MAX_NUM_REFERENCE_FRAMES,
));
}

// TODO: add check if we are after xyb color transform?
// let image = reference_frames[ref_positions[0].reference as usize];
// add checks that use that image here?
// TODO: fill in correct numbers here
ref_positions.push(PatchReferencePosition {reference:0, x0:0, y0:0, xsize:0, ysize: 0})
let x0 = patches_reader.read(br, PATCH_REFERENCE_POSITION_CONTEXT)? as usize;
let y0 = patches_reader.read(br, PATCH_REFERENCE_POSITION_CONTEXT)? as usize;
let ref_pos_xsize = patches_reader.read(br, PATCH_SIZE_CONTEXT)? as usize + 1;
let ref_pos_ysize = patches_reader.read(br, PATCH_SIZE_CONTEXT)? as usize + 1;

}
// TODO: add check : ref_pos.x0 + ref_pos.xsize > ib.xsize()
// TODO: add check : ref_pos.y0 + ref_pos.ysize > ib.ysize()

let id_count = patches_reader.read(br, PATCH_COUNT_CONTEXT)? + 1;
if id_count > max_patches + 1 {
return Err(Error::PatchesTooMany(id_count, max_patches));
}
total_patches += id_count;

if total_patches > max_patches {
return Err(Error::PatchesTooMany(total_patches, max_patches));
}

if next_size < total_patches {
next_size *= 2;
next_size = std::cmp::min(next_size, max_patches);
}
positions.reserve(next_size as usize);
blendings.reserve(next_size as usize * PATCH_BLEND_MODE_NUM_BLEND_MODES as usize);

todo!("implement patch decoding")
for _ in 0..id_count {
let mut pos = PatchPosition {
x: 0,
y: 0,
ref_pos_idx: ref_positions.len(),
};
if positions.is_empty() {
// Read initial position
pos.x = patches_reader.read(br, PATCH_POSITION_CONTEXT)? as usize;
pos.y = patches_reader.read(br, PATCH_POSITION_CONTEXT)? as usize;
} else {
// Read offsets and calculate new position
let delta_x = patches_reader.read_signed(br, PATCH_OFFSET_CONTEXT)?;
if delta_x < 0 && (-delta_x as usize) > positions.last().unwrap().x {
return Err(Error::PatchesInvalidDelta(
"x".to_string(),
positions.last().unwrap().x,
delta_x,
));
}
pos.x = (positions.last().unwrap().x as i32 + delta_x) as usize;

let delta_y = patches_reader.read_signed(br, PATCH_OFFSET_CONTEXT)?;
if delta_y < 0 && (-delta_y as usize) > positions.last().unwrap().y {
return Err(Error::PatchesInvalidDelta(
"y".to_string(),
positions.last().unwrap().y,
delta_y,
));
}
pos.y = (positions.last().unwrap().y as i32 + delta_y) as usize;
}

if pos.x + ref_pos_xsize > xsize as usize {
return Err(Error::PatchesOutOfBounds(
"x".to_string(),
pos.x,
ref_pos_xsize,
xsize,
));
}
if pos.y + ref_pos_ysize > ysize as usize {
return Err(Error::PatchesOutOfBounds(
"y".to_string(),
pos.y,
ref_pos_ysize,
ysize,
));
}

let mut alpha_channel = 0;
let mut clamp = false;
for _ in 0..blendings_stride {
let blend_mode = patches_reader.read(br, PATCH_BLEND_MODE_CONTEXT)? as u8;
if blend_mode >= PATCH_BLEND_MODE_NUM_BLEND_MODES {
return Err(Error::PatchesInvalidBlendMode(
blend_mode,
PATCH_BLEND_MODE_NUM_BLEND_MODES,
));
}

if blend_mode == PATCH_BLEND_MODE_BLEND_ABOVE
|| blend_mode == PATCH_BLEND_MODE_BLEND_BELOW
{
alpha_channel =
patches_reader.read(br, PATCH_ALPHA_CHANNEL_CONTEXT)? as u32;
if alpha_channel >= PATCH_BLEND_MODE_NUM_BLEND_MODES as u32 {
return Err(Error::PatchesInvalidAlphaChannel(
alpha_channel,
num_extra_channels,
));
}
}

if blend_mode == PATCH_BLEND_MODE_ADD || blend_mode == PATCH_BLEND_MODE_MUL {
clamp = patches_reader.read(br, PATCH_CLAMP_CONTEXT)? != 0;
}
blendings.push(PatchBlending {
mode: blend_mode,
alpha_channel,
clamp,
});
}
positions.push(pos);
}

ref_positions.push(PatchReferencePosition {
reference: 0,
x0,
y0,
xsize: ref_pos_xsize,
ysize: ref_pos_ysize,
})
}
Ok(PatchesDictionary {
reference_frames,
positions,
blendings,
ref_positions,
blendings_stride,
})
}
}
15 changes: 11 additions & 4 deletions jxl/src/frame.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,15 @@
use crate::{
bit_reader::BitReader,
error::Result,
features::{noise::Noise, spline::Splines},
features::patches::PatchesDictionary,
features::{noise::Noise, patches::PatchesDictionary, spline::Splines},
headers::{
color_encoding::ColorSpace,
encodings::UnconditionalCoder,
extra_channels::ExtraChannelInfo,
frame_header::{Encoding, FrameHeader, Toc, TocNonserialized},
FileHeader,
},
image::Image,
util::tracing_wrappers::*,
};
use modular::{FullModularImage, Tree};
Expand Down Expand Up @@ -138,10 +138,16 @@ impl Frame {

let patches = if self.header.has_patches() {
info!("decoding patches");
// TODO
let reference_positions: [Option<Box<Image<u8>>>; 4] = Default::default();
// TODO
let num_extra_channels = 0;
Some(PatchesDictionary::read(
br,
self.header.width,
self.header.height,
num_extra_channels,
reference_positions,
)?)
} else {
None
Expand Down Expand Up @@ -224,10 +230,11 @@ mod test_frame {
Ok(result)
}


#[test]
fn patches() -> Result<(), Error> {
let frame = read_frame(include_bytes!("../resources/test/GrayscalePatchesVarDCT.jxl"))?;
let frame = read_frame(include_bytes!(
"../resources/test/GrayscalePatchesVarDCT.jxl"
))?;
let lf_global = frame.lf_global.unwrap();
let patches_dict = lf_global.patches.unwrap();
// TODO continue here
Expand Down
5 changes: 5 additions & 0 deletions jxl/src/image.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,11 @@ pub struct Image<T: ImageDataType> {
data: Vec<T>,
}

impl<T: ImageDataType> Debug for Image<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{:?} {}x{}", T::DATA_TYPE_ID, self.size.0, self.size.1,)
}
}
#[derive(Clone, Copy)]
pub struct ImageRect<'a, T: ImageDataType> {
origin: (usize, usize),
Expand Down

0 comments on commit a762336

Please sign in to comment.