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

//! **The whole decoder, end to end**, through the library entry point
//! [`wai::jpegai_decode::decode`]: `test256.bin` → RGB, compared to the reference
//! `test256_dec.png`. No per-image fixture — every stage runs off the codestream plus
//! normative assets and the sink's model weights.
//!
//! ```text
//! codestream ─┬─ luma:   SOZ→z→i_s→σ/mask→SORP→y_q→resi→ψ→MCM→ŷ→LSBS→decode_y  ─┐
//!             └─ chroma: SOZ→z→i_s→σ/mask→SORS→y_q→resi→ψ→(no MCM)→ŷ→LSBS→decode_uv ┤
//!                                                                                  │
//!   YUV → EFElinear → eICCI → EFEnonlinear → LEF → BT.709 → RGB ───────────────────┘
//! ```
//!
//! Every stage was validated in isolation (the `jpegai_*_calibration` and
//! `jpegai_e2e_{frontend,chroma}` tests). This runs them *composed* through the one
//! public `decode()` call, which is what catches an inter-stage domain or wiring
//! mismatch — the class of bug no per-stage test can see.
//!
//! The result is ~1 LSB off the float reference on a few percent of samples: ψ, the MCM
//! stages and synthesis are 16-bit PTQ. The point is a pinned reconstruction, not
//! bit-equality with the reference PNG — so the asserted bar is max|Δ| ≤ 2 and at most
//! 5% of samples differing (recorded: max|Δ| = 1 on ~4.4%). This test runs on one
//! machine; it does not compare decoded output across machines.
//!
//! env: WAI_JPEGAI_CODESTREAM = test256.bin
//!      WAI_JPEGAI_WEIGHTS_DIR = dir of {hs,hyp}_{y,uv}.json, mcm{0..3}_y.json,
//!                                       synth_{y,uv}.json, icci.json
//!      WAI_JPEGAI_NORM = norm.json  (see wai_jpegai_build_e2e_inputs.py)
//!
//! These inputs are not in the repository. They were extracted by
//! `tools/wai_jpegai_build_e2e_inputs.py` from JPEG-AI reference-software assets
//! (ISO/IEC 6048-1; kept locally in `wai/jpegai-vm-assets`, not committed), and the
//! extraction is published under WAI's versioned models root (`jpegai/v1/`).
//! `tools/wai_jpegai_fetch_e2e_inputs.sh` fetches it and checks every file against
//! the SHA-256 manifest `tools/jpegai-e2e-v1.sha256`; CI (`wai-jpegai-e2e.yml`) does
//! that and runs this file with `WAI_REQUIRE_SOTA_VECTORS=1`. Without the three
//! variables it skips loudly; with the switch it fails instead; a variable naming a
//! missing path, or a weights directory missing a file, fails either way.
#![cfg(feature = "neural_int")]

#[path = "support/require_vectors.rs"]
mod require_vectors;

use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_icci::{Conv, IcciNet, IcciPath, ResBlock};
use wai::jpegai_synthesis::{IntNet, Weights};

fn jload(p: &std::path::Path) -> serde_json::Value {
    serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
}
fn w_of(dir: &str, name: &str) -> Weights {
    Weights::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn int_of(dir: &str, name: &str) -> IntNet {
    IntNet::from_json(&jload(&std::path::Path::new(dir).join(name))).unwrap()
}
fn u8v(v: &serde_json::Value, k: &str) -> Vec<u8> {
    v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect()
}
fn i64v(v: &serde_json::Value, k: &str) -> Vec<i64> {
    v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
}
fn i32v(v: &serde_json::Value, k: &str) -> Vec<i32> {
    v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect()
}
fn f3(v: &serde_json::Value, k: &str) -> [f32; 3] {
    let a = v[k].as_array().unwrap();
    [a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32]
}
fn i3(v: &serde_json::Value, k: &str) -> [i64; 3] {
    let a = v[k].as_array().unwrap();
    [a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()]
}

fn conv_of(c: &serde_json::Value) -> Conv {
    let g = |k: &str| c[k].as_u64().unwrap() as usize;
    Conv {
        w: c["w"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
        b: c["b"].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect(),
        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 codestream_to_rgb_matches_reference_png() {
    let Some([cs, dir, nf]) = require_vectors::env_paths(
        "jpegai_e2e_full::codestream_to_rgb_matches_reference_png",
        ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"],
        "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)",
    ) else {
        return;
    };
    let data = std::fs::read(&cs).unwrap_or_else(|e| panic!("read {cs}: {e}"));
    let nv = jload(std::path::Path::new(&nf));

    let iw = jload(&std::path::Path::new(&dir).join("icci.json"));
    let icci = IcciNet {
        luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]),
        conv_last_v: conv_of(&iw["conv_last_v"]),
        scale_v: iw["scale_v"].as_f64().unwrap() as f32,
    };
    let weights = DecoderWeights {
        luma: PlaneWeights {
            hyper_scale: int_of(&dir, "hs_y.json"),
            hyper_decoder: w_of(&dir, "hyp_y.json"),
            synthesis: w_of(&dir, "synth_y.json"),
        },
        chroma: PlaneWeights {
            hyper_scale: int_of(&dir, "hs_uv.json"),
            hyper_decoder: w_of(&dir, "hyp_uv.json"),
            synthesis: w_of(&dir, "synth_uv.json"),
        },
        mcm_y: [w_of(&dir, "mcm0_y.json"), w_of(&dir, "mcm1_y.json"),
                w_of(&dir, "mcm2_y.json"), w_of(&dir, "mcm3_y.json")],
        icci,
    };
    let norm = NormTables {
        pmf: nv["pmf"].as_array().unwrap().iter()
            .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
        bounds: i64v(&nv, "bounds"),
        per_channel: nv["per_channel"].as_u64().unwrap() as usize,
        cdfs_y: i64v(&nv, "cdfs_y"), cdfs_uv: i64v(&nv, "cdfs_uv"),
        gain_y: i64v(&nv, "gain_y"), gain_uv: i64v(&nv, "gain_uv"),
        lsbs_t0: i32v(&nv, "lsbs_t0"), lsbs_t1: i32v(&nv, "lsbs_t1"),
        lef_thr: i3(&nv, "lef_thr0"), lef_mag: f3(&nv, "lef_mag0"),
    };

    let img = decode(&data, &weights, &norm).expect("supported geometry");
    assert_eq!((img.h, img.w), (256, 256));

    let refrgb = u8v(&nv, "ref_rgb");
    assert_eq!(img.data.len(), refrgb.len(), "RGB size");
    let (mut maxd, mut off) = (0i32, 0usize);
    for (g, r) in img.data.iter().zip(&refrgb) {
        let d = (*g as i32 - *r as i32).abs();
        maxd = maxd.max(d);
        if d != 0 { off += 1; }
    }
    let total = refrgb.len();
    eprintln!("decode() → RGB vs reference PNG: max|Δ| = {maxd}, {off}/{total} samples differ ({:.3}%)",
              100.0 * off as f64 / total as f64);
    // Off-by-one only (the integer/PTQ floor); differing fraction near the per-plane
    // synthesis PTQ floor (luma ~1.4% + chroma ~3.4% → ~4.4% combined). NOT bit-equality.
    assert!(maxd <= 2, "RGB deviation beyond the PTQ floor: max|Δ| = {maxd}");
    assert!(off * 100 <= total * 5, "more than 5% of samples differ: {off}/{total}");
}

fn weights_from(get: &dyn Fn(&str) -> serde_json::Value) -> DecoderWeights {
    let w = |n: &str| Weights::from_json(&get(n)).unwrap();
    let i = |n: &str| IntNet::from_json(&get(n)).unwrap();
    let iw = get("icci.json");
    DecoderWeights {
        luma: PlaneWeights { hyper_scale: i("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
        chroma: PlaneWeights { hyper_scale: i("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
        mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
        icci: IcciNet {
            luma: path_of(&iw["luma"]), chroma: path_of(&iw["chroma"]),
            conv_last_v: conv_of(&iw["conv_last_v"]),
            scale_v: iw["scale_v"].as_f64().unwrap() as f32,
        },
    }
}

fn norm_from(nv: &serde_json::Value) -> NormTables {
    NormTables {
        pmf: nv["pmf"].as_array().unwrap().iter()
            .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
        bounds: i64v(nv, "bounds"),
        per_channel: nv["per_channel"].as_u64().unwrap() as usize,
        cdfs_y: i64v(nv, "cdfs_y"), cdfs_uv: i64v(nv, "cdfs_uv"),
        gain_y: i64v(nv, "gain_y"), gain_uv: i64v(nv, "gain_uv"),
        lsbs_t0: i32v(nv, "lsbs_t0"), lsbs_t1: i32v(nv, "lsbs_t1"),
        lef_thr: i3(nv, "lef_thr0"), lef_mag: f3(nv, "lef_mag0"),
    }
}

/// The JPEG AI exact decode's parameter set, pinned through its listing (SPEC §3.1,
/// `wai.prior.set`). The fetched files are laid out as the reference sink reads them
/// — `norm.json` at the root, the weight files under `weights/` — and the listing
/// of that tree is the text whose length and SHA-256 are pinned below. The pin
/// verifies, the decode runs from the very bytes it verified, and its pixels are the
/// pinned ones when `WAI_JPEGAI_EXPECT_SHA256` is set; one flipped byte in a weight
/// file is `part_mismatch`.
#[test]
fn jpegai_e2e_set_pin_verifies() {
    use wai::codecs::{DeclaredTier, CAP_PRIOR_SET};
    use wai::container::Prior;
    use wai::pin::{listing_of_dir, sha256_hex, verify, Held, READABLE_FORMS};
    const LISTING_SHA256: &str = "3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37";
    let Some([cs, dir, nf]) = require_vectors::env_paths(
        "jpegai_e2e_full::jpegai_e2e_set_pin_verifies",
        ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"],
        "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)",
    ) else {
        return;
    };
    let tree = std::env::temp_dir().join(format!("wai-jpegai-set-{}", std::process::id()));
    std::fs::create_dir_all(tree.join("weights")).unwrap();
    std::fs::copy(&nf, tree.join("norm.json")).unwrap();
    const WEIGHTS: [&str; 11] = ["hs_y.json", "hs_uv.json", "hyp_y.json", "hyp_uv.json", "synth_y.json",
                                 "synth_uv.json", "mcm0_y.json", "mcm1_y.json", "mcm2_y.json", "mcm3_y.json",
                                 "icci.json"];
    for n in WEIGHTS {
        std::fs::copy(std::path::Path::new(&dir).join(n), tree.join("weights").join(n)).unwrap();
    }
    let listing = listing_of_dir(&tree).unwrap();
    assert_eq!(listing.len(), 1019);
    assert_eq!(sha256_hex(&listing), LISTING_SHA256, "the listing of the pinned jpegai/v1 exact bundle");
    std::fs::write(tree.join("jpegai-v1.sha256"), &listing).unwrap();

    let pin = Prior::new("jpegai/v1-exact", LISTING_SHA256, "int-entropy+fixed-order-f32")
        .with_form(CAP_PRIOR_SET)
        .with_tier(DeclaredTier::EntropyConsistency);
    let held = [Held::Listing(tree.join("jpegai-v1.sha256"))];
    let v = verify("wai.image.jpegai", &pin, &held, READABLE_FORMS).expect("the exact bundle verifies");
    let get = |n: &str| -> serde_json::Value { serde_json::from_slice(v.file(&format!("weights/{n}")).unwrap()).unwrap() };
    let weights = weights_from(&get);
    let norm = norm_from(&serde_json::from_slice(v.file("norm.json").unwrap()).unwrap());
    let data = std::fs::read(&cs).unwrap();
    let img = decode(&data, &weights, &norm).expect("supported geometry");
    assert_eq!((img.h, img.w), (256, 256));
    match std::env::var("WAI_JPEGAI_EXPECT_SHA256") {
        Ok(want) if !want.is_empty() => assert_eq!(sha256_hex(&img.data), want, "the pinned pixels"),
        _ => eprintln!("WAI_JPEGAI_EXPECT_SHA256 unset: pixels hash to {}", sha256_hex(&img.data)),
    }

    let p = tree.join("weights/hs_y.json");
    let mut b = std::fs::read(&p).unwrap();
    let at = b.len() / 2;
    b[at] ^= 1;
    std::fs::write(&p, &b).unwrap();
    assert_eq!(verify("wai.image.jpegai", &pin, &held, READABLE_FORMS).unwrap_err().code(), "part_mismatch");
    std::fs::remove_dir_all(&tree).ok();
}