From 954307816dadf393b5f64ad8a07e0f3cae1bbf9b Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Sat, 18 Jul 2026 23:30:54 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20derive=20the=20entropy-decode=20input?= =?UTF-8?q?s=20=E2=80=94=20skip=5Fmask=20+=20rvs=5Fadditive=20(fixture-fre?= =?UTF-8?q?e)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the last two inputs the mANS sgm decoder needed but that were still being read from an e2e fixture. With these, every input of the decode path is computable from the codestream plus normative assets. skip_mask scale_hat = i_s + gain_per_ch ; mask = scale_hat > 382 NOTE the skip decision uses the scale WITHOUT the RVS additive, while sigma_index uses it WITH — the two differ by exactly the rvs values and are easy to conflate. rvs_additive likely = clamp((sum_8x8 |i_s| + 32) >> 6, 0, 4351) band = #{ e in [1003,1352,1937] : likely >= e } rvs = A[grfs_flag ? 0 : 1][band] Band edges are precomputed as idx2log(thr * 2^4) against the reference Log2LinConvertion table, so the 4352-entry table is not needed at decode time; the per-band additives come from the normative scale list scaled by the GRFS branch factor. Band membership is INCLUSIVE at the lower edge — only the luma plane exercises this (one block lands exactly on 1352 and the reference places it in the band above), which is why both planes are tested. Validated against the ISO reference on real test256.bin, both planes, including the composition (derived rvs fed through sigma_index must reproduce the indices the reference entropy coder consumed): luma rvs 40960/40960, mask 40960/40960 (703 coded), sigma 40960/40960 chroma rvs 24576/24576, mask 24576/24576 (537 coded), sigma 24576/24576 Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_pipeline.rs | 99 +++++++++++++++++++ wai/wai-rs/tests/jpegai_derive_calibration.rs | 62 ++++++++++++ 2 files changed, 161 insertions(+) create mode 100644 wai/wai-rs/tests/jpegai_derive_calibration.rs diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index e980fed95..c6fb35adf 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -135,6 +135,105 @@ pub fn dequant(y_q: &[i16], gain_map: &[f32]) -> Vec { /// reference-decoder semantics (see `decode_hyper_scale`). Pure integer ⇒ /// byte-identical everywhere. /// +/// Skip threshold: positions whose scale does not exceed this are not coded. +const THR_SKIP: i64 = 382; + +/// RVS band edges in log-index space — `idx2log(threshold · 2^4)` for the normative +/// thresholds `[4325, 7471, 18667]` against the reference `Log2LinConvertion` table. +/// Precomputing them means the 4352-entry table is not needed at decode time. +const RVS_BAND_EDGES: [i64; 3] = [1003, 1352, 1937]; + +/// Per-band RVS log-domain additive, indexed `[cwg_branch][band]`. Derived from the +/// normative scale list `[115, 125, 146, 160]` scaled by the branch factor +/// (×1.06 when the channel's GRFS flag is set, ×0.98 otherwise) via +/// `a = round(ln(scale/128) / (log_k / 2^7))`, `log_k = (ln 54.82 − ln 0.11)/31`. +const RVS_ADDITIVE: [[i64; 4]; 2] = [ + [-31, 20, 122, 181], // GRFS flag set (scales 122, 132, 155, 170) + [-80, -31, 71, 130], // GRFS flag clear (scales 113, 122, 143, 157) +]; + +/// **Skip mask** — which latent positions carry a coded symbol. +/// +/// ```text +/// scale_hat[c,p] = i_s[c,p] + gain_per_ch[c] // NOTE: without the RVS additive +/// mask[c,p] = scale_hat[c,p] > 382 +/// ``` +/// +/// The skip decision uses the scale *before* the residual-variance additive, whereas +/// [`sigma_index`] uses it *after* — the two quantities differ, and conflating them +/// silently mis-selects which positions are coded. `skip_block_size` is 1 in this +/// profile, so the threshold applies per position with no block expansion. +/// +/// `i_s` is row-major `[channels × positions]`; `gain_per_ch` is per channel. +pub fn skip_mask(i_s: &[i64], gain_per_ch: &[i64], channels: usize) -> Vec { + assert_eq!(gain_per_ch.len(), channels, "gain_per_ch must have one entry per channel"); + assert_eq!(i_s.len() % channels, 0, "i_s length must be a multiple of channels"); + let plane = i_s.len() / channels; + let mut out = vec![false; i_s.len()]; + for c in 0..channels { + let g = gain_per_ch[c]; + for p in 0..plane { + let i = c * plane + p; + out[i] = i_s[i] + g > THR_SKIP; + } + } + out +} + +/// **Residual-variance scale (RVS)** — the per-position log-domain additive that +/// [`sigma_index`] folds in alongside the per-channel gain. +/// +/// ```text +/// likely[c,blk] = clamp( (Σ_{8×8} |i_s| + 32) >> 6, 0, 4351 ) // per 8×8 block +/// band = #{ e ∈ RVS_BAND_EDGES : likely > e } +/// rvs[c,p] = RVS_ADDITIVE[grfs_channel_flag[c] ? 0 : 1][band] +/// ``` +/// +/// The block statistic is a round-to-nearest of the 8×8 sum (an arithmetic shift), the +/// same form [`lsbs_refine`] uses. Band membership is INCLUSIVE at the lower edge — a +/// value landing exactly on an edge belongs to the upper band. (Only the luma plane +/// exercises this: one block lands exactly on edge 1352, and the reference places it in +/// the band above.) `grfs_channel_flags` comes from the picture header +/// (`ComponentHeader::grfs_channel_flags`) and selects the branch. +/// +/// `i_s` is row-major `[channels × h × w]`; returns the same shape. +pub fn rvs_additive(i_s: &[i64], grfs_channel_flags: &[bool], channels: usize, h: usize, w: usize) -> Vec { + const BLOCK: usize = 8; + assert_eq!(i_s.len(), channels * h * w, "i_s must be [channels, h, w]"); + assert_eq!(grfs_channel_flags.len(), channels, "one GRFS flag per channel"); + let plane = h * w; + let mut out = vec![0i64; i_s.len()]; + for c in 0..channels { + let base = c * plane; + let branch = if grfs_channel_flags[c] { 0 } else { 1 }; + let mut by = 0; + while by < h { + let mut bx = 0; + while bx < w { + let mut sum = 0i64; + for dy in 0..BLOCK { + for dx in 0..BLOCK { + let (y, x) = (by + dy, bx + dx); + if y < h && x < w { sum += i_s[base + y * w + x].abs(); } + } + } + let likely = ((sum + 32) >> 6).clamp(0, 4351); + let band = RVS_BAND_EDGES.iter().filter(|&&e| likely >= e).count(); + let value = RVS_ADDITIVE[branch][band]; + for dy in 0..BLOCK { + for dx in 0..BLOCK { + let (y, x) = (by + dy, bx + dx); + if y < h && x < w { out[base + y * w + x] = value; } + } + } + bx += BLOCK; + } + by += BLOCK; + } + } + out +} + /// `i_s` and `rvs` are row-major `[channels × positions]`; `gain_per_ch` is per channel. pub fn sigma_index(i_s: &[i64], gain_per_ch: &[i64], rvs: &[i64], channels: usize) -> Vec { assert_eq!(i_s.len(), rvs.len(), "i_s / rvs length mismatch"); diff --git a/wai/wai-rs/tests/jpegai_derive_calibration.rs b/wai/wai-rs/tests/jpegai_derive_calibration.rs new file mode 100644 index 000000000..745eb8986 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_derive_calibration.rs @@ -0,0 +1,62 @@ +//! Calibration for the two derived entropy-decode inputs — [`skip_mask`] and +//! [`rvs_additive`] — against the ISO reference on real `test256.bin` data, for BOTH +//! the luma and chroma planes. +//! +//! These close the last fixture dependency: with them, every input the mANS `sgm` +//! decoder needs is computable from the codestream plus normative assets. The test +//! therefore also checks the composition — that the derived RVS, fed through +//! [`sigma_index`], reproduces the distribution indices the reference entropy coder +//! actually consumed. +//! +//! env: WAI_JPEGAI_DERIVE_FIXTURE = derive_fixture.json +//! { y|uv: { channels, h, w, i_s, gain_per_ch, grfs, +//! rvs_expected, sigma_expected, mask_expected } } +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::{rvs_additive, sigma_index, skip_mask}; + +#[test] +fn derived_inputs_match_reference() { + let path = match std::env::var("WAI_JPEGAI_DERIVE_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_DERIVE_FIXTURE"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + + for plane in ["y", "uv"] { + let o = &v[plane]; + let g = |k: &str| o[k].as_u64().unwrap() as usize; + let (channels, h, w) = (g("channels"), g("h"), g("w")); + let i64v = |k: &str| o[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::>(); + let boolv = |k: &str| o[k].as_array().unwrap().iter().map(|x| x.as_bool().unwrap()).collect::>(); + + let i_s = i64v("i_s"); + let gain_per_ch = i64v("gain_per_ch"); + let grfs = boolv("grfs"); + let rvs_expected = i64v("rvs_expected"); + let mask_expected = boolv("mask_expected"); + let sigma_expected: Vec = o["sigma_expected"].as_array().unwrap() + .iter().map(|x| x.as_u64().unwrap() as u8).collect(); + + // 1. RVS additive + let rvs = rvs_additive(&i_s, &grfs, channels, h, w); + let bad = rvs.iter().zip(&rvs_expected).filter(|(a, b)| a != b).count(); + eprintln!("[{plane}] rvs_additive: {}/{} exact", rvs.len() - bad, rvs.len()); + assert_eq!(bad, 0, "{plane}: RVS additive diverged from the reference"); + + // 2. skip mask — note it uses the scale WITHOUT the RVS additive + let mask = skip_mask(&i_s, &gain_per_ch, channels); + let bad = mask.iter().zip(&mask_expected).filter(|(a, b)| a != b).count(); + eprintln!("[{plane}] skip_mask: {}/{} exact ({} coded)", + mask.len() - bad, mask.len(), mask.iter().filter(|&&b| b).count()); + assert_eq!(bad, 0, "{plane}: skip mask diverged from the reference"); + + // 3. composition — derived RVS must reproduce the distribution indices the + // reference entropy coder consumed + let sigma = sigma_index(&i_s, &gain_per_ch, &rvs, channels); + let bad = sigma.iter().zip(&sigma_expected).filter(|(a, b)| a != b).count(); + eprintln!("[{plane}] sigma_index: {}/{} exact (from derived rvs)", + sigma.len() - bad, sigma.len()); + assert_eq!(bad, 0, "{plane}: sigma indices diverged when driven by derived RVS"); + } +} -- 2.54.0 (Apple Git-157)