From 5d8a70753ba3aa5f8a7cf576b20fb85a9addaccf Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 16 Jul 2026 07:06:40 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Fbitstream=20=E2=80=94=20full?= =?UTF-8?q?=20picture-header=20parser=20(parse=5Fpicture=5Fheader)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decodes the complete PIH substream into a PictureHeader struct: the fixed prefix (profile, synthesis transforms, level, dimensions, sub-sampling, colour_transform_idx, model_id), then CCS_SGMM::decode_header (multi_threading_z, the quantizer beta F/O headers with independent_beta_uv, region_partitioning), then the per-component tool groups (multi_threading_r, num_chs, and the tools in sub-tool order: use_cube_flags, lsbs_enable, rvs_enable, grfs_enable + per-channel flags), gain_3D_enable, and synthesis_tile_enable per component. Field order and conditionals ported from the reference decode_header sources (ccs_sgmm_tool, quantization, sep_chan_tool, res_var_scale, skip_mode); notably use_cube_flags precedes lsbs_enable, and synthesis_tile_enable trails the component groups (both differ from the header_spec.json listing order). Validated on real test256.bin: 256x256 4:4:4, colour_transform_idx=1, num_chs 160/96, per-component tool flags, and end_bit=437 — i.e. EVERY field consumed exactly, up to the 3-bit pad closing the 49-byte PIH. The grfs channel-flag counts (4 luma, 9 chroma) independently match the sigma-additions gain fixture. Branches not exercised by this 4:4:4 single-tile stream (region sub-splits, independent beta, signalled colour matrix, multi-threading) follow the reference syntax but are untested. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_bitstream.rs | 128 ++++++++++++++++++ .../tests/jpegai_bitstream_calibration.rs | 45 +++++- 2 files changed, 172 insertions(+), 1 deletion(-) diff --git a/wai/wai-rs/src/jpegai_bitstream.rs b/wai/wai-rs/src/jpegai_bitstream.rs index 741955af5..846ebd4e2 100644 --- a/wai/wai-rs/src/jpegai_bitstream.rs +++ b/wai/wai-rs/src/jpegai_bitstream.rs @@ -137,6 +137,134 @@ pub fn find_segment(segs: &[Segment], marker: Marker) -> Option { segs.iter().copied().find(|s| s.marker == marker) } +/// Per-component (luma / chroma) parameters from the picture header. +#[derive(Clone, Debug, Default)] +pub struct ComponentHeader { + pub num_threads_r: u32, + pub num_chs: u32, + pub use_cube_flags: bool, + pub lsbs_enable: bool, + pub rvs_enable: bool, + pub grfs_enable: bool, + /// GRFS channel-selection flags, `num_chs` of them (empty if `grfs_enable` is false). + pub grfs_channel_flags: Vec, + pub synthesis_tile_enable: bool, +} + +/// The decoded JPEG-AI picture header (PIH substream). +#[derive(Clone, Debug, Default)] +pub struct PictureHeader { + pub stream_profile_idc: u32, + pub decoder_profile_id: u32, + pub synthesis_transform_ids: Vec, + pub level_idc: u32, + pub img_width: u32, + pub img_height: u32, + pub diff_display_width: u32, + pub diff_display_height: u32, + pub bit_depth_idc: u32, + /// Luma/chroma spatial subsampling factors (1 = 4:4:4). + pub s_ver: u32, + pub s_hor: u32, + pub c_ver: u32, + pub c_hor: u32, + /// 0 = none, 1 = built-in YUV↔RGB, 2 = signalled 3×3 matrix (+ offset). + pub colour_transform_idx: u32, + pub colour_transform_matrix: Option<[u32; 9]>, + pub colour_transform_offset: Option<[u32; 3]>, + pub model_id: u32, + pub num_threads_z: u32, + /// `beta_displacement_log` (already de-biased by 2048). + pub beta_displacement_log: [i32; 2], + pub independent_beta_uv: bool, + pub region_partitioning: bool, + /// Two components: `[0]` = luma, `[1]` = chroma. + pub components: [ComponentHeader; 2], + pub gain_3d_enable: bool, + /// Bit position just past the last field (before byte-alignment padding). + pub end_bit: usize, +} + +/// Parse the picture header (PIH payload). `start_bit` is the payload bit offset +/// (`Segment::start_bit`). A faithful port of the reference decode order: +/// prefix → `CCS_SGMM::decode_header` (`multi_threading_z`, quantizer β F-header, +/// `region_partitioning`, quantizer β O-header) → the per-component tool groups +/// (`multi_threading_r`, `num_chs`, then the tools in sub-tool order: +/// `use_cube_flags`, `lsbs_enable`, `rvs_enable`, `grfs_enable` + channel flags) → +/// `gain_3D_enable` → `synthesis_tile_enable` per component. +/// +/// 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. +pub fn parse_picture_header(data: &[u8], start_bit: usize) -> PictureHeader { + let mut br = BitReader::at(data, start_bit); + let mut h = PictureHeader::default(); + + h.stream_profile_idc = br.read(4) as u32; + h.decoder_profile_id = br.read(4) as u32; + let num_transforms = br.read(4) as u32 + 1; + h.synthesis_transform_ids = (0..num_transforms).map(|_| br.read(4) as u32).collect(); + h.level_idc = br.read(8) as u32; + h.img_width = br.read(16) as u32 + 64; + h.img_height = br.read(16) as u32 + 64; + h.diff_display_width = br.read(6) as u32; + h.diff_display_height = br.read(6) as u32; + h.bit_depth_idc = br.read(3) as u32; + h.s_ver = br.read(1) as u32 + 1; + h.s_hor = br.read(1) as u32 + 1; + h.c_ver = br.read(1) as u32 + 1; + h.c_hor = br.read(1) as u32 + 1; + h.colour_transform_idx = br.read(2) as u32; + if h.colour_transform_idx == 2 { + let mut m = [0u32; 9]; + for v in m.iter_mut() { *v = br.read(8) as u32; } + let mut o = [0u32; 3]; + for v in o.iter_mut() { *v = br.read(8) as u32; } + h.colour_transform_matrix = Some(m); + h.colour_transform_offset = Some(o); + } + h.model_id = br.read(4) as u32; + + // CCS_SGMM::decode_header + h.num_threads_z = if br.read(1) == 1 { 1 << (br.read(2) + 1) } else { 1 }; + h.beta_displacement_log[0] = br.read(12) as i32 - 2048; // quantizer β F-header + h.region_partitioning = br.read(1) == 1; + if h.region_partitioning { + let _num_ver = br.read(7) + 1; + let _num_hor = br.read(7) + 1; + if br.read(1) == 0 { + let _hyper_overlap = br.read(2); + let _mcm_overlap = br.read(4); + } + } + h.independent_beta_uv = br.read(1) == 1; // quantizer β O-header + h.beta_displacement_log[1] = if h.independent_beta_uv { + br.read(12) as i32 - 2048 + } else { + h.beta_displacement_log[0] + }; + + // per-component tool groups (0 = luma, 1 = chroma) + for c in 0..2 { + let comp = &mut h.components[c]; + comp.num_threads_r = if br.read(1) == 1 { 1 << (br.read(2) + 1) } else { 1 }; + comp.num_chs = br.read(8) as u32; + comp.use_cube_flags = br.read(1) == 1; + comp.lsbs_enable = br.read(1) == 1; + comp.rvs_enable = br.read(1) == 1; + comp.grfs_enable = br.read(1) == 1; + if comp.grfs_enable { + comp.grfs_channel_flags = (0..comp.num_chs).map(|_| br.read(1) == 1).collect(); + } + } + h.gain_3d_enable = br.read(1) == 1; + for c in 0..2 { + h.components[c].synthesis_tile_enable = br.read(1) == 1; + } + h.end_bit = br.bit_pos(); + h +} + #[cfg(test)] mod tests { use super::*; diff --git a/wai/wai-rs/tests/jpegai_bitstream_calibration.rs b/wai/wai-rs/tests/jpegai_bitstream_calibration.rs index 1ce4f8bd7..22b7bd9cf 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, BitReader, Marker}; +use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, BitReader, Marker}; fn codestream() -> Option> { match std::env::var("WAI_JPEGAI_CODESTREAM") { @@ -83,3 +83,46 @@ fn picture_header_leading_fields_match_reference() { assert_eq!(model_id, 0); assert_eq!(br.bit_pos(), 137, "cursor lands where the reference trace does"); } + +#[test] +fn full_picture_header_matches_reference() { + let Some(data) = codestream() else { return }; + let segs = parse_container(&data); + let pih = find_segment(&segs, Marker::Pih).expect("PIH present"); + let h = parse_picture_header(&data, pih.start_bit()); + eprintln!("{h:#?}"); + + assert_eq!(h.stream_profile_idc, 0); + assert_eq!(h.decoder_profile_id, 2); + assert_eq!(h.synthesis_transform_ids, vec![2, 1, 0]); + assert_eq!(h.level_idc, 52); + assert_eq!((h.img_width, h.img_height), (256, 256)); + assert_eq!(h.bit_depth_idc, 0); + assert_eq!((h.s_ver, h.s_hor, h.c_ver, h.c_hor), (1, 1, 1, 1), "4:4:4"); + assert_eq!(h.colour_transform_idx, 1); + assert!(h.colour_transform_matrix.is_none(), "idx=1 carries no matrix"); + assert_eq!(h.model_id, 0); + assert_eq!(h.num_threads_z, 1); + assert_eq!(h.beta_displacement_log, [0, 0], "2048 de-biased, shared across Y/UV"); + assert!(!h.independent_beta_uv); + assert!(!h.region_partitioning); + assert!(!h.gain_3d_enable); + + // per-component (luma 160ch, chroma 96ch); all four tool flags as decoded + assert_eq!(h.components[0].num_chs, 160); + assert_eq!(h.components[1].num_chs, 96); + for c in 0..2 { + let comp = &h.components[c]; + assert!(!comp.use_cube_flags, "c{c} use_cube"); + assert!(comp.lsbs_enable && comp.rvs_enable && comp.grfs_enable, "c{c} lsbs/rvs/grfs"); + assert_eq!(comp.grfs_channel_flags.len(), comp.num_chs as usize, "c{c} grfs flag count == num_chs"); + assert!(!comp.synthesis_tile_enable, "c{c} tile"); + } + // the grfs channel-selection counts match the sigma-additions gain fixture + assert_eq!(h.components[0].grfs_channel_flags.iter().filter(|&&b| b).count(), 4); + assert_eq!(h.components[1].grfs_channel_flags.iter().filter(|&&b| b).count(), 9); + + // the whole header consumed exactly to the padding boundary of the 49-byte PIH + assert_eq!(h.end_bit, 437, "all PIH fields consumed"); + assert!(h.end_bit <= (pih.start + pih.len) * 8, "within the PIH payload"); +} -- 2.54.0 (Apple Git-157)