8386f0f45a
The duplicate ports neural/jpegai_hyper.rs and neural/jpegai_mans.rs, which the float decode path ran its front half with, are deleted (O10): that path runs the shared driver since the switch. The direct pieces of the exact decoder leave production, test-only (legacy_int): the crate-root mANS port (jpegai_mans), the hyper-scale network (jpegai_synthesis::decode_hyper_scale and its convolution) and the pipeline's skip_mask, rvs_additive, rvs_likely and sigma_index. The calibration tests that read them require legacy_int. jpegai_mans_golden runs on E8 and the scale graph alone; jpegai_e2e (the float chain against the reference PNG) takes its front half from the driver, still checking the dumped RVS additions, masks and residuals; the byte-equality crate drops the neural port and keeps E8 and the direct port. Co-Authored-By: Claude Opus 5.5 <[email protected]>
499 lines
25 KiB
Rust
499 lines
25 KiB
Rust
//! Golden validation of the JPEG AI entropy decode against bitstreams produced by
|
||
//! the REFERENCE software's own encoder (ISO/IEC 6048-1 reference, BSD-3,
|
||
//! compiled from gitlab.com/wg1/jpeg-ai/jpeg-ai-reference-software), through the decode
|
||
//! the exact decoder and the float decode path both run: the integer operator set's
|
||
//! `mans_decode` (integer-operators §5E E8) and the scale graph (§13.6.1). The fixture
|
||
//! also carries the reference-built decode tables, so E8's table builder is
|
||
//! byte-compared too. Skipped if the fixture is absent (regenerate with the harness
|
||
//! documented in tools/wai_jpegai_goldens.py).
|
||
#![cfg(feature = "neural")]
|
||
|
||
use std::path::Path;
|
||
|
||
use wai_ops::stage::mans;
|
||
|
||
#[test]
|
||
fn jpegai_mans_matches_reference_software() {
|
||
let fixture = Path::new("../jpegai-vm-assets/jpegai_mans_goldens.json");
|
||
if !fixture.exists() {
|
||
eprintln!("skipping jpegai mans golden: fixture not present");
|
||
return;
|
||
}
|
||
let v: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(fixture).unwrap()).unwrap();
|
||
let as_i64v = |x: &serde_json::Value| -> Vec<i64> {
|
||
x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
|
||
};
|
||
|
||
// 1. E8's table builder reproduces the reference tables bit for bit.
|
||
let pmf: Vec<Vec<i64>> = v["pmf"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
let reference: Vec<Vec<i64>> =
|
||
v["decode_transitions"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
let built = mans::build_latent_tables(&pmf).expect("every row sums to 256");
|
||
for (d, (b, r)) in built.iter().zip(&reference).enumerate() {
|
||
for s in 0..256 {
|
||
assert_eq!(b[s] as i64, r[s], "E8 table mismatch: distribution {d}, state {s}");
|
||
}
|
||
}
|
||
println!(" wai-ops E8 table builder: {} distributions x 256 states bit-exact", built.len());
|
||
|
||
// 2. The latent path (`sgm`) against the reference encoder's bitstream.
|
||
let y = &v["y"];
|
||
let mask: Vec<bool> = y["masks"].as_array().unwrap().iter().map(|m| m.as_bool().unwrap()).collect();
|
||
let sigma: Vec<usize> = as_i64v(&y["sigma"]).iter().map(|&s| s as usize).collect();
|
||
let mem: Vec<u8> = as_i64v(&y["memory"]).iter().map(|&b| b as u8).collect();
|
||
let end = as_i64v(&y["offsets"])[0] as usize;
|
||
let got = mans::decode_sgm(&mem[..end], &pmf, &as_i64v(&v["bounds"]), &[mans::Span { mask: mask.clone(), sigma }]).unwrap();
|
||
let coded = mask.iter().filter(|&&m| m).count();
|
||
for (i, &e) in as_i64v(&y["expected"]).iter().enumerate() {
|
||
if mask[i] {
|
||
assert_eq!(got[0][i], e as i16 as i64, "E8 y symbol {i} diverges");
|
||
}
|
||
}
|
||
println!(" wai-ops E8 sgm: {coded} residuals (incl. escapes) match the reference encoder");
|
||
|
||
// 3. The hyper-latent path (`fac`).
|
||
let z = &v["z"];
|
||
let mem: Vec<u8> = as_i64v(&z["memory"]).iter().map(|&b| b as u8).collect();
|
||
let end = as_i64v(&z["offsets"])[0] as usize;
|
||
let rows: Vec<Vec<i64>> = z["cdfs"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
let want: Vec<Vec<i64>> = z["expected"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
let lens: Vec<usize> = want.iter().map(Vec::len).collect();
|
||
assert_eq!(mans::decode_fac(&mem[..end], &rows, &lens).unwrap(), want, "E8 z (fac) diverges from the reference encoder");
|
||
println!(" wai-ops E8 fac: {} rows x {} hyper-latents match the reference encoder", want.len(), lens[0]);
|
||
}
|
||
|
||
/// End-to-end through all three layers: codestream framing → AE-substream unpacking →
|
||
/// E8's `fac` decode of the reference encoder's golden z-stream.
|
||
#[test]
|
||
fn jpegai_codestream_to_symbols() {
|
||
use wai::neural::jpegai_bitstream::{parse_ae_substream, parse_codestream, MARKER_SOZ};
|
||
|
||
let fixture = std::path::Path::new("../jpegai-vm-assets/jpegai_mans_goldens.json");
|
||
if !fixture.exists() {
|
||
eprintln!("skipping: fixture not present");
|
||
return;
|
||
}
|
||
let v: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(fixture).unwrap()).unwrap();
|
||
let as_i64v = |x: &serde_json::Value| -> Vec<i64> {
|
||
x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
|
||
};
|
||
let z = &v["z"];
|
||
let memory: Vec<u8> = as_i64v(&z["memory"]).iter().map(|&b| b as u8).collect();
|
||
let rows: Vec<Vec<i64>> = z["cdfs"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
let expected: Vec<Vec<i64>> = z["expected"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
|
||
// Frame: SOC | SOZ + ue(size) + payload (single thread: no deltas) | EOC.
|
||
let mut cs: Vec<u8> = 0xFF80u16.to_be_bytes().to_vec();
|
||
cs.extend(MARKER_SOZ.to_be_bytes());
|
||
let n = memory.len() as u64 + 1;
|
||
let len = 64 - n.leading_zeros();
|
||
let total_bits = (2 * len - 1) as usize;
|
||
let padded = total_bits.div_ceil(8) * 8;
|
||
let acc: u128 = (n as u128) << (padded - total_bits);
|
||
for shift in (0..padded).step_by(8).rev() {
|
||
cs.push(((acc >> shift) & 0xFF) as u8);
|
||
}
|
||
cs.extend_from_slice(&memory);
|
||
cs.extend(0xFF81u16.to_be_bytes());
|
||
|
||
let subs = parse_codestream(&cs).expect("codestream parses");
|
||
assert_eq!(subs.len(), 1);
|
||
assert_eq!(subs[0].marker, MARKER_SOZ);
|
||
let (mem, offsets) = parse_ae_substream(subs[0].payload, 1).expect("ae unpacks");
|
||
assert_eq!(offsets.len(), 1);
|
||
let lens: Vec<usize> = expected.iter().map(Vec::len).collect();
|
||
let got = mans::decode_fac(&mem[..offsets[0] as usize], &rows, &lens).unwrap();
|
||
assert_eq!(got, expected, "end-to-end codestream → symbols diverged");
|
||
println!(" codestream → AE unpack → wai-ops E8: {} symbols exact", lens.iter().sum::<usize>());
|
||
}
|
||
|
||
/// The NORMATIVE table set: 32 σ-distributions from the standard's literal
|
||
/// PMF constants (`pdf_r` in the reference `ec_lib_mans.py`), exported via
|
||
/// the reference software's own `store_parameters`. E8's table builder must
|
||
/// reproduce every packed word the runtime decoder consumes.
|
||
#[test]
|
||
fn jpegai_normative_tables_bit_exact() {
|
||
let fixture = std::path::Path::new("../jpegai-vm-assets/jpegai_normative_tables.json");
|
||
if !fixture.exists() {
|
||
eprintln!("skipping: normative fixture not present");
|
||
return;
|
||
}
|
||
let v: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(fixture).unwrap()).unwrap();
|
||
let as_i64v = |x: &serde_json::Value| -> Vec<i64> {
|
||
x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
|
||
};
|
||
let pmf: Vec<Vec<i64>> = v["pmf"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
let packed: Vec<Vec<i64>> = v["packed"].as_array().unwrap().iter().map(&as_i64v).collect();
|
||
assert_eq!(pmf.len(), 32);
|
||
let built = mans::build_latent_tables(&pmf).expect("the normative rows sum to 256");
|
||
for (d, (b, r)) in built.iter().zip(&packed).enumerate() {
|
||
for s in 0..256 {
|
||
assert_eq!(b[s] as i64, r[s], "E8 normative table mismatch: distribution {d}, state {s}");
|
||
}
|
||
}
|
||
println!(" wai-ops E8: the normative tables, 32 distributions x 256 states bit-exact");
|
||
}
|
||
|
||
/// Parse a REAL bitstream produced by the reference encoder (256x256 RGB,
|
||
/// default operating point) — first contact with non-synthetic data.
|
||
#[test]
|
||
fn jpegai_real_bitstream_parses() {
|
||
let path = std::path::Path::new("../jpegai-vm-assets/e2e/test256.bin");
|
||
if !path.exists() {
|
||
eprintln!("skipping: e2e fixture not present");
|
||
return;
|
||
}
|
||
let buf = std::fs::read(path).unwrap();
|
||
let subs = wai::neural::jpegai_bitstream::parse_codestream(&buf)
|
||
.expect("real bitstream must parse");
|
||
let names = |m: u16| match m {
|
||
0xFF82 => "PIH", 0xFF83 => "TON", 0xFF84 => "RDI", 0xFF88 => "SOZ",
|
||
0xFF89 => "SORP", 0xFF8A => "SORS", 0xFF8B => "SOQ", 0xFF8C => "UDI",
|
||
_ => "?",
|
||
};
|
||
let layout: Vec<String> = subs.iter()
|
||
.map(|s| format!("{}:{}", names(s.marker), s.payload.len()))
|
||
.collect();
|
||
println!(" real bitstream substreams: [{}]", layout.join(", "));
|
||
assert!(subs.iter().any(|s| s.marker == 0xFF82), "must contain PIH");
|
||
assert!(subs.iter().any(|s| matches!(s.marker, 0xFF88..=0xFF8B)), "must contain an AE substream");
|
||
// every AE substream must unpack through the aemem layer
|
||
for s in subs.iter().filter(|s| matches!(s.marker, 0xFF88..=0xFF8B)) {
|
||
let unpacked = wai::neural::jpegai_bitstream::parse_ae_substream(s.payload, 1);
|
||
assert!(unpacked.is_some(), "AE substream {} must unpack", names(s.marker));
|
||
}
|
||
}
|
||
|
||
/// Replay every entropy-decode call the reference decoder made on the REAL
|
||
/// test256.bin bitstream through E8, with the normative tables: one `sgm` decode per
|
||
/// latent substream over all its channels' spans, and one `fac` decode over the
|
||
/// hyper-latent substream's rows, luma's then chroma's. Every one of the 258 calls'
|
||
/// symbols must be the reference's.
|
||
#[test]
|
||
fn jpegai_real_substreams_entropy_bit_exact() {
|
||
let assets = std::path::Path::new("../jpegai-vm-assets");
|
||
let fx_path = assets.join("e2e/jpegai_real_entropy_fixture.json");
|
||
let norm_path = assets.join("jpegai_normative_tables.json");
|
||
if !fx_path.exists() || !norm_path.exists() {
|
||
eprintln!("skipping: real-entropy fixtures not present");
|
||
return;
|
||
}
|
||
let norm: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(norm_path).unwrap()).unwrap();
|
||
let pmf: Vec<Vec<i64>> = norm["pmf"].as_array().unwrap().iter()
|
||
.map(|r| r.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect())
|
||
.collect();
|
||
let tables = mans::build_latent_tables(&pmf).unwrap();
|
||
let bounds: Vec<i64> = std::fs::read_to_string(assets.join("me-tANS/bounds.csv")).unwrap()
|
||
.lines().filter(|l| !l.trim().is_empty())
|
||
.map(|l| l.trim().parse().unwrap()).collect();
|
||
|
||
let bin = std::fs::read(assets.join("e2e/test256.bin")).unwrap();
|
||
let subs = wai::neural::jpegai_bitstream::parse_codestream(&bin).unwrap();
|
||
let payload = |m: u16| subs.iter().find(|s| s.marker == m).unwrap().payload;
|
||
let fx: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(&fx_path).unwrap()).unwrap();
|
||
let mut calls = 0usize;
|
||
for (name, marker) in [("SORP", 0xFF89u16), ("SORS", 0xFF8A), ("SOZ", 0xFF88)] {
|
||
let (mem, ends) = wai::neural::jpegai_bitstream::parse_ae_substream(payload(marker), 1).unwrap();
|
||
let st = &fx[name];
|
||
let want_mem: Vec<u8> = st["mem"].as_array().unwrap().iter()
|
||
.map(|n| n.as_u64().unwrap() as u8).collect();
|
||
assert_eq!(mem, &want_mem[..], "{name}: unpacked memory must match dumped");
|
||
let mem = &mem[..ends[0] as usize];
|
||
let list = st["calls"].as_array().unwrap();
|
||
let ints = |c: &serde_json::Value, k: &str| -> Vec<i64> { c[k].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect() };
|
||
if name == "SOZ" {
|
||
let mut rows = vec![];
|
||
let mut want = vec![];
|
||
for call in list {
|
||
let per = call["per_channel"].as_u64().unwrap() as usize;
|
||
rows.extend(ints(call, "pmfs").chunks(63).map(<[i64]>::to_vec));
|
||
want.extend(ints(call, "x").chunks(per).map(<[i64]>::to_vec));
|
||
}
|
||
let lens: Vec<usize> = want.iter().map(Vec::len).collect();
|
||
assert_eq!(mans::decode_fac(mem, &rows, &lens).unwrap(), want, "{name}: E8 fac");
|
||
} else {
|
||
let spans: Vec<mans::Span> = list.iter().map(|c| mans::Span {
|
||
mask: ints(c, "masks").iter().map(|&m| m != 0).collect(),
|
||
sigma: ints(c, "indexes").iter().map(|&s| s as usize).collect(),
|
||
}).collect();
|
||
let got = mans::decode_sgm_tables(mem, &tables, &bounds, &spans).unwrap();
|
||
for (i, (g, c)) in got.iter().zip(list).enumerate() {
|
||
let want: Vec<i64> = ints(c, "x").iter().map(|&v| v as i16 as i64).collect();
|
||
assert_eq!(g, &want, "{name}: E8 sgm span {i}");
|
||
}
|
||
}
|
||
calls += list.len();
|
||
}
|
||
assert_eq!(calls, 258);
|
||
println!(" real-bitstream entropy replay through wai-ops E8: {calls} calls bit-exact across SORP/SORS/SOZ");
|
||
}
|
||
|
||
/// The scale graph (integer-operators §13.6.1, lowered from the float path's bundle form
|
||
/// of the network): z (real, entropy-decoded) → the integer scale `I_s`, bit-exact
|
||
/// against the reference decoder's dumps for both planes. The graph is lowered with gains
|
||
/// of 0, so its output is `I_s` itself.
|
||
#[test]
|
||
fn jpegai_hyper_scale_decoder_int_bit_exact() {
|
||
let assets = std::path::Path::new("../jpegai-vm-assets");
|
||
let w_path = assets.join("hyper_scale_decoder_int_weights.json");
|
||
let fx_path = assets.join("e2e/hyper_scale_io_fixture.json");
|
||
if !w_path.exists() || !fx_path.exists() {
|
||
eprintln!("skipping: hyper-scale fixtures not present");
|
||
return;
|
||
}
|
||
let weights: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(&w_path).unwrap()).unwrap();
|
||
let fx: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(&fx_path).unwrap()).unwrap();
|
||
let ints = |x: &serde_json::Value| -> Vec<i64> {
|
||
x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
|
||
};
|
||
for plane in ["y", "uv"] {
|
||
let zs = &fx[plane];
|
||
let shape: Vec<usize> = zs["z_shape"].as_array().unwrap().iter()
|
||
.map(|n| n.as_u64().unwrap() as usize).collect();
|
||
let (c, h, w) = (shape[1], shape[2], shape[3]);
|
||
let want = ints(&zs["i_s"]);
|
||
let net = wai::jpegai_ops::scale_net_of_blocks(&weights[plane]).unwrap();
|
||
let g = wai::jpegai_ops::load_scale(net, c, &vec![0; c]).unwrap();
|
||
let v = wai::int_ops::run1(&g, wai_ops::Tensor::new(c as u32, h as u32, w as u32, ints(&zs["z"])), &[]).unwrap();
|
||
assert_eq!(v.data, want, "{plane}: the scale graph's I_s");
|
||
println!(" hyper-scale scale graph (wai-ops, acc32) [{plane}]: {} values bit-exact", v.data.len());
|
||
}
|
||
}
|
||
|
||
/// Close the sigma loop on real data through the integer operator set
|
||
/// (integer-operators §13.6.1 steps 2 and 3): z → the plane's scale graph with its gains
|
||
/// → `block_stat`, `band_lookup` with the GRFS flags the picture header signals,
|
||
/// `scale_index` and `coded`. The RVS additions must be the reference's, and every σ
|
||
/// index and mask the one every real `decode_sgm` call consumed.
|
||
#[test]
|
||
fn jpegai_sigma_chain_end_to_end() {
|
||
use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
|
||
use wai_ops::stage;
|
||
let assets = std::path::Path::new("../jpegai-vm-assets");
|
||
let paths = [
|
||
assets.join("hyper_scale_decoder_int_weights.json"),
|
||
assets.join("e2e/hyper_scale_io_fixture.json"),
|
||
assets.join("e2e/sigma_additions_fixture.json"),
|
||
assets.join("e2e/jpegai_real_entropy_fixture.json"),
|
||
];
|
||
if paths.iter().any(|p| !p.exists()) {
|
||
eprintln!("skipping: sigma-chain fixtures not present");
|
||
return;
|
||
}
|
||
let load = |p: &std::path::Path| -> serde_json::Value {
|
||
serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
|
||
};
|
||
let (weights, io, adds, ent) =
|
||
(load(&paths[0]), load(&paths[1]), load(&paths[2]), load(&paths[3]));
|
||
let ints = |x: &serde_json::Value| -> Vec<i64> {
|
||
x.as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
|
||
};
|
||
let bin = std::fs::read(assets.join("e2e/test256.bin")).unwrap();
|
||
let segs = parse_container(&bin);
|
||
let ph = parse_picture_header(&bin, find_segment(&segs, Marker::Pih).unwrap().start_bit());
|
||
let (coded_above, edges, additions) = wai::jpegai_ops::stage_constants();
|
||
let table: Vec<Vec<i64>> = additions.iter().map(|r| r.to_vec()).collect();
|
||
for (k, (plane, stream)) in [("y", "SORP"), ("uv", "SORS")].into_iter().enumerate() {
|
||
let shape: Vec<usize> = io[plane]["z_shape"].as_array().unwrap().iter()
|
||
.map(|n| n.as_u64().unwrap() as usize).collect();
|
||
let gain = ints(&adds[plane]["gain_per_ch"]);
|
||
let rvs = ints(&adds[plane]["rvs"]);
|
||
let c = gain.len();
|
||
let grfs = &ph.components[k].grfs_channel_flags;
|
||
let g = wai::jpegai_ops::load_scale(wai::jpegai_ops::scale_net_of_blocks(&weights[plane]).unwrap(), c, &gain).unwrap();
|
||
let v = wai::int_ops::run1(&g, wai_ops::Tensor::new(c as u32, shape[2] as u32, shape[3] as u32, ints(&io[plane]["z"])), &[]).unwrap();
|
||
let hw = v.data.len() / c;
|
||
let (b, sh, lo, hi) = wai::jpegai_ops::LIKELY;
|
||
let likely = stage::block_stat(&v, b, sh, lo, hi).unwrap().data;
|
||
let calls = ent[stream]["calls"].as_array().unwrap();
|
||
assert_eq!(calls.len(), c, "{plane}: call count vs channels");
|
||
for (i, (&x, &l)) in v.data.iter().zip(&likely).enumerate() {
|
||
let row = if grfs[i / hw] { 0 } else { 1 };
|
||
let r = stage::band_lookup(l, &edges, &table, row).unwrap();
|
||
assert_eq!(r, rvs[i], "{plane}: E7 rvs at {i}");
|
||
let call = &calls[i / hw];
|
||
assert_eq!(stage::scale_index(x + r, 7, 32).unwrap(), call["indexes"][i % hw].as_i64().unwrap(), "{plane}: E5 sigma index at {i}");
|
||
assert_eq!(stage::coded(x, coded_above), call["masks"][i % hw].as_i64().unwrap() != 0, "{plane}: E4 mask at {i}");
|
||
}
|
||
println!(" sigma chain through wai-ops [{plane}]: rvs, sigma indexes and masks of {c} channels reproduced from z");
|
||
}
|
||
}
|
||
|
||
/// EFElinear post-filter: Rust implementation vs the live reference filter
|
||
/// on the real decode (effective signalled coefficients dumped mid-decode).
|
||
#[test]
|
||
fn jpegai_efe_linear_matches_reference() {
|
||
use wai::neural::jpegai_filters::{efe_linear_plane, EfeLinearPlane};
|
||
let fx_path = std::path::Path::new("../jpegai-vm-assets/e2e/efl_fixture.json");
|
||
if !fx_path.exists() {
|
||
eprintln!("skipping: EFElinear fixture not present");
|
||
return;
|
||
}
|
||
let fx: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(fx_path).unwrap()).unwrap();
|
||
let f32s = |x: &serde_json::Value| -> Vec<f32> {
|
||
x.as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
|
||
};
|
||
let arr4 = |x: &serde_json::Value| -> [f32; 4] {
|
||
let v = f32s(x);
|
||
[v[0], v[1], v[2], v[3]]
|
||
};
|
||
let input = f32s(&fx["input"]);
|
||
let want = f32s(&fx["want"]);
|
||
let (h, w) = (256usize, 256usize);
|
||
let plane = |i: usize| &input[i * h * w..(i + 1) * h * w];
|
||
// Y must be untouched
|
||
assert_eq!(plane(0), &want[0..h * w], "EFElinear must not modify Y");
|
||
let mut worst = 0f32;
|
||
for (ci, pj) in [(1usize, 0usize), (2, 1)] {
|
||
let pv = &fx["planes"][pj];
|
||
let p = EfeLinearPlane {
|
||
mean: pv["mean"].as_f64().unwrap() as f32,
|
||
w_uv: arr4(&pv["w_uv"]),
|
||
bias: arr4(&pv["bias"]),
|
||
w_y: arr4(&pv["w_y"]),
|
||
};
|
||
let got = efe_linear_plane(plane(0), plane(ci), h, w, &p);
|
||
for (g, wv) in got.iter().zip(&want[ci * h * w..(ci + 1) * h * w]) {
|
||
worst = worst.max((g - wv).abs());
|
||
}
|
||
}
|
||
assert!(worst < 1e-3, "EFElinear worst error {worst}");
|
||
println!(" EFElinear: both chroma planes reproduced, worst |err| = {worst:.2e}");
|
||
}
|
||
|
||
/// LEF post-filter: Rust vs the live reference on the real decode.
|
||
#[test]
|
||
fn jpegai_lef_matches_reference() {
|
||
use wai::neural::jpegai_filters::lef_apply;
|
||
let fx_path = std::path::Path::new("../jpegai-vm-assets/e2e/lef_fixture.json");
|
||
if !fx_path.exists() {
|
||
eprintln!("skipping: LEF fixture not present");
|
||
return;
|
||
}
|
||
let fx: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(fx_path).unwrap()).unwrap();
|
||
let f32s = |k: &str| -> Vec<f32> {
|
||
fx[k].as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
|
||
};
|
||
let sigma: Vec<i64> = fx["sigma_16"].as_array().unwrap().iter()
|
||
.map(|n| n.as_i64().unwrap()).collect();
|
||
let got = lef_apply(&f32s("y_in"), &sigma, 256, 256, [600, 1200, 2000], [1.1, 1.13, 1.16]);
|
||
let want = f32s("y_want");
|
||
let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
|
||
assert!(worst < 1e-3, "LEF worst error {worst}");
|
||
println!(" LEF: luma sharpening reproduced, worst |err| = {worst:.2e}");
|
||
}
|
||
|
||
/// EFEnonlinear post-filter: Rust vs the live reference on the real decode.
|
||
#[test]
|
||
fn jpegai_efe_nonlinear_matches_reference() {
|
||
use wai::neural::jpegai_filters::efe_nonlinear_plane;
|
||
let fx_path = std::path::Path::new("../jpegai-vm-assets/e2e/efn_fixture.json");
|
||
if !fx_path.exists() {
|
||
eprintln!("skipping: EFEnonlinear fixture not present");
|
||
return;
|
||
}
|
||
let fx: serde_json::Value =
|
||
serde_json::from_slice(&std::fs::read(fx_path).unwrap()).unwrap();
|
||
let f32s = |k: &str| -> Vec<f32> {
|
||
fx[k].as_array().unwrap().iter().map(|n| n.as_f64().unwrap() as f32).collect()
|
||
};
|
||
let wq: Vec<i64> = fx["weights_q15"].as_array().unwrap().iter()
|
||
.map(|n| n.as_i64().unwrap()).collect();
|
||
let got = efe_nonlinear_plane(
|
||
&f32s("y"), &f32s("c_in"),
|
||
&[wq[0], wq[1], wq[2], wq[3], wq[4], wq[5], wq[6], wq[7]],
|
||
fx["luma_min_q"].as_i64().unwrap(), fx["luma_max_q"].as_i64().unwrap());
|
||
let want = f32s("c_want");
|
||
let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
|
||
assert!(worst < 1e-3, "EFEnonlinear worst error {worst}");
|
||
println!(" EFEnonlinear: V plane reproduced, worst |err| = {worst:.2e}");
|
||
}
|
||
|
||
/// Header parser vs the live-decode field trace: every typed field must
|
||
/// match what the reference decoder read from PIH + TON.
|
||
#[test]
|
||
fn jpegai_header_parser_matches_trace() {
|
||
use wai::neural::jpegai_header::parse_headers;
|
||
let assets = std::path::Path::new("../jpegai-vm-assets");
|
||
let bin_path = assets.join("e2e/test256.bin");
|
||
if !bin_path.exists() {
|
||
eprintln!("skipping: bitstream not present");
|
||
return;
|
||
}
|
||
let bin = std::fs::read(&bin_path).unwrap();
|
||
let subs = wai::neural::jpegai_bitstream::parse_codestream(&bin).unwrap();
|
||
let payload = |m: u16| subs.iter().find(|s| s.marker == m).unwrap().payload;
|
||
let h = parse_headers(payload(0xFF82), payload(0xFF83)).expect("parse");
|
||
assert_eq!(h.stream_profile_idc, 0);
|
||
assert_eq!(h.decoder_profile_id, 2);
|
||
assert_eq!(h.synthesis_transform_ids, vec![2, 1, 0]);
|
||
assert_eq!(h.level_idc, 52);
|
||
assert_eq!((h.img_width, h.img_height), (256, 256));
|
||
assert_eq!(h.colour_transform_idx, 1);
|
||
assert_eq!(h.model_id, 0);
|
||
assert_eq!(h.beta_displacement_log, 0);
|
||
assert_eq!(h.models[0].num_chs, 160);
|
||
assert_eq!(h.models[1].num_chs, 96);
|
||
assert_eq!(h.models[0].grfs_channel_flag.len(), 160);
|
||
assert_eq!(h.models[1].grfs_channel_flag.len(), 96);
|
||
assert_eq!(&h.models[0].grfs_channel_flag[..16],
|
||
&[1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]);
|
||
assert!(h.efe_linear_enabled);
|
||
assert_eq!((h.efe_linear.mean_u_q100, h.efe_linear.mean_v_q100), (12561, 13091));
|
||
assert_eq!(h.efe_linear.weights_u, vec![32677]);
|
||
assert_eq!(h.efe_linear.weights_v, vec![32758]);
|
||
assert_eq!(h.efe_linear.weights_u2, vec![32772]);
|
||
assert_eq!(h.efe_linear.weights_v2, vec![32776]);
|
||
assert!(h.icci_enabled && h.icci_short_list);
|
||
assert_eq!(h.icci_signalled_idx, [0, 1]);
|
||
assert!(h.efe_nonlinear_enabled && !h.efn_u_enabled && h.efn_v_enabled);
|
||
assert_eq!((h.efn_min_luma_q100, h.efn_max_luma_q100), (497, 25320));
|
||
assert_eq!(h.efn_weights, vec![32966, 32443, 33031, 32361, 33068, 33006, 31968, 33003]);
|
||
assert!(h.lef_enabled);
|
||
assert_eq!(h.lef_ch_idx, 0);
|
||
println!(" header parser: all typed fields match the live-decode trace");
|
||
}
|
||
|
||
|
||
/// The residual-variance bands, against the reference decoder's own tables: for every
|
||
/// `likely` (0 … 4351) and both GRFS branches, the band the exact decoder's edges give
|
||
/// selects exactly the σ addition and the dequant multiplier the reference's 4352-entry
|
||
/// tables hold at that index (`model_constants.json`, exported from the reference
|
||
/// decoder). The tables change value at 1004, 1353 and 1938.
|
||
#[test]
|
||
fn jpegai_rvs_bands_are_the_reference_tables() {
|
||
use wai::jpegai_pipeline::rvs_tables;
|
||
let (_, edges, additive) = wai::jpegai_ops::stage_constants();
|
||
let path = std::path::Path::new("../jpegai-vm-assets/model_constants.json");
|
||
if !path.exists() {
|
||
eprintln!("skipping: model constants not present");
|
||
return;
|
||
}
|
||
let mc: serde_json::Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
|
||
let ints = |k: &str| -> Vec<i64> {
|
||
mc[k].as_array().unwrap().iter().map(|n| n.as_i64().unwrap()).collect()
|
||
};
|
||
for flag in [false, true] {
|
||
let add = ints(&format!("rvs_sigma_t_flag{}", flag as u8));
|
||
let dq = ints(&format!("rvs_dq_t_flag{}", flag as u8));
|
||
assert_eq!((add.len(), dq.len()), (4352, 4352));
|
||
let (_, dq_table, add_table) = rvs_tables(flag);
|
||
assert_eq!(add_table, additive[if flag { 0 } else { 1 }], "the additive table is its derivation");
|
||
for likely in 0..4352i64 {
|
||
let band = edges.iter().filter(|&&e| likely >= e).count();
|
||
let i = likely as usize;
|
||
assert_eq!(add_table[band], add[i], "σ addition, flag {flag}, likely {likely}");
|
||
assert_eq!(dq_table[band], dq[i], "dequant multiplier, flag {flag}, likely {likely}");
|
||
}
|
||
}
|
||
println!(" RVS bands: both branches' σ additions and dequant multipliers equal the reference's tables at all 4352 indices");
|
||
}
|