wai/wai-rs/src/jpegai_bitstream.rs

@@ -137,6 +137,134 @@ pub fn find_segment(segs: &[Segment], marker: Marker) -> Option<Segment> {
     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<bool>,
+    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<u32>,
+    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::*;

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<Vec<u8>> {
     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");
+}