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<u32>, Vec<u32>) {
+    let group = |cand: i32, fl: u32, br: &mut BitReader| -> (Vec<u32>, Vec<u32>) {
         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

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::<Vec<u8>>();
+    let data = u8v("codestream");
+    let cdfs = u8v("cdfs_y");
+    let per_channel = v["per_channel"].as_u64().unwrap() as usize;
+    let zs: Vec<usize> = 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<i64> = 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<i64> = 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");
+}