open-standards/wai/wai-rs/tests/jpegai_icci_net_calibration.rs

//! Whole-filter calibration for **eICCI** — [`wai::jpegai_icci::IcciNet::forward`] —
//! against the reference's own per-filter dump, on the real `test256` YUV.
//!
//! `IcciPath::forward_plane` is checked separately (`jpegai_icci_calibration`) against
//! `process_comp`, but only per plane and only from pre-computed inputs. This exercises
//! the whole filter: the two-level Haar transform of all three planes into the
//! 48-channel network input, the luma path, the shared chroma trunk with its two heads,
//! and the inverse transform back.
//!
//! Three things it pins down that a per-plane test cannot:
//!
//!   * **Domain.** eICCI runs on [0,1] data — the reference normalises by 255 before the
//!     filter and scales back after. Feeding raw 0..255 breaks the folded BatchNorms and
//!     the ReLUs, and the network then scores *worse than doing nothing*.
//!   * **Depth differs per plane**: 2 residual blocks for luma, 4 for chroma.
//!   * **U and V carry separate residual scales**; sharing one shifts V past a pixel LSB.
//!
//! Every plane is reported against the identity baseline (the filter's own input). eICCI
//! is a real correction here — max|Δ| of 6.45 / 2.56 / 2.72 on Y / U / V — so the output
//! has to sit far closer to the reference than the input does. Without that the test
//! would pass on a no-op, which is exactly how the domain bug first looked plausible.
//!
//! env: WAI_JPEGAI_ICCI_NET_FIXTURE = icci_net_fixture.json
//!   { h, w, luma:{conv_first,blocks[2],conv_last,scale}, chroma:{...,blocks[4]},
//!     conv_last_v, scale_v, in_y/u/v:[f32], out_y/u/v:[f32] }
//!   Weights come from `wai_jpegai_export_weights.py icci`; the in/out planes are the
//!   reference `e2e/pfchain/EfficientICCIFilter_{in,out}_0.npy` dumps.
#![cfg(feature = "neural_int")]

use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};

fn f32v(v: &serde_json::Value) -> Vec<f32> {
    v.as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect()
}
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") }
}
fn path_of(v: &serde_json::Value) -> IcciPath {
    IcciPath {
        conv_first: conv_of(&v["conv_first"]),
        blocks: 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(),
        conv_last: conv_of(&v["conv_last"]),
        scale: v["scale"].as_f64().unwrap() as f32,
    }
}

#[test]
fn icci_net_matches_reference_filter() {
    let path = match std::env::var("WAI_JPEGAI_ICCI_NET_FIXTURE") {
        Ok(p) if std::path::Path::new(&p).exists() => p,
        _ => { eprintln!("skip: set WAI_JPEGAI_ICCI_NET_FIXTURE"); return; }
    };
    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);

    let luma = path_of(&v["luma"]);
    let chroma = path_of(&v["chroma"]);
    assert_eq!(luma.blocks.len(), 2, "the luma path has 2 residual blocks");
    assert_eq!(chroma.blocks.len(), 4, "the chroma path has 4 residual blocks");
    let net = IcciNet {
        luma,
        chroma,
        conv_last_v: conv_of(&v["conv_last_v"]),
        scale_v: v["scale_v"].as_f64().unwrap() as f32,
    };

    // The dumps are in [0,255]; the network is not. Normalise in, scale out.
    let raw = |k: &str| f32v(&v[k]);
    let (in_y, in_u, in_v) = (raw("in_y"), raw("in_u"), raw("in_v"));
    let norm = |p: &[f32]| p.iter().map(|&x| x / 255.0).collect::<Vec<f32>>();
    let (gy, gu, gv) = net.forward(&norm(&in_y), &norm(&in_u), &norm(&in_v), h, w);
    let denorm = |p: Vec<f32>| p.into_iter().map(|x| x * 255.0).collect::<Vec<f32>>();
    let (gy, gu, gv) = (denorm(gy), denorm(gu), denorm(gv));

    for (name, got, exp, inp) in [
        ("Y", &gy, raw("out_y"), &in_y),
        ("U", &gu, raw("out_u"), &in_u),
        ("V", &gv, raw("out_v"), &in_v),
    ] {
        let maxd = got.iter().zip(&exp).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
        let base = inp.iter().zip(&exp).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
        eprintln!("eICCI {name}: max|Δ| vs reference = {maxd:.4e}   (identity baseline {base:.4})");
        assert!(base > 0.5, "{name}: eICCI barely moves this plane — the test is not discriminating");
        assert!(maxd < 1e-2,
                "{name}: IcciNet diverged from the reference filter (max|Δ| = {maxd}); \
                 if it is at or above the identity baseline, check the [0,1] domain");
    }
}