diff --git a/wai/wai-rs/src/jpegai_bitstream.rs b/wai/wai-rs/src/jpegai_bitstream.rs index 846ebd4e2..9818558f3 100644 --- a/wai/wai-rs/src/jpegai_bitstream.rs +++ b/wai/wai-rs/src/jpegai_bitstream.rs @@ -11,6 +11,12 @@ //! bit-granular cursor. Reads past the end return the bits available (zero-extended), //! matching the reference's bounds check. +/// Bits needed to carry `max_value`, per the reference `HeaderCoder::calc_bits_count`: +/// `ceil(log2(max_value + 1))`. `max_value = 0` needs no bits. +pub fn bits_for_max(max_value: u64) -> u32 { + if max_value == 0 { 0 } else { u64::BITS - max_value.leading_zeros() } +} + /// Marker identifiers (big-endian 2-byte prefixes). #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub enum Marker { @@ -90,6 +96,16 @@ impl<'a> BitReader<'a> { v } + /// Read one symbol whose width is **derived from its maximum value**, per the + /// reference `HeaderCoder::calc_bits_count`: `bits = ceil(log2(max_value + 1))`. + /// + /// This is the header's universal width rule — no field has a fixed width. Some are + /// runtime-dependent: a `max_symbol` decoded earlier in a segment sets the width of + /// the weights that follow it (e.g. 65535 ⇒ 16 bits, then 99 ⇒ 7 bits). + pub fn read_max(&mut self, max_value: u64) -> u64 { + self.read(bits_for_max(max_value)) + } + /// Unsigned exp-Golomb, k=0: count leading zeros `n`, then read `n` bits; /// value = `2^n − 1 + suffix`. Used for marker-segment payload lengths. pub fn read_expgolomb_k0(&mut self) -> u64 { @@ -196,6 +212,205 @@ pub struct PictureHeader { /// Branches not exercised by the 4:4:4 single-tile reference stream (region /// sub-splits, `independent_beta_uv`'s second β, the signalled colour matrix, /// multi-threading counts) follow the reference syntax and are decoded but untested. +// ---- tools header (TON) ---- + +/// One `decode_filters` group of the EFE **linear** filter header. +#[derive(Clone, Debug, Default)] +pub struct EflFilters { + /// `best_cand_idx − 1`; `-1` means the plane's filter is absent. + pub cand_u: i32, + pub cand_v: i32, + pub fl_u: u32, + pub fl_v: u32, + pub min_symbol: u32, + pub max_symbol: u32, + pub weights_u: Vec, + pub weights_u2: Vec, + pub weights_v: Vec, + pub weights_v2: Vec, +} + +/// EFE linear filter header: the two plane means plus two filter groups. +#[derive(Clone, Debug, Default)] +pub struct EflHeader { + /// `B1[i] / 100` = the per-plane mean consumed by `jpegai_pipeline::efl_chroma`. + pub b1: [u32; 2], + pub filters_best2: EflFilters, + pub filters_best: EflFilters, +} + +/// eICCI per-tile model selection. +#[derive(Clone, Debug, Default)] +pub struct IcciTile { + /// Per-plane use flags (Y, U, V). + pub use_yuv: [bool; 3], + pub use_short_list: bool, + /// Signalled model indices (into the active list) for the Y and the joint-UV paths. + pub model_idx_y: u32, + pub model_idx_uv: u32, +} + +/// EFE nonlinear filter header (the luma-aided chroma filter's parameters). +#[derive(Clone, Debug, Default)] +pub struct EfnHeader { + pub min_symbol: u32, + pub max_symbol: u32, + pub mask1_enabled: bool, + pub mask2_enabled: bool, + pub nonlinear_u_enabled: bool, + pub nonlinear_v_enabled: bool, + pub tile_width: u32, + pub tile_height: u32, + /// Per-tile luma range — the `luma_min_q` / `luma_max_q` of `efn_filter`. + pub luma_min: Vec, + pub luma_max: Vec, + /// Per-plane filter weights (`A1`), already offset by `min_symbol` — the + /// `weights_q15` of `efn_filter`. + pub weights_u: Vec, + pub weights_v: Vec, +} + +/// The decoded tools header (TON substream). +#[derive(Clone, Debug, Default)] +pub struct ToolsHeader { + pub efl_enabled: bool, + pub efl: Option, + pub icci_enabled: bool, + pub icci_tile_manager_enabled: bool, + pub icci_tiles: Vec, + pub efn_enabled: bool, + pub efn: Option, + pub lef_enabled: bool, + /// LEF's reference channel index. + pub lef_ch_idx: u32, + /// Bit position just past the last field (before byte-alignment padding). + pub end_bit: usize, +} + +/// `cands[i][0]` — the spatial-tiling candidate filter counts of the EFE linear filter. +const EFL_CAND_FILTER_COUNT: [u32; 6] = [1, 2, 2, 3, 3, 4]; + +/// One `decode_filters` group. `scale_2x2` selects the 4:4:4 path, where a single +/// plane's coefficients are decoded and replicated rather than coded four times. +fn parse_efl_filters(br: &mut BitReader, scale_2x2: bool) -> EflFilters { + let mut f = EflFilters::default(); + f.cand_u = br.read_max(15) as i32 - 1; + f.cand_v = br.read_max(15) as i32 - 1; + if f.cand_u > -1 { f.fl_u = br.read_max(9) as u32; } + if f.cand_v > -1 { f.fl_v = br.read_max(9) as u32; } + f.min_symbol = br.read_max((1 << 16) - 1) as u32; + f.max_symbol = br.read_max((1 << 16) - 1) as u32; + let mut group = |cand: i32, fl: u32, br: &mut BitReader| -> (Vec, Vec) { + if cand <= -1 { return (Vec::new(), Vec::new()); } + let taps = (fl * fl) as usize; + // shape[0]==4 is coded as a single plane under 4:4:4 and replicated + let n_main = if scale_2x2 { taps } else { EFL_CAND_FILTER_COUNT[cand as usize] as usize * taps }; + let main = (0..n_main).map(|_| br.read_max(f.max_symbol as u64) as u32).collect(); + let second = (0..taps).map(|_| br.read_max(f.max_symbol as u64) as u32).collect(); + (main, second) + }; + let (u, u2) = group(f.cand_u, f.fl_u, br); + f.weights_u = u; + f.weights_u2 = u2; + let (v, v2) = group(f.cand_v, f.fl_v, br); + f.weights_v = v; + f.weights_v2 = v2; + f +} + +/// Parse the **post-filter chain** of the tools header (TON), in filters-factory order: +/// EFE-linear → ICCI → eICCI → EFE-nonlinear → LEF. Each tool is gated by its own enable +/// flag (the recursive header walk), and every field's width is derived from its +/// `max_symbol_value` — including the runtime-dependent ones, where a just-decoded +/// `max_symbol` sets the width of the weights that follow. +/// +/// `start_bit` is the bit position of the **first filter's enable flag**, not the TON +/// payload start: `_stream_header_part` defaults to `"tool_header"`, so tools other than +/// the post-filters also emit into TON ahead of the filter chain (in the reference +/// stream two such bits precede EFE-linear's flag). Locating that prefix generally means +/// walking the whole tool tree; this function parses the chain from a given offset. +/// +/// `models_count` is the size of the eICCI model list the stream selects from (it sets +/// the width of the signalled model indices); `num_tiles` is the eICCI tile count. +pub fn parse_tools_header(data: &[u8], start_bit: usize, scale_2x2: bool, models_count: u32, num_tiles: usize) -> ToolsHeader { + let mut br = BitReader::at(data, start_bit); + let mut t = ToolsHeader::default(); + + t.efl_enabled = br.read_max(1) == 1; + if t.efl_enabled { + let mut h = EflHeader::default(); + h.b1[0] = br.read_max((1 << 15) - 1) as u32; + h.b1[1] = br.read_max((1 << 15) - 1) as u32; + h.filters_best2 = parse_efl_filters(&mut br, scale_2x2); + h.filters_best = parse_efl_filters(&mut br, scale_2x2); + t.efl = Some(h); + } + + t.icci_enabled = br.read_max(1) == 1; + if t.icci_enabled { + t.icci_tile_manager_enabled = br.read_max(1) == 1; + for _ in 0..num_tiles.max(1) { + let mut tile = IcciTile::default(); + for u in tile.use_yuv.iter_mut() { *u = br.read_max(1) == 1; } + if tile.use_yuv.iter().any(|&b| b) { + tile.use_short_list = br.read_max(1) == 1; + if tile.use_yuv[0] { + tile.model_idx_y = br.read_max(models_count as u64) as u32; + } + if tile.use_yuv[1] || tile.use_yuv[2] { + tile.model_idx_uv = br.read_max(models_count as u64) as u32; + } + } + t.icci_tiles.push(tile); + } + } + + t.efn_enabled = br.read_max(1) == 1; + if t.efn_enabled { + let mut e = EfnHeader::default(); + e.min_symbol = br.read_max((1 << 16) - 1) as u32; + e.max_symbol = br.read_max((1 << 16) - 1) as u32; + e.mask1_enabled = br.read_max(1) == 1; + e.mask2_enabled = br.read_max(1) == 1; + if e.mask1_enabled || e.mask2_enabled { + let _b_size = br.read_max((1 << 10) - 1); + let h = br.read_max((1 << 10) - 1) as usize; + let w = br.read_max((1 << 10) - 1) as usize; + // W5 masks: only the enabled one(s) carry h×w symbols (max 2 ⇒ 2 bits) + for enabled in [e.mask1_enabled, e.mask2_enabled] { + let n = if enabled { h * w } else { 0 }; + for _ in 0..n { br.read_max(2); } + } + } + e.nonlinear_u_enabled = br.read_max(1) == 1; + e.nonlinear_v_enabled = br.read_max(1) == 1; + if e.nonlinear_u_enabled || e.nonlinear_v_enabled { + e.tile_width = br.read_max((1 << 16) - 1) as u32; + e.tile_height = br.read_max((1 << 16) - 1) as u32; + let n = br.read_max((1 << 16) - 1) as usize; // numTiles + e.luma_min = (0..n).map(|_| br.read_max((1 << 16) - 1) as u32).collect(); + e.luma_max = (0..n).map(|_| br.read_max((1 << 16) - 1) as u32).collect(); + // candNum + A1 are coded per ENABLED plane (U first, then V) + for (enabled, dst) in [(e.nonlinear_u_enabled, 0), (e.nonlinear_v_enabled, 1)] { + if !enabled { continue; } + let cand_num = br.read_max((1 << 16) - 1) as usize; + let w: Vec = (0..cand_num) + .map(|_| br.read_max(e.max_symbol as u64) as u32 + e.min_symbol) + .collect(); + if dst == 0 { e.weights_u = w; } else { e.weights_v = w; } + } + } + t.efn = Some(e); + } + + t.lef_enabled = br.read_max(1) == 1; + if t.lef_enabled { + t.lef_ch_idx = br.read_max(255) as u32; + } + t.end_bit = br.bit_pos(); + t +} + pub fn parse_picture_header(data: &[u8], start_bit: usize) -> PictureHeader { let mut br = BitReader::at(data, start_bit); let mut h = PictureHeader::default(); diff --git a/wai/wai-rs/tests/jpegai_bitstream_calibration.rs b/wai/wai-rs/tests/jpegai_bitstream_calibration.rs index 22b7bd9cf..9d71a1de4 100644 --- a/wai/wai-rs/tests/jpegai_bitstream_calibration.rs +++ b/wai/wai-rs/tests/jpegai_bitstream_calibration.rs @@ -9,7 +9,7 @@ //! env: WAI_JPEGAI_CODESTREAM = path to test256.bin #![cfg(feature = "neural_int")] -use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, BitReader, Marker}; +use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, parse_tools_header, BitReader, Marker}; fn codestream() -> Option> { match std::env::var("WAI_JPEGAI_CODESTREAM") { @@ -126,3 +126,76 @@ fn full_picture_header_matches_reference() { assert_eq!(h.end_bit, 437, "all PIH fields consumed"); assert!(h.end_bit <= (pih.start + pih.len) * 8, "within the PIH payload"); } + +#[test] +fn tools_header_matches_reference() { + let Some(data) = codestream() else { return }; + let segs = parse_container(&data); + let ton = find_segment(&segs, Marker::Ton).expect("TON present"); + let pih = parse_picture_header(&data, find_segment(&segs, Marker::Pih).unwrap().start_bit()); + + // 4:4:4 ⇒ scale_ver = scale_hor = 3 − c/s = 2 ⇒ the single-plane filter path + let scale_2x2 = (3 - pih.c_ver / pih.s_ver == 2) && (3 - pih.c_hor / pih.s_hor == 2); + assert!(scale_2x2, "test256 is 4:4:4"); + + // The post-filter chain starts 2 bits into TON: `_stream_header_part` defaults to + // "tool_header", so other tools emit into TON ahead of the filters. models_count + // sets the signalled-index width (2 bits here). + let filters_bit = ton.start_bit() + 2; + let t = parse_tools_header(&data, filters_bit, scale_2x2, 3, 1); + eprintln!("{t:#?}"); + + // EFE linear: the two plane means feeding efl_chroma + assert!(t.efl_enabled); + let efl = t.efl.as_ref().unwrap(); + assert_eq!(efl.b1, [12561, 13091], "B1/100 = 125.61 / 130.91 = the EFL plane means"); + let f1 = &efl.filters_best2; + assert_eq!((f1.cand_u, f1.cand_v), (0, 0)); + assert_eq!((f1.fl_u, f1.fl_v), (1, 1)); + assert_eq!((f1.min_symbol, f1.max_symbol), (0, 65535)); + assert_eq!((&f1.weights_u[..], &f1.weights_u2[..]), (&[32677u32][..], &[32772u32][..])); + assert_eq!((&f1.weights_v[..], &f1.weights_v2[..]), (&[32758u32][..], &[32776u32][..])); + let f2 = &efl.filters_best; + // the second group's max_symbol (99) narrows the weights to 7 bits + assert_eq!((f2.min_symbol, f2.max_symbol), (32677, 99)); + assert_eq!((&f2.weights_u[..], &f2.weights_u2[..]), (&[0u32][..], &[95u32][..])); + assert_eq!((&f2.weights_v[..], &f2.weights_v2[..]), (&[81u32][..], &[99u32][..])); + + // eICCI: the per-tile model selection (the last input the eICCI net needed) + assert!(t.icci_enabled); + assert!(!t.icci_tile_manager_enabled); + assert_eq!(t.icci_tiles.len(), 1); + let tile = &t.icci_tiles[0]; + assert_eq!(tile.use_yuv, [true, true, true]); + assert!(tile.use_short_list); + assert_eq!((tile.model_idx_y, tile.model_idx_uv), (0, 1)); + + // EFE nonlinear: exactly the inputs efn_filter consumes + assert!(t.efn_enabled); + let efn = t.efn.as_ref().unwrap(); + assert_eq!((efn.min_symbol, efn.max_symbol), (0, 65535)); + assert!(!efn.mask1_enabled && !efn.mask2_enabled); + assert!(!efn.nonlinear_u_enabled && efn.nonlinear_v_enabled, "only V is filtered"); + assert_eq!((efn.tile_width, efn.tile_height), (256, 256)); + assert_eq!((&efn.luma_min[..], &efn.luma_max[..]), (&[497u32][..], &[25320u32][..]), + "== efn_fixture luma_min_q / luma_max_q"); + assert!(efn.weights_u.is_empty(), "U disabled ⇒ no candNum/A1 for U"); + assert_eq!(efn.weights_v, vec![32966, 32443, 33031, 32361, 33068, 33006, 31968, 33003], + "== efn_fixture weights_q15"); + + assert!(t.lef_enabled); + assert_eq!(t.lef_ch_idx, 0); + + // the whole tools header consumed exactly, up to the pad closing the 63-byte TON + assert_eq!(t.end_bit, 973, "all TON fields consumed"); + assert!(t.end_bit <= (ton.start + ton.len) * 8); +} + +#[test] +fn header_widths_follow_calc_bits_count() { + use wai::jpegai_bitstream::bits_for_max; + // bits = ceil(log2(max+1)) — the reference HeaderCoder::calc_bits_count + for (max, want) in [(0u64, 0u32), (1, 1), (2, 2), (3, 2), (4, 3), (9, 4), (15, 4), (99, 7), (255, 8), (65535, 16)] { + assert_eq!(bits_for_max(max), want, "bits_for_max({max})"); + } +}