diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index 8c4a87a41..bd635fd05 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -20,7 +20,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. [`lef_filter`] is the LEF adaptive-sharpness filter -//! — a scale-gated unsharp mask on the luma plane. +//! — a scale-gated unsharp mask on the luma plane. [`efl_chroma`] is the EFE *linear* +//! filter — a per-plane linear cross-component (luma-into-chroma) combine. //! //! **Latent reconstruction.** [`reconstruct_yhat`] is the SGMM autoregressive spine: //! it turns the dequantized residual back into the synthesis-input latent `y_hat` by @@ -192,6 +193,22 @@ pub fn sigma_index(i_s: &[i64], gain_per_ch: &[i64], rvs: &[i64], channels: usiz /// 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). +/// **EFE linear filter** (post-filter) — a per-plane linear cross-component combine +/// that refines one chroma plane `x` using the co-located luma `y`: +/// +/// ```text +/// out = w_uv·(x − mean) + w_y·y + mean +/// ``` +/// +/// `w_uv` is the chroma self-gain about its plane mean, `w_y` the luma cross-gain, and +/// `mean` the chroma plane mean (the reference's `bias`). The luma plane itself is left +/// unchanged. Float arithmetic (IEEE-754 deterministic). `x`/`y` are the two planes, +/// same length. +pub fn efl_chroma(x: &[f32], y: &[f32], w_uv: f32, w_y: f32, mean: f32) -> Vec { + assert_eq!(x.len(), y.len(), "chroma / luma length mismatch"); + x.iter().zip(y).map(|(&xi, &yi)| w_uv * (xi - mean) + w_y * yi + mean).collect() +} + 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]"); @@ -393,4 +410,17 @@ mod tests { // 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]); } + + #[test] + fn efl_chroma_cross_component_combine() { + // out = w_uv·(x-mean) + w_y·y + mean. w_uv=1 → x + w_y·(y) offset about mean. + let x = [130.0f32, 120.0]; + let y = [100.0f32, 100.0]; + let o = efl_chroma(&x, &y, 1.0, 0.01, 125.0); + // 1·(130-125) + 0.01·100 + 125 = 5 + 1 + 125 = 131 ; 1·(-5)+1+125 = 121 + assert!((o[0] - 131.0).abs() < 1e-4 && (o[1] - 121.0).abs() < 1e-4, "{o:?}"); + // w_uv=0 collapses to the luma cross-term about the mean + let o0 = efl_chroma(&x, &y, 0.0, 0.5, 125.0); + assert!((o0[0] - (0.5 * 100.0 + 125.0)).abs() < 1e-4, "{}", o0[0]); + } } diff --git a/wai/wai-rs/tests/jpegai_efl_calibration.rs b/wai/wai-rs/tests/jpegai_efl_calibration.rs new file mode 100644 index 000000000..486e577ea --- /dev/null +++ b/wai/wai-rs/tests/jpegai_efl_calibration.rs @@ -0,0 +1,33 @@ +//! Calibration for [`wai::jpegai_pipeline::efl_chroma`] — the EFE linear per-plane +//! cross-component filter. Runs on the reference fixture (luma + each chroma plane with +//! its `w_uv`/`w_y`/`mean`) and compares against the reference filtered plane. Float +//! arithmetic ⇒ machine-precision match expected. +//! +//! env: WAI_JPEGAI_EFL_FIXTURE = efl_fixture.json +//! { y:[f64], planes:[ { x:[f64], want:[f64], w_uv, w_y, mean } ] } +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::efl_chroma; + +#[test] +fn efl_chroma_matches_reference() { + let path = match std::env::var("WAI_JPEGAI_EFL_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_EFL_FIXTURE to efl_fixture.json"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let f32arr = |a: &serde_json::Value| a.as_array().unwrap().iter().map(|x| x.as_f64().unwrap() as f32).collect::>(); + let y = f32arr(&v["y"]); + + for (pi, p) in v["planes"].as_array().unwrap().iter().enumerate() { + let x = f32arr(&p["x"]); + let want = f32arr(&p["want"]); + let (w_uv, w_y, mean) = (p["w_uv"].as_f64().unwrap() as f32, p["w_y"].as_f64().unwrap() as f32, p["mean"].as_f64().unwrap() as f32); + let got = efl_chroma(&x, &y, w_uv, w_y, mean); + assert_eq!(got.len(), want.len(), "plane {pi}: size mismatch"); + let mut maxd = 0f32; + for i in 0..got.len() { maxd = maxd.max((got[i] - want[i]).abs()); } + eprintln!("efl_chroma plane {pi} vs reference: max|Δ| = {maxd:.6}"); + assert!(maxd < 1e-2, "plane {pi}: EFL diverged (max|Δ| = {maxd})"); + } +}