wai/wai-rs/src/jpegai_icci.rs

@@ -242,6 +242,57 @@ impl IcciPath {
     }
 }
 
+/// The whole eICCI filter for one tile: three planes in, three refined planes out.
+///
+/// Shape follows the exported model, **not** the reference Python class (which says
+/// `nf=64, nbY=8, nbUV=8`): the shipped nets are `nf=48`, **2** luma blocks and **4**
+/// chroma blocks. The chroma trunk is shared — the reference traces it once per chroma
+/// plane, computing identical features twice; we evaluate it once and feed both heads.
+///
+/// ```text
+/// (Y16,U16,V16), yuv48 = dwt2(Y), dwt2(U), dwt2(V), cat(..)
+/// Y_o = idwt2( conv_last_Y(luma_trunk(yuv48))  · scale + Y16 )
+/// fea = chroma_trunk(yuv48)                      // evaluated once
+/// U_o = idwt2( conv_last_u(fea) · scale + U16 )
+/// V_o = idwt2( conv_last_v(fea) · scale + V16 )
+/// ```
+#[derive(Clone)]
+pub struct IcciNet {
+    /// luma: `conv_first_Y` (BN pre-folded by the exporter) + its residual blocks
+    pub luma: IcciPath,
+    /// chroma trunk: `conv_firstUV` (fold `BNUV` in) + its residual blocks; `conv_last`
+    /// is the U head and `scale` applies to both heads.
+    pub chroma: IcciPath,
+    /// the V head — shares `chroma`'s trunk, differs only in this final convolution
+    pub conv_last_v: Conv,
+}
+
+impl IcciNet {
+    /// Refine one tile's three planes. `y`/`u`/`v` are `[h × w]` (4:4:4); `h` and `w`
+    /// must be multiples of 4 (two wavelet levels).
+    pub fn forward(&self, y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize)
+        -> (Vec<f32>, Vec<f32>, Vec<f32>) {
+        let (y16, u16, v16, yuv48, h2, w2) = icci_dwt_planes(y, u, v, h, w);
+
+        let y_o = self.luma.forward_plane(&yuv48, &y16, h2, w2);
+
+        // shared chroma trunk, evaluated once
+        let mut fea = self.chroma.conv_first.forward(&yuv48, h2, w2);
+        for t in fea.iter_mut() { if *t < 0.0 { *t = 0.0; } }
+        for blk in &self.chroma.blocks {
+            fea = blk.forward(&fea, h2, w2);
+        }
+        let head = |conv: &Conv, band: &[f32]| -> Vec<f32> {
+            let res = conv.forward(&fea, h2, w2);
+            let refined: Vec<f32> = res.iter().zip(band).map(|(&r, &s)| r * self.chroma.scale + s).collect();
+            haar_idwt2(&refined, h2, w2).0
+        };
+        let u_o = head(&self.chroma.conv_last, &u16);
+        let v_o = head(&self.conv_last_v, &v16);
+        (y_o, u_o, v_o)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

wai/wai-rs/tests/jpegai_icci_calibration.rs

@@ -0,0 +1,112 @@
+//! 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.
+//!
+//! 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.
+//!
+//! 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
+//!
+//! env: WAI_JPEGAI_ICCI_WEIGHTS = icci_model0.json  { name: {dims, v:[f32]} }
+//!      WAI_JPEGAI_ICCI_ONNX    = post_filters/model_0.onnx
+#![cfg(feature = "neural")]
+
+use ort::session::{builder::GraphOptimizationLevel, Session};
+use ort::value::Value;
+use wai::jpegai_icci::{Conv, IcciNet, 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 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 }
+}
+
+#[test]
+fn icci_net_matches_onnx_reference() {
+    let wj = match std::env::var("WAI_JPEGAI_ICCI_WEIGHTS") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_WEIGHTS to icci_model0.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;
+
+    // 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 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),
+    };
+
+    // 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 (ry, ru, rv) = net.forward(&y, &u, &vv, h, w);
+
+    // 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})");
+    }
+}