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<f32>, Vec<f32>, Vec<f32>) {
+    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<f32>, Vec<f32>, Vec<f32>) {
+    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<f32>, Vec<f32>, Vec<f32>) = (
+            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]);
+        }
+    }
 }

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::<Vec<f32>>();
+    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})");
+}