LfGlobal decoding: LF quantization factors, scaffolding

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<T, E = Error> = std::result::Result<T, E>;

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<LfGlobalState>,
 }
 
 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(())
+    }
 }

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<LfQuantFactors> {
+        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,
+        })
+    }
+}

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]

jxl/src/headers/encodings.rs

@@ -386,7 +386,7 @@ impl<Config, T: UnconditionalCoder<Config>> DefaultedCoder<Config> 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 {

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

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,

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,

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<u32>,
 }
 
-#[derive(UnconditionalCoder, Debug)]
+#[derive(UnconditionalCoder, Debug, Clone)]
 pub struct Preview {
     div8: bool,
     #[condition(div8)]

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]