wai/wai-rs/src/jpegai_pipeline.rs

@@ -17,6 +17,10 @@
 //! correction that turns the reconstructed latent `yhat` into the synthesis input
 //! `y_lsbs` — a per-8×8-block table lookup keyed on the same `i_s`.
 //!
+//! **Post-filters (pixel domain).** [`efn_filter`] is the EFE luma-aided adaptive
+//! nonlinear chroma filter — a piecewise-linear, luma-guided refinement of a chroma
+//! plane applied after synthesis.
+//!
 //! **Latent reconstruction.** [`reconstruct_yhat`] is the SGMM autoregressive spine:
 //! it turns the dequantized residual back into the synthesis-input latent `y_hat` by
 //! running the four MCM context stages over a 2×2 checkerboard, each stage predicting
@@ -157,6 +161,38 @@ pub fn sigma_index(i_s: &[i64], gain_per_ch: &[i64], rvs: &[i64], channels: usiz
 /// `y_lsbs` to float precision (IEEE-754 deterministic).
 ///
 /// `yhat`/`resi`/`i_s` are row-major `[channels × h × w]`.
+/// **EFE luma-aided adaptive nonlinear chroma filter** (post-filter, single tile).
+/// Refines a chroma plane `c_in` using the co-located luma `y` (both same resolution):
+///
+/// ```text
+/// w[i]  = (weights_q15[i] − 32767) / 2^11          // deinteger, wP = 16
+/// min   = luma_min_q / 100 ;  gap = (luma_max_q/100 − min) / 8
+/// c_out = relu( w[0]·(y − min) + c_in )
+///       + Σ_{j=1..7} w[j] · relu( y − min − j·gap )
+/// ```
+///
+/// The two-conv/ReLU form in the reference collapses to this closed form: conv1 builds
+/// one luma-gated channel `w[0]·(y−min)+c_in` plus seven luma-only histogram channels
+/// `y − min − j·gap`; ReLU; conv2 sums them with weights `1, w[1..7]`. Float arithmetic
+/// (IEEE-754 deterministic). `y`/`c_in` are row-major `[h × w]`; returns the filtered
+/// chroma plane. (Single-tile / 4:4:4; caller downsamples luma for subsampled formats.)
+pub fn efn_filter(y: &[f32], c_in: &[f32], weights_q15: &[i64], luma_min_q: i64, luma_max_q: i64) -> Vec<f32> {
+    assert_eq!(y.len(), c_in.len(), "y / c_in length mismatch");
+    assert!(weights_q15.len() >= 8, "EFN needs 8 weights per tile");
+    let w: Vec<f32> = weights_q15.iter().map(|&q| (q - 32767) as f32 / 2048.0).collect();
+    let min = luma_min_q as f32 / 100.0;
+    let gap = (luma_max_q as f32 / 100.0 - min) / 8.0;
+    let relu = |x: f32| if x > 0.0 { x } else { 0.0 };
+    (0..y.len()).map(|i| {
+        let yi = y[i];
+        let mut acc = relu(w[0] * (yi - min) + c_in[i]);
+        for j in 1..8 {
+            acc += w[j] * relu(yi - min - j as f32 * gap);
+        }
+        acc
+    }).collect()
+}
+
 pub fn lsbs_refine(yhat: &[f32], resi: &[f32], i_s: &[i64], t0: &[i32], t1: &[i32],
                    channels: usize, h: usize, w: usize) -> Vec<f32> {
     const BLOCK: usize = 8;
@@ -275,4 +311,17 @@ mod tests {
         let o2 = lsbs_refine(&yhat, &resi, &big, &t0b, &vec![0i32; 4352], 1, 8, 8);
         assert!((o2[0] - (10.0 + (10.0 - 2.0))).abs() < 1e-6, "table index must clamp to 4351");
     }
+
+    #[test]
+    fn efn_filter_closed_form_and_relu_gates() {
+        // deinteger(32767) = 0, deinteger(32767 + 2048) = 1. Build w = [1, 1, 0,0,0,0,0,0].
+        // min = 100/100 = 1 ; gap = (900/100 - 1)/8 = 1.  So:
+        //   c_out = relu(1·(y-1) + c_in) + 1·relu(y-1-1) + 0·...
+        let w = [32767 + 2048, 32767 + 2048, 32767, 32767, 32767, 32767, 32767, 32767];
+        // pixel A: y=5, c_in=10 → relu(4+10)=14 ; + relu(5-2)=3 → 17
+        // pixel B: y=0, c_in=10 → relu(-1+10)=9 ; + relu(0-2)=0 → 9   (luma gate off)
+        let out = efn_filter(&[5.0, 0.0], &[10.0, 10.0], &w, 100, 900);
+        assert!((out[0] - 17.0).abs() < 1e-4, "got {}", out[0]);
+        assert!((out[1] - 9.0).abs() < 1e-4, "got {}", out[1]);
+    }
 }

wai/wai-rs/tests/jpegai_efn_calibration.rs

@@ -0,0 +1,41 @@
+//! Calibration for [`wai::jpegai_pipeline::efn_filter`] — the EFE luma-aided adaptive
+//! nonlinear chroma post-filter. Runs on the reference fixture (luma `y`, input chroma
+//! `c_in`, the 8 Q15 tile weights, luma min/max) and compares the filtered plane to the
+//! reference output `c_want`. Float arithmetic ⇒ machine-precision match expected.
+//!
+//! env: WAI_JPEGAI_EFN_FIXTURE = efn_fixture.json
+//!   { y:[f64], c_in:[f64], c_want:[f64], weights_q15:[i64], luma_min_q, luma_max_q }
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_pipeline::efn_filter;
+
+#[test]
+fn efn_filter_matches_reference() {
+    let path = match std::env::var("WAI_JPEGAI_EFN_FIXTURE") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_EFN_FIXTURE to efn_fixture.json"); return; }
+    };
+    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
+    let f32v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect::<Vec<f32>>();
+    let y = f32v("y");
+    let c_in = f32v("c_in");
+    let c_want = f32v("c_want");
+    let weights: Vec<i64> = v["weights_q15"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
+    let (lmin, lmax) = (v["luma_min_q"].as_i64().unwrap(), v["luma_max_q"].as_i64().unwrap());
+
+    let got = efn_filter(&y, &c_in, &weights, lmin, lmax);
+    assert_eq!(got.len(), c_want.len(), "output size mismatch");
+
+    let mut maxd = 0f32;
+    let mut mse = 0f64;
+    for i in 0..got.len() {
+        let d = (got[i] - c_want[i]).abs();
+        maxd = maxd.max(d);
+        mse += (d as f64) * (d as f64);
+    }
+    mse /= got.len() as f64;
+    let sig: f64 = c_want.iter().map(|&x| (x as f64) * (x as f64)).sum::<f64>() / got.len() as f64;
+    let snr = if mse > 0.0 { 10.0 * (sig / mse).log10() } else { f64::INFINITY };
+    eprintln!("efn_filter vs reference: max|Δ| = {maxd:.6}   SNR = {snr:.2} dB");
+    assert!(maxd < 1e-2, "EFN filter diverged from the reference (max|Δ| = {maxd})");
+}