Implement a first version of modular channel decoding.

Does not support the Weighted predictor or prev-channel properties yet,
and is not optimized.

Author: veluca93

jxl/src/frame/modular.rs

@@ -7,7 +7,7 @@ use std::fmt::Debug;
 
 use crate::{
     bit_reader::BitReader,
-    error::{Error, Result},
+    error::Result,
     headers::{
         extra_channels::ExtraChannelInfo, frame_header::FrameHeader, modular::GroupHeader,
         JxlHeader,
@@ -16,10 +16,12 @@ use crate::{
     util::{tracing_wrappers::*, CeilLog2},
 };
 
+mod decode;
 mod predict;
 mod transforms;
 mod tree;
 
+use decode::{decode_modular_section, ModularStreamId};
 pub use predict::Predictor;
 use transforms::{make_grids, TransformStepChunk};
 pub use tree::Tree;
@@ -63,9 +65,9 @@ impl ChannelInfo {
 enum ModularGridKind {
     // Single big channel.
     None,
-    // 2048x2048 image-pixels.
+    // 2048x2048 image-pixels (if modular_group_shift == 1).
     Lf,
-    // 256x256 image-pixels.
+    // 256x256 image-pixels (if modular_group_shift == 1).
     Hf,
 }
 
@@ -78,6 +80,7 @@ struct ModularBuffer {
     auxiliary_data: Option<Image<i32>>,
     remaining_uses: usize,
     used_by_transforms: Vec<usize>,
+    size: (usize, usize),
 }
 
 #[allow(dead_code)]
@@ -144,50 +147,78 @@ impl FullModularImage {
             });
         }
 
+        if channels.is_empty() {
+            return Ok(Self {
+                buffer_info: vec![],
+                transform_steps: vec![],
+                section_buffer_indices: vec![vec![]; 2 + frame_header.passes.num_passes as usize],
+            });
+        }
+
         trace!("reading modular header");
         let header = GroupHeader::read(br)?;
 
-        if header.use_global_tree && global_tree.is_none() {
-            return Err(Error::NoGlobalTree);
-        }
-
         let (mut buffer_info, transform_steps) =
             transforms::meta_apply_transforms(&channels, &header.transforms)?;
 
         // Assign each (channel, group) pair present in the bitstream to the section in which it will be decoded.
         let mut section_buffer_indices: Vec<Vec<usize>> = vec![];
 
+        let mut sorted_buffers: Vec<_> = buffer_info
+            .iter()
+            .enumerate()
+            .filter_map(|(i, b)| {
+                if b.is_coded {
+                    Some((b.channel_id, i))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        sorted_buffers.sort_by_key(|x| x.0);
+
         section_buffer_indices.push(
-            buffer_info
+            sorted_buffers
                 .iter()
-                .enumerate()
-                .filter(|x| x.1.is_coded)
-                .take_while(|x| x.1.info.is_meta_or_small(frame_header.group_dim()))
-                .map(|x| x.0)
+                .take_while(|x| {
+                    buffer_info[x.1]
+                        .info
+                        .is_meta_or_small(frame_header.group_dim())
+                })
+                .map(|x| x.1)
                 .collect(),
         );
 
         section_buffer_indices.push(
-            buffer_info
+            sorted_buffers
                 .iter()
-                .enumerate()
-                .filter(|x| x.1.is_coded)
-                .skip_while(|x| x.1.info.is_meta_or_small(frame_header.group_dim()))
-                .filter(|x| x.1.info.is_shift_in_range(3, usize::MAX))
-                .map(|x| x.0)
+                .skip_while(|x| {
+                    buffer_info[x.1]
+                        .info
+                        .is_meta_or_small(frame_header.group_dim())
+                })
+                .filter(|x| buffer_info[x.1].info.is_shift_in_range(3, usize::MAX))
+                .map(|x| x.1)
                 .collect(),
         );
 
         for pass in 0..frame_header.passes.num_passes as usize {
             let (min_shift, max_shift) = frame_header.passes.downsampling_bracket(pass);
             section_buffer_indices.push(
-                buffer_info
+                sorted_buffers
                     .iter()
-                    .enumerate()
-                    .filter(|x| x.1.is_coded)
-                    .filter(|x| !x.1.info.is_meta_or_small(frame_header.group_dim()))
-                    .filter(|x| x.1.info.is_shift_in_range(min_shift, max_shift))
-                    .map(|x| x.0)
+                    .filter(|x| {
+                        !buffer_info[x.1]
+                            .info
+                            .is_meta_or_small(frame_header.group_dim())
+                    })
+                    .filter(|x| {
+                        buffer_info[x.1]
+                            .info
+                            .is_shift_in_range(min_shift, max_shift)
+                    })
+                    .map(|x| x.1)
                     .collect(),
             );
         }
@@ -206,7 +237,15 @@ impl FullModularImage {
             &mut buffer_info,
         );
 
-        // TODO(veluca93): read global channels
+        decode_modular_section(
+            &mut buffer_info,
+            &section_buffer_indices[0],
+            0,
+            ModularStreamId::GlobalData.get_id(frame_header),
+            &header,
+            global_tree,
+            br,
+        )?;
 
         Ok(FullModularImage {
             buffer_info,

jxl/src/frame/modular/decode.rs

@@ -0,0 +1,156 @@
+// 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,
+    entropy_coding::decode::Reader,
+    error::{Error, Result},
+    frame::quantizer::NUM_QUANT_TABLES,
+    headers::{frame_header::FrameHeader, modular::GroupHeader},
+    image::Image,
+    util::tracing_wrappers::*,
+};
+
+use super::{
+    predict::WeightedPredictorState, tree::NUM_NONREF_PROPERTIES, ModularBufferInfo, Tree,
+};
+
+#[allow(unused)]
+pub enum ModularStreamId {
+    GlobalData,
+    VarDCTLF(usize),
+    ModularLF(usize),
+    LFMeta(usize),
+    QuantTable(usize),
+    ModularHF { pass: usize, group: usize },
+}
+
+impl ModularStreamId {
+    pub fn get_id(&self, frame_header: &FrameHeader) -> usize {
+        match self {
+            Self::GlobalData => 0,
+            Self::VarDCTLF(g) => 1 + g,
+            Self::ModularLF(g) => 1 + frame_header.num_lf_groups() + g,
+            Self::LFMeta(g) => 1 + frame_header.num_lf_groups() * 2 + g,
+            Self::QuantTable(q) => 1 + frame_header.num_lf_groups() * 3 + q,
+            Self::ModularHF { pass, group } => {
+                1 + frame_header.num_lf_groups() * 3
+                    + NUM_QUANT_TABLES
+                    + frame_header.num_groups() * *pass
+                    + *group
+            }
+        }
+    }
+}
+
+#[allow(clippy::too_many_arguments)]
+#[instrument(level = "debug", skip(buffers, reader, br))]
+fn decode_modular_channel(
+    buffers: &mut [ModularBufferInfo],
+    buffer_indices: &[usize],
+    index: usize,
+    grid_index: usize,
+    stream_id: usize,
+    header: &GroupHeader,
+    tree: &Tree,
+    reader: &mut Reader,
+    br: &mut BitReader,
+) -> Result<()> {
+    debug!("reading channel");
+    let size = {
+        let b = &mut buffers[buffer_indices[index]].buffer_grid[grid_index];
+        if b.data.is_none() {
+            b.data = Some(Image::new(b.size)?)
+        }
+        b.size
+    };
+
+    let chan = buffers[buffer_indices[index]].channel_id;
+    let mut wp_state = WeightedPredictorState::new(header);
+    for y in 0..size.1 {
+        let mut property_buffer = [0; 256];
+        property_buffer[0] = chan as i32;
+        property_buffer[1] = stream_id as i32;
+        for x in 0..size.0 {
+            let prediction_result = tree.predict(
+                buffers,
+                buffer_indices,
+                index,
+                grid_index,
+                &mut wp_state,
+                x,
+                y,
+                &mut property_buffer,
+            );
+            let dec = reader.read_signed(br, prediction_result.context as usize)?;
+            let val =
+                prediction_result.guess + (prediction_result.multiplier as i64) * (dec as i64);
+            buffers[buffer_indices[index]].buffer_grid[grid_index]
+                .data
+                .as_mut()
+                .unwrap()
+                .as_rect_mut()
+                .row(y)[x] = val as i32;
+            trace!(y, x, val, dec, ?property_buffer, ?prediction_result);
+            // TODO(veluca): update WP errors.
+        }
+    }
+
+    Ok(())
+}
+
+pub fn decode_modular_section(
+    buffers: &mut [ModularBufferInfo],
+    buffer_indices: &[usize],
+    grid_index: usize,
+    stream_id: usize,
+    header: &GroupHeader,
+    global_tree: &Option<Tree>,
+    br: &mut BitReader,
+) -> Result<()> {
+    if buffer_indices.is_empty() {
+        return Ok(());
+    }
+    if header.use_global_tree && global_tree.is_none() {
+        return Err(Error::NoGlobalTree);
+    }
+    let local_tree = if !header.use_global_tree {
+        Some(Tree::read(br, 1024)?)
+    } else {
+        None
+    };
+    let tree = if header.use_global_tree {
+        global_tree.as_ref().unwrap()
+    } else {
+        local_tree.as_ref().unwrap()
+    };
+
+    if tree.max_property() >= NUM_NONREF_PROPERTIES - 2 {
+        todo!(
+            "WP and reference properties are not implemented yet, max property: {}",
+            tree.max_property()
+        );
+    }
+
+    let mut reader = tree.histograms.make_reader(br)?;
+
+    for i in 0..buffer_indices.len() {
+        decode_modular_channel(
+            buffers,
+            buffer_indices,
+            i,
+            grid_index,
+            stream_id,
+            header,
+            tree,
+            &mut reader,
+            br,
+        )?;
+    }
+
+    reader.check_final_state()?;
+
+    Ok(())
+}

jxl/src/frame/modular/predict.rs

@@ -3,7 +3,10 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-use crate::error::{Error, Result};
+use crate::{
+    error::{Error, Result},
+    headers::modular::GroupHeader,
+};
 use num_derive::FromPrimitive;
 use num_traits::FromPrimitive;
 
@@ -37,3 +40,18 @@ impl TryFrom<u32> for Predictor {
 impl Predictor {
     pub const NUM_PREDICTORS: u32 = Predictor::AverageAll as u32 + 1;
 }
+
+#[derive(Debug)]
+pub struct WeightedPredictorState;
+
+impl WeightedPredictorState {
+    pub fn new(_header: &GroupHeader) -> Self {
+        // TODO(veluca): implement the weighted predictor.
+        Self
+    }
+
+    pub fn predict_and_property(&self) -> (i64, i32) {
+        // TODO(veluca): implement the weighted predictor.
+        (0, 0)
+    }
+}

jxl/src/frame/modular/transforms.rs

@@ -463,11 +463,32 @@ pub fn make_grids(
     // Create grids.
     for g in buffer_info.iter_mut() {
         g.buffer_grid = get_grid_indices(g.grid_kind)
-            .map(|_| ModularBuffer {
-                data: None,
-                auxiliary_data: None,
-                remaining_uses: 0,
-                used_by_transforms: vec![],
+            .map(|(x, y)| {
+                let chan_size = g.info.size;
+                let size = match g.grid_kind {
+                    ModularGridKind::None => chan_size,
+                    ModularGridKind::Lf => {
+                        let dx = frame_header.lf_group_dim() >> g.info.shift.unwrap().0;
+                        let bx = x as usize * dx;
+                        let dy = frame_header.lf_group_dim() >> g.info.shift.unwrap().1;
+                        let by = y as usize * dy;
+                        ((chan_size.0 - bx).min(dx), (chan_size.1 - by).min(dy))
+                    }
+                    ModularGridKind::Hf => {
+                        let dx = frame_header.group_dim() >> g.info.shift.unwrap().0;
+                        let bx = x as usize * dx;
+                        let dy = frame_header.group_dim() >> g.info.shift.unwrap().1;
+                        let by = y as usize * dy;
+                        ((chan_size.0 - bx).min(dx), (chan_size.1 - by).min(dy))
+                    }
+                };
+                ModularBuffer {
+                    data: None,
+                    auxiliary_data: None,
+                    remaining_uses: 0,
+                    used_by_transforms: vec![],
+                    size,
+                }
             })
             .collect();
     }