diff --git a/wai/wai-rs/tests/jpegai_icci_calibration.rs b/wai/wai-rs/tests/jpegai_icci_calibration.rs index 3afb86a74..974e4e0d5 100644 --- a/wai/wai-rs/tests/jpegai_icci_calibration.rs +++ b/wai/wai-rs/tests/jpegai_icci_calibration.rs @@ -1,112 +1,66 @@ -//! Calibration for the eICCI post-filter ([`wai::jpegai_icci::IcciNet`]) against the -//! shipped ONNX model, which is the reference: it takes the three planes and returns the -//! three refined planes, with the wavelet transform inside the graph. +//! Calibration for [`wai::jpegai_icci`] — the eICCI post-filter's plane path — against +//! GROUND TRUTH captured from the ISO reference decoder running on `test256.bin`. //! -//! The network is float, so this is a faithful f32 port rather than a quantised one — -//! we expect a machine-precision match, not a PTQ-floor one. +//! The fixture holds the reference's own `process_dwt444_2` outputs (the 48-channel +//! network input and the plane's 16 subbands) together with the selected model's +//! weights, and the reference's `process_comp` output. Weights arrive with BatchNorm +//! already folded (`w·γ/√(σ²+ε)`, `(b−μ)·γ/√(σ²+ε)+β`, ε=1e-5) — the same fold +//! [`Conv::fold_bn`] performs. //! -//! Weight layout (traced from the export, which dedupes identical tensors and pre-folds -//! most BatchNorms): -//! • `onnx::Conv_8299` luma stem, 48→48 3×3 (BN already folded) -//! • `onnx::Conv_8302/8305`, `8308/8311` the 2 luma residual blocks ([1,3] then [3,1]) -//! • `conv_last_Y` luma head, 48→16 3×3 -//! • `conv_firstUV` + `BNUV` chroma stem (BNUV is γ=1,β=0,μ=0,σ²=1 ⇒ x/√(1+ε)) -//! • `8314/8317 … 8332/8335` the 4 chroma residual blocks -//! • `conv_last_u`, `conv_last_v` the chroma heads -//! • every bias is zero; every scale is `scaleY` (0.1), deduped by the exporter +//! NOTE: the weights must come from the reference's PyTorch `state_dict`, NOT from +//! `post_filters/model_*.onnx` — that graph wraps `forward`, which the decoder never +//! calls, and its initializers are pre-folded and incomplete (no BNY, no scaleU/scaleV, +//! no per-residual-block scale). //! -//! env: WAI_JPEGAI_ICCI_WEIGHTS = icci_model0.json { name: {dims, v:[f32]} } -//! WAI_JPEGAI_ICCI_ONNX = post_filters/model_0.onnx -#![cfg(feature = "neural")] +//! env: WAI_JPEGAI_ICCI_FIXTURE = icci_y_fixture.json +#![cfg(feature = "neural_int")] -use ort::session::{builder::GraphOptimizationLevel, Session}; -use ort::value::Value; -use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock}; +use wai::jpegai_icci::{Conv, IcciPath, ResBlock}; -const EPS: f32 = 1e-5; // ONNX BatchNormalization default - -fn conv_from(v: &serde_json::Value, name: &str, ph: usize, pw: usize) -> Conv { - let t = &v[name]; - let dims: Vec = t["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); - let w: Vec = t["v"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(); - let (cout, cin, kh, kw) = (dims[0], dims[1], dims[2], dims[3]); - Conv { w, b: vec![0f32; cout], cin, cout, kh, kw, ph, pw } +fn f32v(v: &serde_json::Value) -> Vec { + v.as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect() } - -fn block(v: &serde_json::Value, a: &str, b: &str, scale: f32) -> ResBlock { - ResBlock { conv1: conv_from(v, a, 0, 1), conv2: conv_from(v, b, 1, 0), scale } +fn conv_of(v: &serde_json::Value) -> Conv { + let g = |k: &str| v[k].as_u64().unwrap() as usize; + Conv { w: f32v(&v["w"]), b: f32v(&v["b"]), cin: g("cin"), cout: g("cout"), + kh: g("kh"), kw: g("kw"), ph: g("ph"), pw: g("pw") } } #[test] -fn icci_net_matches_onnx_reference() { - let wj = match std::env::var("WAI_JPEGAI_ICCI_WEIGHTS") { +fn icci_plane_path_matches_reference() { + let path = match std::env::var("WAI_JPEGAI_ICCI_FIXTURE") { Ok(p) if std::path::Path::new(&p).exists() => p, - _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_WEIGHTS to icci_model0.json"); return; } + _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_FIXTURE to icci_y_fixture.json"); return; } }; - let onnx = match std::env::var("WAI_JPEGAI_ICCI_ONNX") { - Ok(p) if std::path::Path::new(&p).exists() => p, - _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_ONNX to model_0.onnx"); return; } - }; - let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap(); - let scale = v["scaleY"]["v"][0].as_f64().unwrap() as f32; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let (h, w) = (v["h"].as_u64().unwrap() as usize, v["w"].as_u64().unwrap() as usize); - // chroma stem: fold BNUV (γ=1, β=0, μ=0, σ²=1) into conv_firstUV → w/√(1+ε) - let mut chroma_stem = conv_from(&v, "conv_firstUV.weight", 1, 1); - let n = chroma_stem.cout; - chroma_stem = chroma_stem.fold_bn(&vec![1.0; n], &vec![0.0; n], &vec![0.0; n], &vec![1.0; n], EPS); + let blocks: Vec = v["blocks"].as_array().unwrap().iter().map(|b| ResBlock { + conv1: conv_of(&b["c1"]), + conv2: conv_of(&b["c2"]), + scale: b["scale"].as_f64().unwrap() as f32, + }).collect(); + assert_eq!(blocks.len(), 2, "the Y path has two residual blocks"); - let net = IcciNet { - luma: IcciPath { - conv_first: conv_from(&v, "onnx::Conv_8299", 1, 1), - blocks: vec![ - block(&v, "onnx::Conv_8302", "onnx::Conv_8305", scale), - block(&v, "onnx::Conv_8308", "onnx::Conv_8311", scale), - ], - conv_last: conv_from(&v, "conv_last_Y.weight", 1, 1), - scale, - }, - chroma: IcciPath { - conv_first: chroma_stem, - blocks: vec![ - block(&v, "onnx::Conv_8314", "onnx::Conv_8317", scale), - block(&v, "onnx::Conv_8320", "onnx::Conv_8323", scale), - block(&v, "onnx::Conv_8326", "onnx::Conv_8329", scale), - block(&v, "onnx::Conv_8332", "onnx::Conv_8335", scale), - ], - conv_last: conv_from(&v, "conv_last_u.weight", 1, 1), - scale, - }, - conv_last_v: conv_from(&v, "conv_last_v.weight", 1, 1), + let path_net = IcciPath { + conv_first: conv_of(&v["conv_first"]), + blocks, + conv_last: conv_of(&v["conv_last"]), + scale: v["scale"].as_f64().unwrap() as f32, }; - // deterministic planes at a tile size that is a multiple of 4 - let (h, w) = (32usize, 32usize); - let plane = h * w; - let mk = |seed: usize| -> Vec { - (0..plane).map(|i| (((i * 37 + seed * 11) % 251) as f32 - 125.0) / 255.0).collect() - }; - let (y, u, vv) = (mk(0), mk(1), mk(2)); + let yuv48 = f32v(&v["yuv48"]); + let sub16 = f32v(&v["sub16"]); + let expected = f32v(&v["expected"]); + assert_eq!(yuv48.len(), 48 * h * w, "network input is 48 channels at h×w"); + assert_eq!(sub16.len(), 16 * h * w, "plane subbands are 16 channels"); - let (ry, ru, rv) = net.forward(&y, &u, &vv, h, w); + let got = path_net.forward_plane(&yuv48, &sub16, h, w); + assert_eq!(got.len(), expected.len(), "output is the full-resolution plane"); - // reference: the ONNX itself (planes in → refined planes out) - let mut s = Session::builder().unwrap() - .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap() - .with_optimization_level(GraphOptimizationLevel::Level3).unwrap() - .commit_from_file(&onnx).unwrap(); - let arr = |d: &Vec| Value::from_array(ndarray::Array4::from_shape_vec((1, 1, h, w), d.clone()).unwrap()).unwrap().into_dyn(); - let out = s.run(vec![ - ("Y".to_string(), arr(&y)), - ("U".to_string(), arr(&u)), - ("V".to_string(), arr(&vv)), - ]).unwrap(); - - for (i, (name, mine)) in [("Y_o", &ry), ("U_o", &ru), ("V_o", &rv)].iter().enumerate() { - let want: Vec = out[i].try_extract_array::().unwrap().iter().copied().collect(); - assert_eq!(mine.len(), want.len(), "{name}: size mismatch"); - let mut maxd = 0f32; - for j in 0..want.len() { maxd = maxd.max((mine[j] - want[j]).abs()); } - eprintln!("eICCI {name} vs ONNX: max|Δ| = {maxd:.7}"); - assert!(maxd < 1e-3, "{name}: eICCI diverged from the ONNX (max|Δ| = {maxd})"); - } + let mut maxd = 0f32; + for i in 0..got.len() { maxd = maxd.max((got[i] - expected[i]).abs()); } + eprintln!("eICCI plane path vs reference process_comp: max|Δ| = {maxd:.3e}"); + // float arithmetic on a ~100k-MAC-per-pixel path; ε-level agreement expected + assert!(maxd < 1e-4, "eICCI plane path diverged from the reference (max|Δ| = {maxd})"); }