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"); + } +}