From 4420a6927a7afac44ba619bf192f9b324c596813 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Tue, 14 Jul 2026 05:13:39 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Fpipeline::lef=5Ffilter=20?= =?UTF-8?q?=E2=80=94=20LEF=20adaptive-sharpness=20post-filter=20(float-exa?= =?UTF-8?q?ct)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second pixel-domain post-filter — a scale-gated unsharp mask on the luma plane: n = y_in / 255 blur = clamp( conv3x3(n, [[5,5,5],[5,24,5],[5,5,5]]/64), 0, 1 ) // valid; border keeps n factor = mag[k] where region in [thr[k], thr[k+1]) (k=0,1) or region >= thr[2]; else 1 out = clamp( blur + (n - blur)*factor, 0, 1 ) * 255 The 3x3 conv is "valid" (no padding) — interior pixels blur, the 1-pixel border passes through — matching the reference's pad-and-mask trick. `region` gates the sharpening strength: the scale grid (sigma, 16x16) nearest-upsampled to pixel resolution. Validated against the reference fixture (input luma + 16x16 scale + thr/mag → sharpened luma): max|Δ| = 1.5e-5 (f32-vs-f64), SNR 153 dB. Plus unit tests for the flat-plane invariance, the border pass-through, and the factor>1 sharpening gate. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_pipeline.rs | 71 +++++++++++++++++++++- wai/wai-rs/tests/jpegai_lef_calibration.rs | 43 +++++++++++++ 2 files changed, 113 insertions(+), 1 deletion(-) create mode 100644 wai/wai-rs/tests/jpegai_lef_calibration.rs diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index 2f16e8772..8c4a87a41 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/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 { + 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 = 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 { 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]); + } } diff --git a/wai/wai-rs/tests/jpegai_lef_calibration.rs b/wai/wai-rs/tests/jpegai_lef_calibration.rs new file mode 100644 index 000000000..4416e6c54 --- /dev/null +++ b/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::>(); + 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::() / 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})"); +} -- 2.54.0 (Apple Git-157)