wai-rs: validate jpegai_icci against the reference decoder (ground truth, 2.4e-07) #26
1 changed files with 48 additions and 94 deletions
|
|
@ -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<usize> = t["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
|
||||
let w: Vec<f32> = 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<f32> {
|
||||
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<ResBlock> = 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<f32> {
|
||||
(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<f32>| 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<f32> = out[i].try_extract_array::<f32>().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})");
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue