diff --git a/jxl/src/error.rs b/jxl/src/error.rs index 8fc6959..b74d2c3 100644 --- a/jxl/src/error.rs +++ b/jxl/src/error.rs @@ -103,6 +103,8 @@ pub enum Error { PipelineChannelUnused(usize), #[error("Trying to copy rects of different size, src: {0}x{1} dst {2}x{3}")] CopyOfDifferentSize(usize, usize, usize, usize), + #[error("LF quantization factor is too small: {0}")] + LfQuantFactorTooSmall(f32), } pub type Result = std::result::Result; diff --git a/jxl/src/frame.rs b/jxl/src/frame.rs index 14e365f..c0166f6 100644 --- a/jxl/src/frame.rs +++ b/jxl/src/frame.rs @@ -8,14 +8,41 @@ use crate::{ error::Result, headers::{ encodings::UnconditionalCoder, - frame_header::{FrameHeader, Toc, TocNonserialized}, + frame_header::{Encoding, FrameHeader, Toc, TocNonserialized}, FileHeader, }, + util::tracing_wrappers::*, }; +use quantizer::LfQuantFactors; + +mod quantizer; + +#[derive(Debug, PartialEq, Eq)] +pub enum Section { + LfGlobal, + Lf(usize), + HfGlobal, + Hf(usize, usize), // group, pass +} + +pub struct LfGlobalState { + // TODO(veluca93): patches + // TODO(veluca93): splines + // TODO(veluca93): noise + #[allow(dead_code)] + lf_quant: LfQuantFactors, + // TODO(veluca93), VarDCT: HF quant matrices + // TODO(veluca93), VarDCT: block context map + // TODO(veluca93), VarDCT: LF color correlation + // TODO(veluca93): Modular data +} pub struct Frame { header: FrameHeader, toc: Toc, + #[allow(dead_code)] + file_header: FileHeader, + lf_global: Option, } impl Frame { @@ -28,14 +55,16 @@ impl Frame { &(), br, &TocNonserialized { - num_entries: num_toc_entries, + num_entries: num_toc_entries as u32, }, ) .unwrap(); br.jump_to_byte_boundary()?; Ok(Self { header: frame_header, + file_header: file_header.clone(), toc, + lf_global: None, }) } @@ -69,4 +98,52 @@ impl Frame { .collect() } } + + #[instrument(level = "debug", skip(self), ret)] + pub fn get_section_idx(&self, section: Section) -> usize { + if self.header.num_toc_entries() == 1 { + 0 + } else { + match section { + Section::LfGlobal => 0, + Section::Lf(a) => 1 + a, + Section::HfGlobal => self.header.num_dc_groups() + 1, + Section::Hf(group, pass) => { + 2 + self.header.num_dc_groups() + self.header.num_groups() * pass + group + } + } + } + } + + #[instrument(skip_all)] + pub fn decode_lf_global(&mut self, br: &mut BitReader) -> Result<()> { + assert!(self.lf_global.is_none()); + + if self.header.has_patches() { + info!("decoding patches"); + todo!("patches not implemented"); + } + + if self.header.has_splines() { + info!("decoding splines"); + todo!("splines not implemented"); + } + + if self.header.has_noise() { + info!("decoding noise"); + todo!("noise not implemented"); + } + + let lf_quant = LfQuantFactors::new(br)?; + debug!(?lf_quant); + + if self.header.encoding == Encoding::VarDCT { + info!("decoding VarDCT info"); + todo!("VarDCT not implemented"); + } + + self.lf_global = Some(LfGlobalState { lf_quant }); + + Ok(()) + } } diff --git a/jxl/src/frame/quantizer.rs b/jxl/src/frame/quantizer.rs new file mode 100644 index 0000000..3c355a8 --- /dev/null +++ b/jxl/src/frame/quantizer.rs @@ -0,0 +1,42 @@ +// 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::{ + bit_reader::BitReader, + error::{Error, Result}, + headers::encodings::{Empty, UnconditionalCoder}, +}; + +#[derive(Debug)] +#[allow(dead_code)] +pub struct LfQuantFactors { + pub quant_factors: [f32; 3], + pub inv_quant_factors: [f32; 3], +} + +impl LfQuantFactors { + pub fn new(br: &mut BitReader) -> Result { + let mut quant_factors = [0.0f32; 3]; + if br.read(1)? == 1 { + quant_factors[0] = 1.0 / 4096.0; + quant_factors[1] = 1.0 / 512.0; + quant_factors[2] = 1.0 / 256.0; + } else { + for i in 0..3 { + quant_factors[i] = f32::read_unconditional(&(), br, &Empty {})? / 128.0; + if quant_factors[i] < 1e-8 { + return Err(Error::LfQuantFactorTooSmall(quant_factors[i])); + } + } + } + + let inv_quant_factors = quant_factors.map(f32::recip); + + Ok(LfQuantFactors { + quant_factors, + inv_quant_factors, + }) + } +} diff --git a/jxl/src/headers/color_encoding.rs b/jxl/src/headers/color_encoding.rs index 8c25802..2284a21 100644 --- a/jxl/src/headers/color_encoding.rs +++ b/jxl/src/headers/color_encoding.rs @@ -54,7 +54,7 @@ pub enum RenderingIntent { Absolute, } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct CustomXY { #[default(0)] #[coder(u2S(Bits(19), Bits(19) + 524288, Bits(20) + 1048576, Bits(21) + 2097152))] @@ -68,7 +68,7 @@ pub struct CustomTransferFunctionNonserialized { color_space: ColorSpace, } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] #[nonserialized(CustomTransferFunctionNonserialized)] #[validate] pub struct CustomTransferFunction { @@ -104,7 +104,7 @@ impl CustomTransferFunction { } } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] #[validate] pub struct ColorEncoding { #[all_default] diff --git a/jxl/src/headers/encodings.rs b/jxl/src/headers/encodings.rs index aae300c..ee54a2c 100644 --- a/jxl/src/headers/encodings.rs +++ b/jxl/src/headers/encodings.rs @@ -386,7 +386,7 @@ impl> DefaultedCoder for T { // TODO(veluca93): this will likely need to be implemented differently if // there are extensions. -#[derive(Debug, PartialEq, Default)] +#[derive(Debug, PartialEq, Default, Clone)] pub struct Extensions {} impl UnconditionalCoder<()> for Extensions { diff --git a/jxl/src/headers/frame_header.rs b/jxl/src/headers/frame_header.rs index 81a6072..8ab3781 100644 --- a/jxl/src/headers/frame_header.rs +++ b/jxl/src/headers/frame_header.rs @@ -25,7 +25,7 @@ enum FrameType { } #[derive(UnconditionalCoder, Copy, Clone, PartialEq, Debug, FromPrimitive)] -enum Encoding { +pub enum Encoding { VarDCT = 0, Modular = 1, } @@ -270,7 +270,7 @@ pub struct FrameHeader { #[coder(Bits(1))] #[default(Encoding::VarDCT)] - encoding: Encoding, + pub encoding: Encoding, #[default(0)] flags: u64, @@ -431,23 +431,23 @@ impl FrameHeader { self.group_dim() * BLOCK_DIM } - pub fn num_toc_entries(&self) -> u32 { - const GROUP_DIM: u32 = 256; - const BLOCK_DIM: u32 = 8; - const H_SHIFT: [u32; 4] = [0, 1, 1, 0]; - const V_SHIFT: [u32; 4] = [0, 1, 0, 1]; + fn group_counts(&self) -> (usize, usize) { + const GROUP_DIM: usize = 256; + const BLOCK_DIM: usize = 8; + const H_SHIFT: [usize; 4] = [0, 1, 1, 0]; + const V_SHIFT: [usize; 4] = [0, 1, 0, 1]; let mut maxhs = 0; let mut maxvs = 0; for ch in self.jpeg_upsampling { maxhs = maxhs.max(H_SHIFT[ch as usize]); maxvs = maxvs.max(V_SHIFT[ch as usize]); } - let xsize = self.width; - let ysize = self.height; + let xsize = self.width as usize; + let ysize = self.height as usize; let xsize_blocks = xsize.div_ceil(BLOCK_DIM << maxhs) << maxhs; let ysize_blocks = ysize.div_ceil(BLOCK_DIM << maxvs) << maxvs; - let group_dim: u32 = self.group_dim(); + let group_dim = self.group_dim() as usize; let xsize_groups = xsize.div_ceil(group_dim); let ysize_groups = ysize.div_ceil(group_dim); @@ -457,10 +457,24 @@ impl FrameHeader { let num_groups = xsize_groups * ysize_groups; let num_dc_groups = xsize_dc_groups * ysize_dc_groups; + (num_groups, num_dc_groups) + } + + pub fn num_groups(&self) -> usize { + self.group_counts().0 + } + + pub fn num_dc_groups(&self) -> usize { + self.group_counts().1 + } + + pub fn num_toc_entries(&self) -> usize { + let (num_groups, num_dc_groups) = self.group_counts(); + if num_groups == 1 && self.passes.num_passes == 1 { 1 } else { - 2 + num_dc_groups + num_groups + 2 + num_dc_groups + num_groups * self.passes.num_passes as usize } } @@ -469,6 +483,18 @@ impl FrameHeader { / (animation.tps_numerator as f64) } + pub fn has_patches(&self) -> bool { + self.flags & Flags::ENABLE_PATCHES != 0 + } + + pub fn has_noise(&self) -> bool { + self.flags & Flags::ENABLE_NOISE != 0 + } + + pub fn has_splines(&self) -> bool { + self.flags & Flags::ENABLE_SPLINES != 0 + } + fn check(&self, nonserialized: &FrameHeaderNonserialized) -> Result<(), Error> { if self.upsampling > 1 { if let Some((info, upsampling)) = nonserialized diff --git a/jxl/src/headers/image_metadata.rs b/jxl/src/headers/image_metadata.rs index 9790790..6e347c5 100644 --- a/jxl/src/headers/image_metadata.rs +++ b/jxl/src/headers/image_metadata.rs @@ -11,7 +11,7 @@ use crate::{ use jxl_macros::UnconditionalCoder; use num_derive::FromPrimitive; -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] pub struct Signature; impl Signature { @@ -45,7 +45,7 @@ pub enum Orientation { Rotate270 = 8, } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct Animation { #[coder(u2S(100, 1000, Bits(10) + 1, Bits(30) + 1))] pub tps_numerator: u32, @@ -56,7 +56,7 @@ pub struct Animation { pub have_timecodes: bool, } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] #[validate] pub struct ToneMapping { #[all_default] @@ -93,7 +93,7 @@ impl ToneMapping { // TODO(firsching): remove once we use this! #[allow(dead_code)] -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct ImageMetadata { #[all_default] all_default: bool, diff --git a/jxl/src/headers/mod.rs b/jxl/src/headers/mod.rs index 25819dc..1565c40 100644 --- a/jxl/src/headers/mod.rs +++ b/jxl/src/headers/mod.rs @@ -21,7 +21,7 @@ pub use image_metadata::*; pub use size::Size; pub use transform_data::*; -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct FileHeader { #[allow(dead_code)] signature: Signature, diff --git a/jxl/src/headers/size.rs b/jxl/src/headers/size.rs index f8c3917..92f4820 100644 --- a/jxl/src/headers/size.rs +++ b/jxl/src/headers/size.rs @@ -19,7 +19,7 @@ enum AspectRatio { Ratio2Over1 = 7, } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct Size { small: bool, #[condition(small)] @@ -38,7 +38,7 @@ pub struct Size { xsize: Option, } -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct Preview { div8: bool, #[condition(div8)] diff --git a/jxl/src/headers/transform_data.rs b/jxl/src/headers/transform_data.rs index 8a8a72f..9149e5e 100644 --- a/jxl/src/headers/transform_data.rs +++ b/jxl/src/headers/transform_data.rs @@ -14,7 +14,7 @@ pub struct CustomTransformDataNonserialized { // TODO(firsching): remove once we use this! #[allow(dead_code)] -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] pub struct OpsinInverseMatrix { #[all_default] #[default(true)] @@ -320,7 +320,7 @@ const DEFAULT_KERN_8: [f32; 210] = [ // TODO(firsching): remove once we use this! #[allow(dead_code)] -#[derive(UnconditionalCoder, Debug)] +#[derive(UnconditionalCoder, Debug, Clone)] #[nonserialized(CustomTransformDataNonserialized)] pub struct CustomTransformData { #[all_default] diff --git a/jxl_cli/src/main.rs b/jxl_cli/src/main.rs index 24e6ef3..1d524b1 100644 --- a/jxl_cli/src/main.rs +++ b/jxl_cli/src/main.rs @@ -5,7 +5,7 @@ use jxl::bit_reader::BitReader; use jxl::container::{ContainerParser, ParseEvent}; -use jxl::frame::Frame; +use jxl::frame::{Frame, Section}; use jxl::headers::FileHeader; use jxl::icc::read_icc; use std::env; @@ -28,13 +28,15 @@ fn parse_jxl_codestream(data: &[u8]) -> Result<(), jxl::error::Error> { }; loop { - let frame = Frame::new(&mut br, &file_header)?; + let mut frame = Frame::new(&mut br, &file_header)?; br.jump_to_byte_boundary()?; - let section_readers = frame.sections(&mut br)?; + let mut section_readers = frame.sections(&mut br)?; println!("read frame with {} sections", section_readers.len()); + frame.decode_lf_global(&mut section_readers[frame.get_section_idx(Section::LfGlobal)])?; + if frame.header().is_last { break; }