35f9ba5175
The entropy-call schedule's content-dependent part — the per-position coding mask for the residual (SORP/SORS) decode — is now DERIVED, replacing the dumped schedule fixture. JPEG AI's skip_mode codes a position iff its skip-scale σ exceeds the skip threshold; the key was that the skip mask uses the σ surface BEFORE the rvs addition (sigma_skip), while build_indexes uses it AFTER (sigma_final) — so on the bucketed index the decision looks ambiguous, but on sigma_skip it is exact: mask = (sigma_skip > 382). Traced bit-exact from jpeg-ai-ref (coding_tools/skip_ls/skip_mode.py generate_skip_mask; common_modules.py shows scale_log excl/incl rvs). Validated against the reference's own dumped masks: 0 violations across all 65536 positions (SORP/y split 381|383, SORS/uv 382|384, strict >). Cube-flag skip would OR extra coded positions in, but rides the header substream (not the residual stream) and is absent in this bitstream. The mask and the channel count now come from the σ surface; the decode asserts the derived mask equals the fixture's per call while the fixture still exists. Full e2e (real bitstream -> RGB vs reference PNG) still green. Co-Authored-By: Claude Opus 4.8 <[email protected]>
466 lines
21 KiB
Rust
466 lines
21 KiB
Rust
//! 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 →
|
||
//! mANS entropy → integer σ-net → 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};
|
||
use wai::neural::jpegai_hyper::{build_indexes, HyperScaleDecoderInt, IntConvBlock};
|
||
use wai::neural::jpegai_mans::{build_decode_transitions, MansDecoder};
|
||
|
||
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()
|
||
.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
|
||
}
|
||
|
||
fn make_net(weights: &serde_json::Value, plane: &str) -> HyperScaleDecoderInt {
|
||
HyperScaleDecoderInt {
|
||
blocks: weights[plane]["blocks"].as_array().unwrap().iter().map(|b| IntConvBlock {
|
||
out_ch: b["out_ch"].as_u64().unwrap() as usize,
|
||
in_ch: b["in_ch"].as_u64().unwrap() as usize,
|
||
k: b["k"].as_u64().unwrap() as usize,
|
||
weight: i64s(&b["weight"]),
|
||
bias: i64s(&b["bias"]),
|
||
shift: i64s(&b["shift"]).iter().map(|&v| v as u32).collect(),
|
||
}).collect(),
|
||
d2s_block: weights[plane]["d2s_block"].as_u64().unwrap() as usize,
|
||
clip_max: weights[plane]["clip_max"].as_i64().unwrap(),
|
||
}
|
||
}
|
||
|
||
/// 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"));
|
||
|
||
// ---- normative entropy tables ----
|
||
let pmf: Vec<Vec<i64>> = norm["pmf"].as_array().unwrap().iter().map(i64s).collect();
|
||
let cdf: Vec<Vec<i64>> = pmf.iter().map(|row| {
|
||
let mut c = vec![0i64; row.len() + 1];
|
||
for (i, &p) in row.iter().enumerate() { c[i + 1] = c[i] + p; }
|
||
c
|
||
}).collect();
|
||
let transitions = build_decode_transitions(&pmf, &cdf);
|
||
let bounds: Vec<u8> = 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();
|
||
|
||
// ---- 1. codestream + entropy decode ----
|
||
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"));
|
||
|
||
// SOZ → z (factorized; per-channel CDFs are model-derived, from fixture)
|
||
let (mem_z, ends_z) = parse_ae_substream(payload(0xFF88), 1).unwrap();
|
||
let mut dec_z = MansDecoder::new(mem_z, ends_z[0]);
|
||
dec_z.set_sgm_transitions(&transitions, &bounds);
|
||
let mut z_dec: Vec<Vec<u8>> = Vec::new();
|
||
for call in ent["SOZ"]["calls"].as_array().unwrap() {
|
||
let cdfs: Vec<u8> = i64s(&call["pmfs"]).iter().map(|&v| v as u8).collect();
|
||
let channels = call["channels"].as_u64().unwrap() as usize;
|
||
let per = call["per_channel"].as_u64().unwrap() as usize;
|
||
let mut z = vec![0u8; channels * per];
|
||
dec_z.decode_factorize(&cdfs, &mut z, channels, per);
|
||
z_dec.push(z);
|
||
}
|
||
|
||
// ---- 2. z → integer σ-net → σ surfaces → bucket indexes ----
|
||
let mut sigma_final: Vec<Vec<i64>> = Vec::new(); // per plane [C*256]
|
||
let mut sigma_skip: Vec<Vec<i64>> = Vec::new();
|
||
let mut indexes: Vec<Vec<u8>> = Vec::new();
|
||
let mut z_net_planes: Vec<Vec<i64>> = Vec::new();
|
||
for (pi, plane) in ["y", "uv"].iter().enumerate() {
|
||
let off = i64s(&route["z_offsets"][plane]);
|
||
let per = z_dec[pi].len() / off.len();
|
||
let z_net: Vec<i64> = z_dec[pi].iter().enumerate()
|
||
.map(|(i, &v)| v as i64 - off[i / per]).collect();
|
||
let net = make_net(&weights, plane);
|
||
let i_s = net.forward(&z_net, 4, 4);
|
||
// σ gain vector = model constant mlog + header beta displacement
|
||
let gain: Vec<i64> = i64s(&consts["tools_0"][format!("gain_mlog_{plane}")])
|
||
.iter().map(|&g| g + hdr.beta_displacement_log).collect();
|
||
let hw = i_s.len() / gain.len();
|
||
let skip: Vec<i64> = i_s.iter().enumerate().map(|(i, &v)| v + gain[i / hw]).collect();
|
||
// rvs σ-addition = grfs-gated table lookup at the per-block likely index
|
||
let t_f1 = i64s(&consts["rvs_sigma_t_flag1"]);
|
||
let t_f0 = i64s(&consts["rvs_sigma_t_flag0"]);
|
||
let grfs = &hdr.models[pi].grfs_channel_flag;
|
||
let fin: Vec<i64> = skip.iter().enumerate().map(|(i, &v)| {
|
||
let ch = i / hw;
|
||
let (y, x) = ((i % 256) / 16, i % 16);
|
||
let mut sum = 0i64;
|
||
let (by, bx) = (y / 8 * 8, x / 8 * 8);
|
||
for yy in 0..8 {
|
||
for xx in 0..8 {
|
||
sum += skip[ch * 256 + (by + yy) * 16 + bx + xx];
|
||
}
|
||
}
|
||
let lk = ((sum + 32) >> 6).clamp(0, 4351) as usize;
|
||
let t = if grfs[ch] == 1 { &t_f1 } else { &t_f0 };
|
||
v + t[lk]
|
||
}).collect();
|
||
// cross-check vs the dumped fixture while both exist
|
||
let rvs_fix = i64s(&adds[plane]["rvs"]);
|
||
for (i, &f) in fin.iter().enumerate() {
|
||
assert_eq!(f, skip[i] + rvs_fix[i], "derived rvs mismatch at {i}");
|
||
}
|
||
indexes.push(build_indexes(&fin));
|
||
sigma_skip.push(skip);
|
||
sigma_final.push(fin);
|
||
z_net_planes.push(z_net);
|
||
}
|
||
|
||
// ---- 3. residual entropy decode (SORP=Y, SORS=UV) ----
|
||
// The per-position coding mask is DERIVED, not read from a schedule fixture:
|
||
// JPEG AI's skip_mode codes a position iff its skip-scale σ (the σ surface
|
||
// BEFORE the rvs addition — `sigma_skip`, not `sigma_final`) exceeds the
|
||
// skip threshold. `mask = (sigma_skip > 382)`; the channel count is the σ
|
||
// surface's, both from the bitstream. (Cube-flag skip would OR extra coded
|
||
// positions in; this stream carries none — those flags ride the header
|
||
// substream, not the residual stream, so they never shift its bit position.)
|
||
const SKIP_THR: i64 = 382; // jpeg-ai-ref skip_mode.thr_skip (unscaled σ precision)
|
||
let mut resi: Vec<Vec<i16>> = Vec::new();
|
||
for (pi, (marker, stream)) in [(0xFF89u16, "SORP"), (0xFF8A, "SORS")].iter().enumerate() {
|
||
let (mem, ends) = parse_ae_substream(payload(*marker), 1).unwrap();
|
||
let mut dec = MansDecoder::new(mem, ends[0]);
|
||
dec.set_sgm_transitions(&transitions, &bounds);
|
||
let nch = sigma_skip[pi].len() / 256; // channels = latent C, from the σ surface
|
||
let fixture_calls = ent[stream]["calls"].as_array().unwrap(); // cross-check only
|
||
let mut plane = vec![0i16; nch * 256];
|
||
for c in 0..nch {
|
||
let masks: Vec<bool> =
|
||
(0..256).map(|p| sigma_skip[pi][c * 256 + p] > SKIP_THR).collect();
|
||
// cross-check the derived schedule against the dumped one while it exists.
|
||
let fx: Vec<bool> = i64s(&fixture_calls[c]["masks"]).iter().map(|&v| v != 0).collect();
|
||
assert_eq!(masks, fx, "derived skip mask != fixture at {stream} ch {c}");
|
||
let mut r = vec![0i16; 256];
|
||
dec.decode_sgm(&indexes[pi][c * 256..(c + 1) * 256], &mut r, &masks);
|
||
plane[c * 256..(c + 1) * 256].copy_from_slice(&r);
|
||
}
|
||
resi.push(plane);
|
||
}
|
||
|
||
// ---- 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}");
|
||
}
|