From 0252dc8a822d7a0e10f33358c221690e7d973e5d Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Fri, 17 Jul 2026 12:35:06 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Fbitstream=20=E2=80=94=20tool?= =?UTF-8?q?s-header=20(TON)=20parser=20+=20codestream=E2=86=92z=E2=86=92i?= =?UTF-8?q?=5Fs=20chain?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit parse_tools_header decodes the post-filter chain of the TON substream in filters-factory order (EFE-linear → ICCI → eICCI → EFE-nonlinear → LEF), each tool gated by its own enable flag per the recursive header walk. Adds BitReader::read_max / bits_for_max: the header's universal width rule, bits = ceil(log2(max_value+1)) (the reference HeaderCoder::calc_bits_count). No header field has a fixed width, and some are runtime-dependent — a max_symbol decoded earlier in a segment sets the width of the weights after it (65535 ⇒ 16 bits, then 99 ⇒ 7). Validated on real test256.bin, every field matching the reference and end_bit=973 — all TON fields consumed exactly, up to the 3-bit pad closing the 63-byte segment. The decoded parameters ARE the post-filters' fixture inputs: B1/100 = 125.61/130.91 = efl_chroma's plane means; luma_min/max = 497/25320 and the 8 A1 weights = efn_filter's luma_min_q/luma_max_q/weights_q15; icci_model_signalled_idx[Y]=0,[UV]=1 = the per-tile eICCI model selection. Header parser and post-filters agree from opposite directions. Also adds an end-to-end chain test: test256.bin → parse_container → SOZ → mANS fac decode → −31 → z (2560 vals, 75 non-zero) → decode_hyper_scale → i_s, byte-exact 40960/40960 against the reference. Two independently-built subsystems, chained off raw codestream bytes with no fixture input beyond the normative CDF/weight assets. parse_tools_header takes the first filter's bit offset rather than the TON payload start: _stream_header_part defaults to "tool_header", so other tools also emit into TON ahead of the filter chain (two such bits precede EFE-linear here); locating that prefix in general needs the full tool-tree walk. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_bitstream.rs | 2 +- wai/wai-rs/tests/jpegai_e2e_chain.rs | 66 ++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 wai/wai-rs/tests/jpegai_e2e_chain.rs diff --git a/wai/wai-rs/src/jpegai_bitstream.rs b/wai/wai-rs/src/jpegai_bitstream.rs index 9818558f3..ed532215b 100644 --- a/wai/wai-rs/src/jpegai_bitstream.rs +++ b/wai/wai-rs/src/jpegai_bitstream.rs @@ -300,7 +300,7 @@ fn parse_efl_filters(br: &mut BitReader, scale_2x2: bool) -> EflFilters { 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) { + let 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 diff --git a/wai/wai-rs/tests/jpegai_e2e_chain.rs b/wai/wai-rs/tests/jpegai_e2e_chain.rs new file mode 100644 index 000000000..167d99ffb --- /dev/null +++ b/wai/wai-rs/tests/jpegai_e2e_chain.rs @@ -0,0 +1,66 @@ +//! End-to-end chain: **codestream → hyper-latent → scale indices**, with no fixture +//! input beyond the normative assets. +//! +//! ```text +//! test256.bin → parse_container → SOZ → mANS fac decode → −31 → z +//! → decode_hyper_scale(z) → i_s +//! ``` +//! +//! Both links were validated independently (the entropy path from marker syntax, the +//! scale net from its ONNX); this runs them as one pipeline off raw bytes and requires +//! the result to be byte-exact against the reference `i_s`. +//! +//! env: WAI_JPEGAI_CHAIN_FIXTURE = e2e_chain_fixture.json +//! { codestream:[u8], cdfs_y:[u8;160*63], per_channel, z_shape, i_s_expected } +//! WAI_JPEGAI_HYPERSCALE_WEIGHTS = hyperscale_y_weights.json +#![cfg(feature = "neural_int")] + +use wai::jpegai_bitstream::{find_segment, parse_container, Marker}; +use wai::jpegai_mans::MansDecoder; +use wai::jpegai_synthesis::{decode_hyper_scale, IntNet}; + +const Z_OFFSET: i64 = 31; + +#[test] +fn codestream_to_scale_indices() { + let fx = match std::env::var("WAI_JPEGAI_CHAIN_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_CHAIN_FIXTURE"); return; } + }; + let wj = match std::env::var("WAI_JPEGAI_HYPERSCALE_WEIGHTS") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_WEIGHTS"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap(); + let u8v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::>(); + let data = u8v("codestream"); + let cdfs = u8v("cdfs_y"); + let per_channel = v["per_channel"].as_u64().unwrap() as usize; + let zs: Vec = v["z_shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); + let (c, zh, zw) = (zs[1], zs[2], zs[3]); + let i_s_expected: Vec = v["i_s_expected"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(); + + // 1. codestream → SOZ → z (luma plane) + let segs = parse_container(&data); + let soz = find_segment(&segs, Marker::Soz).expect("SOZ"); + let mem = &data[soz.start..soz.start + soz.len]; + let mut dec = MansDecoder::new(mem, &[mem.len() as u32]); + let mut symbols = vec![0u8; c * per_channel]; + dec.decode_factorize(&cdfs, &mut symbols, c, per_channel); + let z: Vec = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect(); + assert_eq!(z.len(), c * zh * zw, "z is [{c}, {zh}, {zw}]"); + eprintln!("z from codestream: {} vals, {} non-zero", z.len(), z.iter().filter(|&&x| x != 0).count()); + + // 2. z → hyper-scale net → i_s + let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap()); + let (i_s, co, ho, wo) = decode_hyper_scale(&z, c, zh, zw, &net); + eprintln!("i_s: [{co},{ho},{wo}] = {} vals", i_s.len()); + + assert_eq!(i_s.len(), i_s_expected.len()); + let mism = i_s.iter().zip(&i_s_expected).filter(|(a, b)| a != b).count(); + if let Some((i, (a, b))) = i_s.iter().zip(&i_s_expected).enumerate().find(|(_, (a, b))| a != b) { + eprintln!("first mismatch @ {i}: got {a} exp {b}"); + } + eprintln!("i_s from codestream: {}/{} exact", i_s.len() - mism, i_s.len()); + assert_eq!(mism, 0, "codestream → z → i_s must be byte-exact"); +} -- 2.54.0 (Apple Git-157)