wai/wai-rs/src/jpegai_pipeline.rs

@@ -19,7 +19,8 @@
 //!
 //! **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.
+//! plane applied after synthesis. [`lef_filter`] is the LEF adaptive-sharpness filter
+//! — a scale-gated unsharp mask on the luma plane.
 //!
 //! **Latent reconstruction.** [`reconstruct_yhat`] is the SGMM autoregressive spine:
 //! it turns the dequantized residual back into the synthesis-input latent `y_hat` by
@@ -176,6 +177,57 @@ pub fn sigma_index(i_s: &[i64], gain_per_ch: &[i64], rvs: &[i64], channels: usiz
 /// `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.)
+/// **LEF adaptive-sharpness filter** (post-filter) — a scale-gated unsharp mask on a
+/// pixel plane (luma). `y_in` is `[h × w]` in pixel units; `sigma` is the reference
+/// scale grid `[sh × sw]` (nearest-upsampled to pixel resolution to gate the strength).
+///
+/// ```text
+/// n        = y_in / 255
+/// blur     = clamp( conv3×3(n, [[5,5,5],[5,24,5],[5,5,5]]/64), 0, 1 )   // valid; border keeps n
+/// factor   = mag[k] where region ∈ [thr[k], thr[k+1]) (k=0,1) or region ≥ thr[2]; else 1
+/// out      = clamp( blur + (n − blur)·factor, 0, 1 ) · 255
+/// ```
+///
+/// The 3×3 convolution is "valid" (no padding): interior pixels are the blurred value,
+/// the 1-pixel border passes through unfiltered — matching the reference's pad-and-mask
+/// trick. `region` is `sigma` sampled nearest: `region[i,j] = sigma[⌊i·sh/h⌋, ⌊j·sw/w⌋]`.
+/// Float arithmetic (IEEE-754 deterministic).
+pub fn lef_filter(y_in: &[f32], sigma: &[f32], sh: usize, sw: usize,
+                  thr: [f32; 3], mag: [f32; 3], h: usize, w: usize) -> Vec<f32> {
+    assert_eq!(y_in.len(), h * w, "y_in must be [h × w]");
+    assert_eq!(sigma.len(), sh * sw, "sigma must be [sh × sw]");
+    const K: [[f32; 3]; 3] = [[5.0, 5.0, 5.0], [5.0, 24.0, 5.0], [5.0, 5.0, 5.0]];
+    let n: Vec<f32> = y_in.iter().map(|&v| v / 255.0).collect();
+    let mut out = vec![0f32; h * w];
+    for i in 0..h {
+        let si = (i * sh / h).min(sh - 1);
+        for j in 0..w {
+            let idx = i * w + j;
+            // valid 3×3 blur for interior pixels; border keeps the input
+            let result = if i >= 1 && i + 1 < h && j >= 1 && j + 1 < w {
+                let mut acc = 0f32;
+                for dy in 0..3 {
+                    for dx in 0..3 {
+                        acc += K[dy][dx] / 64.0 * n[(i + dy - 1) * w + (j + dx - 1)];
+                    }
+                }
+                acc.clamp(0.0, 1.0)
+            } else {
+                n[idx]
+            };
+            let sj = (j * sw / w).min(sw - 1);
+            let region = sigma[si * sw + sj];
+            let factor = if region >= thr[0] && region < thr[1] { mag[0] }
+                else if region >= thr[1] && region < thr[2] { mag[1] }
+                else if region >= thr[2] { mag[2] }
+                else { 1.0 };
+            let res = (result + (n[idx] - result) * factor).clamp(0.0, 1.0);
+            out[idx] = res * 255.0;
+        }
+    }
+    out
+}
+
 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");
@@ -324,4 +376,21 @@ mod tests {
         assert!((out[0] - 17.0).abs() < 1e-4, "got {}", out[0]);
         assert!((out[1] - 9.0).abs() < 1e-4, "got {}", out[1]);
     }
+
+    #[test]
+    fn lef_filter_border_blur_and_scale_gate() {
+        // 4×4 constant plane → blur == input; border passes through; a constant plane is
+        // unchanged regardless of factor (input == blurred → diff = 0).
+        let flat = vec![100.0f32; 16];
+        let sig = vec![1000.0f32; 4]; // 2×2 grid
+        let o = lef_filter(&flat, &sig, 2, 2, [500.0, 900.0, 1300.0], [1.1, 1.2, 1.3], 4, 4);
+        assert!(o.iter().all(|&v| (v - 100.0).abs() < 1e-3), "flat plane must be unchanged: {o:?}");
+
+        // sharpening gate: a bright center pixel; region 1000 → thr[1..2) → mag[1]=2.0.
+        // interior center gets blurred then sharpened back by factor 2 (over-shoots).
+        let mut img = vec![0.0f32; 9]; img[4] = 64.0; // 3×3, center = 64
+        let o2 = lef_filter(&img, &vec![1000.0f32; 1], 1, 1, [500.0, 900.0, 1300.0], [1.0, 2.0, 3.0], 3, 3);
+        // n_center=64/255; blur=24/64*n; res = blur + (n-blur)*2 ; must exceed the plain input
+        assert!(o2[4] > 64.0, "factor>1 must sharpen the center up: {}", o2[4]);
+    }
 }

wai/wai-rs/tests/jpegai_lef_calibration.rs

@@ -0,0 +1,43 @@
+//! Calibration for [`wai::jpegai_pipeline::lef_filter`] — the LEF adaptive-sharpness
+//! post-filter. Runs on the reference fixture (input luma, the 16×16 scale grid, the
+//! three thresholds/magnitudes) and compares against the reference `y_want`. Float
+//! arithmetic ⇒ machine-precision match expected.
+//!
+//! env: WAI_JPEGAI_LEF_FIXTURE = lef_fixture.json
+//!   { y_in:[f64], y_want:[f64], sigma:[f64], sh, sw, thr:[f64;3], mag:[f64;3], h, w }
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_pipeline::lef_filter;
+
+#[test]
+fn lef_filter_matches_reference() {
+    let path = match std::env::var("WAI_JPEGAI_LEF_FIXTURE") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_LEF_FIXTURE to lef_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_in = f32v("y_in");
+    let y_want = f32v("y_want");
+    let sigma = f32v("sigma");
+    let thr_v = f32v("thr");
+    let mag_v = f32v("mag");
+    let g = |k: &str| v[k].as_u64().unwrap() as usize;
+    let (sh, sw, h, w) = (g("sh"), g("sw"), g("h"), g("w"));
+
+    let got = lef_filter(&y_in, &sigma, sh, sw, [thr_v[0], thr_v[1], thr_v[2]], [mag_v[0], mag_v[1], mag_v[2]], h, w);
+    assert_eq!(got.len(), y_want.len(), "output size mismatch");
+
+    let mut maxd = 0f32;
+    let mut mse = 0f64;
+    for i in 0..got.len() {
+        let d = (got[i] - y_want[i]).abs();
+        maxd = maxd.max(d);
+        mse += (d as f64) * (d as f64);
+    }
+    mse /= got.len() as f64;
+    let sig: f64 = y_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!("lef_filter vs reference: max|Δ| = {maxd:.6}   SNR = {snr:.2} dB");
+    assert!(maxd < 1e-2, "LEF filter diverged from the reference (max|Δ| = {maxd})");
+}