wai-rs: TON tools-header parser + codestream→z→i_s chain #21
2 changed files with 67 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
66
wai/wai-rs/tests/jpegai_e2e_chain.rs
Normal file
66
wai/wai-rs/tests/jpegai_e2e_chain.rs
Normal file
|
|
@ -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");
|
||||
}
|
||||
Loading…
Reference in a new issue