diff --git a/wai/wai-rs/src/jpegai_pipeline.rs b/wai/wai-rs/src/jpegai_pipeline.rs index bd635fd05..e980fed95 100644 --- a/wai/wai-rs/src/jpegai_pipeline.rs +++ b/wai/wai-rs/src/jpegai_pipeline.rs @@ -23,6 +23,10 @@ //! — 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. //! +//! **Colour combine.** [`yuv_to_rgb`] / [`rgb_to_yuv`] convert the final decoded planes +//! between YUV and RGB (BT.709 full-range, data normalised to [0,1]) — the last pixel +//! transform before the output image. +//! //! **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 @@ -193,6 +197,50 @@ 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). +// BT.709 full-range luma coefficients (Rec. ITU-R BT.709). +const KR: f32 = 0.2126; +const KG: f32 = 0.7152; +const KB: f32 = 0.0722; +const KBY: f32 = 1.8556; // 2·(1−Kb) +const KRY: f32 = 1.5748; // 2·(1−Kr) + +/// **YUV → RGB** colour combine (BT.709 full-range). All planes normalised to [0,1]; +/// chroma carries the `+0.5` offset. `y`/`u`/`v` are the three planes (same length); +/// returns `(r, g, b)`. The last pixel transform before the output image. +/// +/// ```text +/// r = y + Kry·(v−½) +/// b = y + Kby·(u−½) +/// g = y − (Kb·Kby/Kg)·(u−½) − (Kr·Kry/Kg)·(v−½) +/// ``` +pub fn yuv_to_rgb(y: &[f32], u: &[f32], v: &[f32]) -> (Vec, Vec, Vec) { + assert!(y.len() == u.len() && u.len() == v.len(), "plane length mismatch"); + let gu = KB * KBY / KG; + let gv = KR * KRY / KG; + let (mut r, mut g, mut b) = (vec![0f32; y.len()], vec![0f32; y.len()], vec![0f32; y.len()]); + for i in 0..y.len() { + let (uu, vv) = (u[i] - 0.5, v[i] - 0.5); + r[i] = y[i] + KRY * vv; + b[i] = y[i] + KBY * uu; + g[i] = y[i] - gu * uu - gv * vv; + } + (r, g, b) +} + +/// **RGB → YUV** (BT.709 full-range), the inverse of [`yuv_to_rgb`]. Planes in [0,1]; +/// chroma returned with the `+0.5` offset. Used for round-trip checks / re-encoding. +pub fn rgb_to_yuv(r: &[f32], g: &[f32], b: &[f32]) -> (Vec, Vec, Vec) { + assert!(r.len() == g.len() && g.len() == b.len(), "plane length mismatch"); + let (mut y, mut u, mut v) = (vec![0f32; r.len()], vec![0f32; r.len()], vec![0f32; r.len()]); + for i in 0..r.len() { + let yi = KR * r[i] + KG * g[i] + KB * b[i]; + y[i] = yi; + u[i] = (b[i] - yi) / KBY + 0.5; + v[i] = (r[i] - yi) / KRY + 0.5; + } + (y, u, v) +} + /// **EFE linear filter** (post-filter) — a per-plane linear cross-component combine /// that refines one chroma plane `x` using the co-located luma `y`: /// @@ -423,4 +471,23 @@ mod tests { 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]); } + + #[test] + fn yuv_rgb_round_trips_and_greyscale() { + // grey (u=v=0.5) → r=g=b=y + let (r, g, b) = yuv_to_rgb(&[0.4], &[0.5], &[0.5]); + assert!((r[0] - 0.4).abs() < 1e-6 && (g[0] - 0.4).abs() < 1e-6 && (b[0] - 0.4).abs() < 1e-6); + + // round-trip rgb → yuv → rgb is the identity to float precision + let (rr, gg, bb): (Vec, Vec, Vec) = ( + vec![0.10, 0.55, 0.90, 0.33], vec![0.20, 0.60, 0.05, 0.77], vec![0.80, 0.15, 0.40, 0.50], + ); + let (y, u, v) = rgb_to_yuv(&rr, &gg, &bb); + let (r2, g2, b2) = yuv_to_rgb(&y, &u, &v); + for i in 0..rr.len() { + assert!((r2[i] - rr[i]).abs() < 1e-5, "r {} vs {}", r2[i], rr[i]); + assert!((g2[i] - gg[i]).abs() < 1e-5, "g {} vs {}", g2[i], gg[i]); + assert!((b2[i] - bb[i]).abs() < 1e-5, "b {} vs {}", b2[i], bb[i]); + } + } } diff --git a/wai/wai-rs/tests/jpegai_color_calibration.rs b/wai/wai-rs/tests/jpegai_color_calibration.rs new file mode 100644 index 000000000..c6ea7b1fc --- /dev/null +++ b/wai/wai-rs/tests/jpegai_color_calibration.rs @@ -0,0 +1,35 @@ +//! Calibration for [`wai::jpegai_pipeline::yuv_to_rgb`] — the BT.709 colour combine. +//! Compares against a reference-formula golden (deterministic YUV → RGB computed with +//! the exact Rec. ITU-R BT.709 full-range coefficients). Float arithmetic ⇒ +//! machine-precision match expected. +//! +//! env: WAI_JPEGAI_COLOR_FIXTURE = color_fixture.json { yuv:[f64], rgb:[f64], c, h, w } +//! (yuv/rgb are 3 planes stacked, row-major) +#![cfg(feature = "neural_int")] + +use wai::jpegai_pipeline::yuv_to_rgb; + +#[test] +fn yuv_to_rgb_matches_bt709_reference() { + let path = match std::env::var("WAI_JPEGAI_COLOR_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_COLOR_FIXTURE to color_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 yuv = f32v("yuv"); + let rgb = f32v("rgb"); + let plane = (v["h"].as_u64().unwrap() * v["w"].as_u64().unwrap()) as usize; + + let (y, u, vp) = (&yuv[0..plane], &yuv[plane..2 * plane], &yuv[2 * plane..3 * plane]); + let (r, g, b) = yuv_to_rgb(y, u, vp); + + let mut maxd = 0f32; + for (i, plane_vals) in [&r, &g, &b].iter().enumerate() { + for j in 0..plane { + maxd = maxd.max((plane_vals[j] - rgb[i * plane + j]).abs()); + } + } + eprintln!("yuv_to_rgb vs BT.709 reference: max|Δ| = {maxd:.7}"); + assert!(maxd < 1e-5, "colour combine diverged from BT.709 (max|Δ| = {maxd})"); +}