wai/wai-rs/src/jpegai_pipeline.rs

@@ -135,6 +135,105 @@ pub fn dequant(y_q: &[i16], gain_map: &[f32]) -> Vec<f32> {
 /// 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<bool> {
+    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<i64> {
+    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<u8> {
     assert_eq!(i_s.len(), rvs.len(), "i_s / rvs length mismatch");

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::<Vec<i64>>();
+        let boolv = |k: &str| o[k].as_array().unwrap().iter().map(|x| x.as_bool().unwrap()).collect::<Vec<bool>>();
+
+        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<u8> = 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");
+    }
+}