open-standards/wai/wai-rs/tests/jpegai_e2e.rs
Transaction Science 48becebcab wai-rs: the jpegai tests on ort rc.15, in the process's one environment
8a7492b23 moved ort to 2.0.0-rc.15 and gave ort_session.rs the environment
rc.15 asks every session builder for; five test files build sessions of their
own (jpegai_e2e and the four calibration tests) and still called
`Session::builder()` with none. They build only with `neural` (and
`legacy_int`), which no push-triggered job builds, and a checked list cut at
ten lines hid them.

rc.15 allows one environment per process, and these tests call into wai in the
same process, so they take wai's: `ort_session::environment` is public as
`wai::ort_environment` (a new path, compatible). All five build and run under
neural,semantic,legacy_int; without their JPEG AI fixtures they return early,
as before. jpegai_decode_path builds and passes under its features. clippy
reports nothing in the files this touches.

Changelog: none

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-11-12 00:15:14 -05:00

403 lines
18 KiB
Rust
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//! End-to-end JPEG AI decode in Rust: real bitstream → RGB pixels, no Python.
//!
//! Every stage was validated in isolation first (see jpegai_mans_golden.rs
//! and jpegai-vm-assets/README.md); this test chains them: codestream parse →
//! the entropy decodes and the integer σ-net through the integer operator set
//! (`jpegai_ops`, integer-operators §13.6.1) → dequant → MCM value path (ORT stage nets) →
//! LSBS → synthesis (ORT, 16×16-traced) → EFElinear → eICCI (ORT wrappers) →
//! EFEnonlinear → LEF → BT.709, compared against the reference decoder's PNG.
//!
//! The signalled parameters now come from a real PIH/TON header decode:
//! `parse_headers` reads the bitstream's own headers and supplies the beta
//! displacement, the EFElinear coefficients + means, the eICCI choices, the
//! EFEnonlinear weights + luma masks, the LEF model, and the per-region model
//! flags — all driven into the decode below, and independently checked against
//! the reference header trace by `jpegai_header_parser_matches_trace`. What
//! remains fixtured is sink-supplied model state (the ORT nets, the integer
//! σ-net weights, the normative tables, the per-checkpoint constants) plus a
//! few intermediate reference dumps kept as scaffolding for stage cross-checks
//! (the per-position dequant gain map, the z factorized-CDF offsets, the LSBS
//! thresholds); the header-signalled parameters are no longer among them.
//!
//! Needs both `neural` (the ORT stage/synthesis nets + the jpegai modules) and
//! `codecs` (the `image` crate, to load the reference PNG it compares against).
#![cfg(all(feature = "neural", feature = "codecs"))]
use std::path::{Path, PathBuf};
use ort::session::{builder::GraphOptimizationLevel, Session};
use ort::value::Value;
use wai::neural::jpegai_bitstream::{parse_ae_substream, parse_codestream};
use wai::neural::jpegai_header::parse_headers;
use wai::neural::jpegai_filters::{efe_linear_plane, efe_nonlinear_plane, lef_apply, EfeLinearPlane};
fn assets() -> PathBuf {
PathBuf::from("../jpegai-vm-assets")
}
fn load_json(p: &Path) -> serde_json::Value {
serde_json::from_slice(&std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display())))
.unwrap()
}
fn i64s(v: &serde_json::Value) -> Vec<i64> {
v.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
}
fn f32s(v: &serde_json::Value) -> Vec<f32> {
v.as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
}
fn sess(p: &Path) -> Session {
Session::builder(wai::ort_environment().unwrap())
.unwrap()
.with_intra_threads(1)
.unwrap()
.with_inter_threads(1)
.unwrap()
.with_optimization_level(GraphOptimizationLevel::Level3)
.unwrap()
.commit_from_file(p)
.unwrap_or_else(|e| panic!("load {}: {e}", p.display()))
}
/// Run a single-input ONNX model on a [1,C,H,W] f32 tensor; returns (data, shape).
fn run1(s: &mut Session, x: Vec<f32>, c: usize, h: usize, w: usize) -> (Vec<f32>, Vec<usize>) {
let name = s.inputs().first().map(|i| i.name().to_string()).unwrap();
let arr = ndarray::Array4::from_shape_vec((1, c, h, w), x).unwrap();
let v = Value::from_array(arr).unwrap();
let out = s.run(ort::inputs![name.as_str() => v]).unwrap();
let view = out[0].try_extract_array::<f32>().unwrap();
let shape = view.shape().to_vec();
(view.iter().copied().collect(), shape)
}
/// Run a two-input MCM stage net: (p, prev) — both [1,C,8,8].
fn run_mcm(s: &mut Session, p: Vec<f32>, pc: usize, prev: Vec<f32>, vc: usize) -> Vec<f32> {
let names: Vec<String> = s.inputs().iter().map(|i| i.name().to_string()).collect();
let pa = ndarray::Array4::from_shape_vec((1, pc, 8, 8), p).unwrap();
let mut feeds: Vec<(std::borrow::Cow<'static, str>, ort::session::SessionInputValue<'static>)> =
vec![(names[0].clone().into(), Value::from_array(pa).unwrap().into())];
if names.len() > 1 {
let va = ndarray::Array4::from_shape_vec((1, vc, 8, 8), prev).unwrap();
feeds.push((names[1].clone().into(), Value::from_array(va).unwrap().into()));
}
let out = s.run(feeds).unwrap();
out[0].try_extract_array::<f32>().unwrap().iter().copied().collect()
}
/// Split a [C,16,16] plane into 4 stage slices [C,8,8] in decode order
/// (TL, BR, TR, BL — the VM's down_shuffle part order).
fn down_shuffle(x: &[f32], c: usize) -> [Vec<f32>; 4] {
let phase = |py: usize, px: usize| -> Vec<f32> {
let mut out = Vec::with_capacity(c * 64);
for ch in 0..c {
for y in 0..8 {
for xx in 0..8 {
out.push(x[ch * 256 + (2 * y + py) * 16 + 2 * xx + px]);
}
}
}
out
};
[phase(0, 0), phase(1, 1), phase(0, 1), phase(1, 0)]
}
/// Inverse: stage slices (decode order TL,BR,TR,BL) → [C,16,16].
fn up_shuffle(st: &[Vec<f32>; 4], c: usize) -> Vec<f32> {
let mut out = vec![0f32; c * 256];
for (si, (py, px)) in [(0usize, (0usize, 0usize)), (1, (1, 1)), (2, (0, 1)), (3, (1, 0))] {
for ch in 0..c {
for y in 0..8 {
for xx in 0..8 {
out[ch * 256 + (2 * y + py) * 16 + 2 * xx + px] = st[si][ch * 64 + y * 8 + xx];
}
}
}
}
out
}
/// LSBS scaling: σ_skip-driven table lookup (validated in round 6b).
fn lsbs(yhat: &[f32], resi_dq: &[f32], sigma_skip: &[i64], c: usize, t0: &[i64], t1: &[i64]) -> Vec<f32> {
// likely: per-channel 8×8 block sum of the SIGNED skip surface, (Σ+32)>>6
let mut out = vec![0f32; c * 256];
for ch in 0..c {
let mut blocks = [[0i64; 2]; 2];
for y in 0..16 {
for x in 0..16 {
blocks[y / 8][x / 8] += sigma_skip[ch * 256 + y * 16 + x];
}
}
for y in 0..16 {
for x in 0..16 {
let lk = ((blocks[y / 8][x / 8] + 32) >> 6).clamp(0, 4351) as usize;
let i = ch * 256 + y * 16 + x;
let mean_full = yhat[i] - resi_dq[i];
let additive = (t0[lk] as f32 * mean_full + t1[lk] as f32 * resi_dq[i]) / 8192.0;
out[i] = yhat[i] + additive;
}
}
}
out
}
/// Diagnostic: max abs diff vs a checkpoint dumped from the live reference.
fn ckdiff(cp: &serde_json::Value, key: &str, got: &[f32]) {
if cp.get(key).is_none() {
return;
}
let want = f32s(&cp[key]);
let n = want.len().min(got.len());
let mut worst = 0f32;
for i in 0..n {
worst = worst.max((got[i] - want[i]).abs());
}
eprintln!(" [ck] {key}: max|err| {worst:.4e} (got {} want {})", got.len(), want.len());
}
#[test]
fn jpegai_end_to_end_decode() {
let a = assets();
for need in ["e2e/test256.bin", "e2e/route_fixture.json", "model16/CCS_SGMM"] {
if !a.join(need).exists() {
eprintln!("skipping: {need} not present");
return;
}
}
let cp = load_json(&a.join("e2e/checkpoints_fixture.json"));
// ---- fixtures ----
let ent = load_json(&a.join("e2e/jpegai_real_entropy_fixture.json"));
let weights = load_json(&a.join("hyper_scale_decoder_int_weights.json"));
let adds = load_json(&a.join("e2e/sigma_additions_fixture.json"));
let gains = load_json(&a.join("e2e/dequant_gain_map_fixture.json"));
let route = load_json(&a.join("e2e/route_fixture.json"));
let norm = load_json(&a.join("jpegai_normative_tables.json"));
// ---- 1. codestream + headers ----
let bin = std::fs::read(a.join("e2e/test256.bin")).unwrap();
let subs = parse_codestream(&bin).unwrap();
let payload = |m: u16| subs.iter().find(|s| s.marker == m).unwrap().payload;
// ---- headers: typed parse (PIH + TON, separate cursors) ----
let hdr = parse_headers(payload(0xFF82), payload(0xFF83)).expect("headers");
assert_eq!((hdr.img_width, hdr.img_height), (256, 256));
let consts = load_json(&a.join("model_constants.json"));
// ---- 2–3. SOZ → z → the scale graph → σ, mask → the residuals, through the integer
// operator set as both decode paths run them (integer-operators §13.6.1), from this
// bundle's members: the normative PMFs and bounds, the factorized CDFs, the gains and
// the scale network. The derived quantities are checked against the dumped ones
// while both exist. ----
let fac = load_json(&a.join("jpegai_factorized_cdfs.json"));
let cdfs = |p: &str| -> Vec<i64> { fac[p].as_array().unwrap().iter().flat_map(i64s).collect() };
for plane in ["y", "uv"] {
assert!(i64s(&route["z_offsets"][plane]).iter().all(|&o| o == 31), "the hyper-latent offsets are §13.6.1's");
}
let bounds: Vec<i64> = std::fs::read_to_string(a.join("me-tANS/bounds.csv")).unwrap()
.lines().filter(|l| !l.trim().is_empty()).map(|l| l.trim().parse().unwrap()).collect();
let gain_of = |p: &str| -> Vec<i64> { i64s(&consts["tools_0"][format!("gain_mlog_{p}")]).iter().map(|&g| g + hdr.beta_displacement_log).collect() };
let net = wai_ops::lower::jpegai::JpegaiNet {
norm: wai_ops::lower::jpegai::Norm {
pmf: norm["pmf"].as_array().unwrap().iter().map(i64s).collect(),
bounds,
per_channel: fac["per_channel"].as_i64().unwrap(),
cdfs_y: cdfs("y"),
cdfs_uv: cdfs("uv"),
gain_y: gain_of("y"),
gain_uv: gain_of("uv"),
},
hs_y: wai::jpegai_ops::scale_net_of_blocks(&weights["y"]).unwrap(),
hs_uv: wai::jpegai_ops::scale_net_of_blocks(&weights["uv"]).unwrap(),
};
let sp = wai::jpegai_ops::load_net(net).expect("the bundle's scale path loads under acc32");
let (mem_z, ends_z) = parse_ae_substream(payload(0xFF88), 1).unwrap();
let mut sigma_final: Vec<Vec<i64>> = Vec::new(); // per plane [C*256]
let mut sigma_skip: Vec<Vec<i64>> = Vec::new();
let mut z_net_planes: Vec<Vec<i64>> = Vec::new();
let mut resi: Vec<Vec<i16>> = Vec::new();
for (pi, (plane, marker, stream, p)) in [("y", 0xFF89u16, "SORP", wai::jpegai_ops::Plane::Luma), ("uv", 0xFF8A, "SORS", wai::jpegai_ops::Plane::Chroma)].into_iter().enumerate() {
let (mem, ends) = parse_ae_substream(payload(marker), 1).unwrap();
let grfs: Vec<bool> = hdr.models[pi].grfs_channel_flag.iter().map(|&f| f == 1).collect();
let ints = wai::jpegai_ops::plane(&sp, p, &mem_z[..ends_z[0] as usize], &mem[..ends[0] as usize], &grfs).expect("the plane decodes");
let fin: Vec<i64> = ints.v.iter().zip(&ints.rvs).map(|(v, r)| v + r).collect();
let rvs_fix = i64s(&adds[plane]["rvs"]);
for (i, &f) in fin.iter().enumerate() {
assert_eq!(f, ints.v[i] + rvs_fix[i], "derived rvs mismatch at {i}");
}
let calls = ent[stream]["calls"].as_array().unwrap();
for (c, call) in calls.iter().enumerate() {
let fx: Vec<bool> = i64s(&call["masks"]).iter().map(|&v| v != 0).collect();
assert_eq!(&ints.mask[c * 256..(c + 1) * 256], &fx[..], "derived skip mask != fixture at {stream} ch {c}");
let want: Vec<i64> = i64s(&call["x"]).iter().map(|&v| v as i16 as i64).collect();
assert_eq!(&ints.y_q[c * 256..(c + 1) * 256], &want[..], "{stream} ch {c}: the residuals");
}
resi.push(ints.y_q.iter().map(|&q| q as i16).collect());
sigma_final.push(fin);
z_net_planes.push(ints.z);
sigma_skip.push(ints.v);
}
// ---- 4. dequant (per-position gain map) ----
let mut resi_dq: Vec<Vec<f32>> = Vec::new();
for (pi, plane) in ["y", "uv"].iter().enumerate() {
let g = f32s(&gains[plane]["gain"]);
resi_dq.push(resi[pi].iter().zip(&g).map(|(&r, &gv)| r as f32 * gv).collect());
}
// ---- 5. hyper decoder (ORT) → ψ ----
let mut psi: Vec<Vec<f32>> = Vec::new();
for (pi, (plane, c)) in [("y", 160usize), ("uv", 96)].iter().enumerate() {
let mut s = sess(&a.join(format!("model/CCS_SGMM/tools_0/model_{plane}/common_modules/hyper_decoder.onnx")));
let zf: Vec<f32> = z_net_planes[pi].iter().map(|&v| v as f32).collect();
let (out, _) = run1(&mut s, zf, *c, 4, 4);
psi.push(out);
}
// ---- 6. value path: Y via MCM stages, UV via ψ up-shuffle ----
let resi_y_st = down_shuffle(&resi_dq[0], 160);
let mut yhat_st: Vec<Vec<f32>> = Vec::new();
for s_i in 0..4 {
let mut st = sess(&a.join(format!("model/CCS_SGMM/tools_0/model_y/common_modules/MCM/stage{s_i}.onnx")));
let p: Vec<f32> = psi[0][s_i * 160 * 64..(s_i + 1) * 160 * 64].to_vec();
let prev: Vec<f32> = yhat_st.iter().flatten().copied().collect();
let prev_c = 160 * yhat_st.len().max(1);
let prev = if yhat_st.is_empty() { vec![0f32; 160 * 64] } else { prev };
let mu = run_mcm(&mut st, p, 160, prev, prev_c);
yhat_st.push(resi_y_st[s_i].iter().zip(&mu).map(|(a, b)| a + b).collect());
}
let yhat_y = up_shuffle(&[yhat_st[0].clone(), yhat_st[1].clone(), yhat_st[2].clone(), yhat_st[3].clone()], 160);
ckdiff(&cp, "resi_dq_y", &resi_dq[0]);
ckdiff(&cp, "psi_y", &psi[0]);
ckdiff(&cp, "yhat_y", &yhat_y);
// UV: μ = up_shuffle of ψ chunks (decode order = chunks 0,2,3,1 → TL,BR,TR,BL)
let ch = |i: usize| psi[1][i * 96 * 64..(i + 1) * 96 * 64].to_vec();
let mu_uv = up_shuffle(&[ch(0), ch(1), ch(2), ch(3)], 96);
let yhat_uv: Vec<f32> = resi_dq[1].iter().zip(&mu_uv).map(|(a, b)| a + b).collect();
ckdiff(&cp, "yhat_uv", &yhat_uv);
// ---- 7. LSBS ----
let t0 = i64s(&route["lsbs_t0"]);
let t1 = i64s(&route["lsbs_t1"]);
let y_ls = lsbs(&yhat_y, &resi_dq[0], &sigma_skip[0], 160, &t0, &t1);
let uv_ls = lsbs(&yhat_uv, &resi_dq[1], &sigma_skip[1], 96, &t0, &t1);
ckdiff(&cp, "y_lsbs", &y_ls);
ckdiff(&cp, "uv_lsbs", &uv_ls);
// ---- 8. synthesis (ORT, 16×16-traced) ----
let mut sy = sess(&a.join("model16/CCS_SGMM/tools_0/model_y/synthesis16.onnx"));
let (ypix, _) = run1(&mut sy, y_ls.clone(), 160, 16, 16); // [1,1,256,256]
let mut suv = sess(&a.join("model16/CCS_SGMM/tools_0/model_uv/synthesis16.onnx"));
let mut uv_in = y_ls;
uv_in.extend_from_slice(&uv_ls);
let (uvpix, _) = run1(&mut suv, uv_in, 256, 16, 16); // [1,2,256,256]
let (h, w) = (256usize, 256usize);
let mut yp = ypix;
let mut up: Vec<f32> = uvpix[0..h * w].to_vec();
let mut vp: Vec<f32> = uvpix[h * w..2 * h * w].to_vec();
ckdiff(&cp, "ypix", &yp);
ckdiff(&cp, "uvpix", &uvpix);
// ---- 9. EFElinear (parameters straight from the tool header) ----
let deint = |q: u64| (q as f32 - 32767.0) / 2048.0;
assert!(hdr.efe_linear_enabled);
for (ci, wq, w2q, mean_q) in [
(1usize, hdr.efe_linear.weights_u[0], hdr.efe_linear.weights_u2[0], hdr.efe_linear.mean_u_q100),
(2, hdr.efe_linear.weights_v[0], hdr.efe_linear.weights_v2[0], hdr.efe_linear.mean_v_q100),
] {
let mean = mean_q as f32 / 100.0;
let p = EfeLinearPlane {
mean,
w_uv: [deint(wq) + 1.0; 4],
bias: [mean; 4],
w_y: [deint(w2q); 4],
};
let src = if ci == 1 { &up } else { &vp };
let out = efe_linear_plane(&yp, src, h, w, &p);
if ci == 1 { up = out } else { vp = out }
}
{
let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
ckdiff(&cp, "efl_out", &all);
}
// ---- 10. eICCI (ORT wrappers; [0,1] domain) ----
let mut yuv01 = Vec::with_capacity(3 * h * w);
for pl in [&yp, &up, &vp] {
yuv01.extend(pl.iter().map(|&v| v / 255.0));
}
let crop = |data: &[f32], shape: &[usize], ch: usize| -> Vec<f32> {
let (oh, ow) = (shape[2], shape[3]);
let mut out = Vec::with_capacity(h * w);
for y in 0..h {
for x in 0..w {
out.push(data[(ch * oh + y) * ow + x] * 255.0);
}
}
out
};
let mut si_y = sess(&a.join("e2e/pfchain/icci_y_m6.onnx"));
let (fy, fysh) = run1(&mut si_y, yuv01.clone(), 3, h, w);
let mut si_uv = sess(&a.join("e2e/pfchain/icci_uv_m3.onnx"));
let (fuv, fuvsh) = run1(&mut si_uv, yuv01, 3, h, w);
yp = crop(&fy, &fysh, 0);
up = crop(&fuv, &fuvsh, 0);
vp = crop(&fuv, &fuvsh, 1);
{
let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
ckdiff(&cp, "icci_out", &all);
}
// ---- 11. EFEnonlinear (parameters from the tool header) ----
assert!(hdr.efe_nonlinear_enabled && hdr.efn_v_enabled && !hdr.efn_u_enabled);
let wq: Vec<i64> = hdr.efn_weights.iter().map(|&v| v as i64).collect();
vp = efe_nonlinear_plane(&yp, &vp,
&[wq[0], wq[1], wq[2], wq[3], wq[4], wq[5], wq[6], wq[7]],
hdr.efn_min_luma_q100 as i64, hdr.efn_max_luma_q100 as i64);
{
let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
ckdiff(&cp, "efn_out", &all);
}
// ---- 12. LEF (op-point tables at header model_id; σ channel from header) ----
assert!(hdr.lef_enabled);
let mid = hdr.model_id as usize;
let thr_t = i64s(&consts["lef_thr"][mid]);
let mag_t = f32s(&consts["lef_mag"][mid]);
let c0 = hdr.lef_ch_idx as usize * 256;
yp = lef_apply(&yp, &sigma_final[0][c0..c0 + 256], h, w,
[thr_t[0], thr_t[1], thr_t[2]], [mag_t[0], mag_t[1], mag_t[2]]);
{
let mut all = yp.clone(); all.extend_from_slice(&up); all.extend_from_slice(&vp);
ckdiff(&cp, "lef_out", &all);
}
// ---- 13. BT.709 → RGB → compare with reference PNG ----
let refimg = image::open(a.join("e2e/test256_dec.png")).unwrap().to_rgb8();
let (kr, kg, kb): (f32, f32, f32) = (0.2126, 0.7152, 0.0722);
let (kby, kry): (f32, f32) = (1.8556, 1.5748);
let mut mismatched = 0usize;
let mut maxdiff = 0i32;
for i in 0..h * w {
let y01 = yp[i] / 255.0;
let u01 = up[i] / 255.0 - 0.5;
let v01 = vp[i] / 255.0 - 0.5;
let rgb = [
y01 + kry * v01,
y01 - (kb * kby / kg) * u01 - (kr * kry / kg) * v01,
y01 + kby * u01,
];
let px = refimg.get_pixel((i % w) as u32, (i / w) as u32);
for (c, &val) in rgb.iter().enumerate() {
let got = (val * 255.0).round_ties_even().clamp(0.0, 255.0) as i32;
let d = (got - px[c] as i32).abs();
if d > 0 { mismatched += 1; }
maxdiff = maxdiff.max(d);
}
}
let total = 3 * h * w;
println!(" e2e decode: {mismatched}/{total} samples differ, max diff {maxdiff}");
assert!(maxdiff <= 1, "pixel deviation beyond rounding: {maxdiff}");
assert!(mismatched * 1000 <= total, "more than 0.1% mismatched: {mismatched}/{total}");
}