41db782eaf
Start the JPEG-AI decode-pipeline glue that wires the byte-exact stages together. - jpegai_pipeline::dequant — resi[c,p] = y_q[c,p] · gain_map[c,p], where the per-position gain map already folds the per-channel gain (scaler_vec) and the residual-variance scale (RVS). Deterministic f32. - e2e entropy→dequant: decode the real test256.bin y-latent (40960 symbols) with the mANS decoder, dequantize, and match the reference resi_dq_y checkpoint — BYTE-EXACT (max|Δ| = 0). Proves the mANS output IS the dequant input. - e2e synthesis on the real reference latent: decode_y(y_lsbs)→ypix at 58.85 dB and decode_uv(concat(y_lsbs,uv_lsbs))→uvpix at 58.88 dB — the full synthesis stage validated on real pipeline data, and the UV input assembly (y‖uv) pinned. All feature-gated on the checkpoint/fixtures. Builds toward decode(bitstream)→pixels. Co-Authored-By: Claude Opus 4.8 <[email protected]>
59 lines
3.3 KiB
Rust
59 lines
3.3 KiB
Rust
//! End-to-end integration: the entropy → dequant spine segment on real `test256.bin`
|
||
//! data. Decode the y-latent symbols with the mANS decoder (`MansDecoder::decode_sgm`
|
||
//! over the real SORP stream), then dequantize by the per-position gain map, and
|
||
//! require BYTE-EXACT (f32) equality to the reference `resi_dq_y` checkpoint. This
|
||
//! proves the mANS output is exactly the dequant input and the dequant arithmetic is
|
||
//! `y_q · gain_map`.
|
||
//!
|
||
//! env: WAI_JPEGAI_MANS_FIXTURE = mans_sorp_fixture.json {dt,bounds,mem,offsets,calls}
|
||
//! WAI_JPEGAI_DEQUANT_FIXTURE = jpegai_dequant_fixture.json {gain_map[f32], resi_dq[f32]}
|
||
#![cfg(feature = "neural_int")]
|
||
|
||
use wai::jpegai_mans::MansDecoder;
|
||
use wai::jpegai_pipeline::dequant;
|
||
|
||
#[test]
|
||
fn entropy_to_dequant_byte_exact() {
|
||
let mf = match std::env::var("WAI_JPEGAI_MANS_FIXTURE") {
|
||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||
_ => { eprintln!("skip: set WAI_JPEGAI_MANS_FIXTURE"); return; }
|
||
};
|
||
let df = match std::env::var("WAI_JPEGAI_DEQUANT_FIXTURE") {
|
||
Ok(p) if std::path::Path::new(&p).exists() => p,
|
||
_ => { eprintln!("skip: set WAI_JPEGAI_DEQUANT_FIXTURE"); return; }
|
||
};
|
||
let m: serde_json::Value = serde_json::from_slice(&std::fs::read(&mf).unwrap()).unwrap();
|
||
let d: serde_json::Value = serde_json::from_slice(&std::fs::read(&df).unwrap()).unwrap();
|
||
|
||
// --- entropy: decode the whole y-latent (160 calls × 256) via mANS ---
|
||
let dt: Vec<u32> = m["dt"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect();
|
||
let bounds: Vec<u8> = m["bounds"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
|
||
let mem: Vec<u8> = m["mem"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
|
||
let offsets: Vec<u32> = m["offsets"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect();
|
||
let mut dec = MansDecoder::new(&mem, &offsets);
|
||
dec.set_sgm_transitions(&dt, &bounds);
|
||
let mut y_q: Vec<i16> = Vec::new();
|
||
for call in m["calls"].as_array().unwrap() {
|
||
let sigma: Vec<u8> = call["sigma"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
|
||
let mask: Vec<bool> = call["mask"].as_array().unwrap().iter().map(|x| x.as_bool().unwrap()).collect();
|
||
let mut out = vec![0i16; sigma.len()];
|
||
dec.decode_sgm(&sigma, &mut out, &mask);
|
||
y_q.extend_from_slice(&out);
|
||
}
|
||
|
||
// --- dequant: resi = y_q · gain_map ---
|
||
let gain_map: Vec<f32> = d["gain_map"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect();
|
||
let resi_ref: Vec<f32> = d["resi_dq"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect();
|
||
let resi = dequant(&y_q, &gain_map);
|
||
|
||
assert_eq!(resi.len(), resi_ref.len(), "size mismatch");
|
||
let (mut maxd, mut mism) = (0f32, 0usize);
|
||
for i in 0..resi.len() {
|
||
let dlt = (resi[i] - resi_ref[i]).abs();
|
||
// exact-in-f32: allow only 1-ULP-scale rounding vs the stored f64->f32 golden
|
||
if dlt > 1e-4 { mism += 1; }
|
||
maxd = maxd.max(dlt);
|
||
}
|
||
eprintln!("entropy→dequant vs reference resi_dq_y: {} vals, max|Δ|={maxd:.2e}, mismatched(>1e-4)={mism}", resi.len());
|
||
assert_eq!(mism, 0, "entropy→dequant must be byte-exact (f32); {mism} mismatches");
|
||
}
|