wai/wai-rs/src/jpegai_icci.rs

@@ -0,0 +1,171 @@
+//! **eICCI** (enhanced in-loop cross-component improvement) — the DWT-domain post-filter
+//! of the JPEG-AI decode path. This module holds the wavelet transform its CNN operates
+//! in; the network itself (per-tile model ensemble) builds on top.
+//!
+//! The transform is a **Haar wavelet on 2×2 blocks** with ½ scaling, applied **two
+//! levels**: a plane `[1, H, W]` becomes `[4, H/2, W/2]` and then `[16, H/4, W/4]`.
+//! Stacking the three planes' 16 subbands gives the 48 channels the network's first
+//! convolution consumes. Perfect reconstruction: [`haar_idwt`] inverts [`haar_dwt`]
+//! exactly (`LL−HL−LH+HH = 4a` ⇒ `rec1 = a`), so `idwt2 ∘ dwt2 = id` to float precision.
+//!
+//! Channel layout matches the reference (`cat` along the channel axis): one level emits
+//! `[LL(c), HL(c), LH(c), HH(c)]`; two levels emit the four level-1 bands each expanded
+//! into its own four level-2 bands, i.e. `x16[i*4 .. i*4+4]` is the level-2 decomposition
+//! of level-1 band `i`.
+
+/// One-level Haar DWT: `[c, h, w]` → `[4c, h/2, w/2]`, channels grouped `LL, HL, LH, HH`.
+/// `h` and `w` must be even. Returns the coefficients and the output dims.
+pub fn haar_dwt(x: &[f32], c: usize, h: usize, w: usize) -> (Vec<f32>, usize, usize) {
+    assert_eq!(x.len(), c * h * w, "input must be [c, h, w]");
+    assert!(h % 2 == 0 && w % 2 == 0, "DWT needs even dims, got {h}×{w}");
+    let (ho, wo) = (h / 2, w / 2);
+    let plane_o = ho * wo;
+    let mut out = vec![0f32; 4 * c * plane_o];
+    for ch in 0..c {
+        let base = ch * h * w;
+        for i in 0..ho {
+            for j in 0..wo {
+                // the 2×2 block, pre-scaled by ½ (matching the reference)
+                let a = x[base + (2 * i) * w + 2 * j] / 2.0; // even row, even col
+                let b = x[base + (2 * i + 1) * w + 2 * j] / 2.0; // odd row, even col
+                let cc = x[base + (2 * i) * w + 2 * j + 1] / 2.0; // even row, odd col
+                let d = x[base + (2 * i + 1) * w + 2 * j + 1] / 2.0; // odd row, odd col
+                let p = i * wo + j;
+                out[(ch) * plane_o + p] = a + b + cc + d; // LL
+                out[(c + ch) * plane_o + p] = -a - b + cc + d; // HL
+                out[(2 * c + ch) * plane_o + p] = -a + b - cc + d; // LH
+                out[(3 * c + ch) * plane_o + p] = a - b - cc + d; // HH
+            }
+        }
+    }
+    (out, ho, wo)
+}
+
+/// One-level Haar IDWT: `[4c, h, w]` → `[c, 2h, 2w]` — the exact inverse of [`haar_dwt`].
+pub fn haar_idwt(x: &[f32], c: usize, h: usize, w: usize) -> (Vec<f32>, usize, usize) {
+    assert_eq!(x.len(), 4 * c * h * w, "input must be [4c, h, w]");
+    let (ho, wo) = (h * 2, w * 2);
+    let plane_i = h * w;
+    let mut out = vec![0f32; c * ho * wo];
+    for ch in 0..c {
+        let base = ch * ho * wo;
+        for i in 0..h {
+            for j in 0..w {
+                let p = i * w + j;
+                let ll = x[(ch) * plane_i + p] / 2.0;
+                let hl = x[(c + ch) * plane_i + p] / 2.0;
+                let lh = x[(2 * c + ch) * plane_i + p] / 2.0;
+                let hh = x[(3 * c + ch) * plane_i + p] / 2.0;
+                out[base + (2 * i) * wo + 2 * j] = ll - hl - lh + hh; // even row, even col
+                out[base + (2 * i + 1) * wo + 2 * j] = ll - hl + lh - hh; // odd row, even col
+                out[base + (2 * i) * wo + 2 * j + 1] = ll + hl - lh - hh; // even row, odd col
+                out[base + (2 * i + 1) * wo + 2 * j + 1] = ll + hl + lh + hh; // odd row, odd col
+            }
+        }
+    }
+    (out, ho, wo)
+}
+
+/// Two-level Haar DWT of a single plane: `[1, h, w]` → `[16, h/4, w/4]`.
+/// Level 1 gives four bands; each is decomposed again, and the results are concatenated
+/// in band order — `out[i*4 .. i*4+4]` is the level-2 decomposition of level-1 band `i`.
+pub fn haar_dwt2(x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
+    assert_eq!(x.len(), h * w, "input must be a single [h, w] plane");
+    let (x4, h1, w1) = haar_dwt(x, 1, h, w); // [4, h/2, w/2]
+    let plane1 = h1 * w1;
+    let (h2, w2) = (h1 / 2, w1 / 2);
+    let plane2 = h2 * w2;
+    let mut out = vec![0f32; 16 * plane2];
+    for band in 0..4 {
+        let (sub, _, _) = haar_dwt(&x4[band * plane1..(band + 1) * plane1], 1, h1, w1); // [4, h/4, w/4]
+        out[band * 4 * plane2..(band + 1) * 4 * plane2].copy_from_slice(&sub);
+    }
+    (out, h2, w2)
+}
+
+/// Two-level Haar IDWT: `[16, h, w]` → `[1, 4h, 4w]` — the exact inverse of [`haar_dwt2`].
+pub fn haar_idwt2(x16: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
+    assert_eq!(x16.len(), 16 * h * w, "input must be [16, h, w]");
+    let plane = h * w;
+    let (h1, w1) = (h * 2, w * 2);
+    let plane1 = h1 * w1;
+    let mut x4 = vec![0f32; 4 * plane1];
+    for band in 0..4 {
+        let (rec, _, _) = haar_idwt(&x16[band * 4 * plane..(band + 1) * 4 * plane], 1, h, w); // [1, 2h, 2w]
+        x4[band * plane1..(band + 1) * plane1].copy_from_slice(&rec);
+    }
+    haar_idwt(&x4, 1, h1, w1) // [1, 4h, 4w]
+}
+
+/// Stack the three planes' two-level subbands into the 48-channel tensor the eICCI
+/// network consumes: `cat(Y16, U16, V16)` at `[48, h/4, w/4]`. All planes must be the
+/// same `h × w` (4:4:4).
+pub fn icci_dwt_planes(y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize)
+    -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>, usize, usize) {
+    let (y16, h2, w2) = haar_dwt2(y, h, w);
+    let (u16, _, _) = haar_dwt2(u, h, w);
+    let (v16, _, _) = haar_dwt2(v, h, w);
+    let mut yuv48 = Vec::with_capacity(48 * h2 * w2);
+    yuv48.extend_from_slice(&y16);
+    yuv48.extend_from_slice(&u16);
+    yuv48.extend_from_slice(&v16);
+    (y16, u16, v16, yuv48, h2, w2)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn ramp(n: usize) -> Vec<f32> {
+        (0..n).map(|i| ((i * 37 % 251) as f32) - 125.0).collect()
+    }
+
+    #[test]
+    fn haar_dwt_single_block_matches_hand_math() {
+        // one 2×2 block: a=(0,0)=1, c=(0,1)=2, b=(1,0)=3, d=(1,1)=4  (row-major [1,2,3,4])
+        let (o, ho, wo) = haar_dwt(&[1.0, 2.0, 3.0, 4.0], 1, 2, 2);
+        assert_eq!((ho, wo), (1, 1));
+        // halves: a=.5 b=1.5 c=1 d=2  →  LL=5, HL=1, LH=2, HH=0
+        assert!((o[0] - 5.0).abs() < 1e-6, "LL {}", o[0]);
+        assert!((o[1] - 1.0).abs() < 1e-6, "HL {}", o[1]);
+        assert!((o[2] - 2.0).abs() < 1e-6, "LH {}", o[2]);
+        assert!((o[3] - 0.0).abs() < 1e-6, "HH {}", o[3]);
+    }
+
+    #[test]
+    fn haar_idwt_inverts_dwt_exactly() {
+        let (h, w) = (8usize, 6usize);
+        let x = ramp(h * w);
+        let (d, dh, dw) = haar_dwt(&x, 1, h, w);
+        let (r, rh, rw) = haar_idwt(&d, 1, dh, dw);
+        assert_eq!((rh, rw), (h, w));
+        for i in 0..x.len() {
+            assert!((r[i] - x[i]).abs() < 1e-4, "idwt∘dwt != id at {i}: {} vs {}", r[i], x[i]);
+        }
+    }
+
+    #[test]
+    fn two_level_round_trips_and_has_16_bands() {
+        let (h, w) = (16usize, 16usize);
+        let x = ramp(h * w);
+        let (x16, h2, w2) = haar_dwt2(&x, h, w);
+        assert_eq!((h2, w2), (4, 4));
+        assert_eq!(x16.len(), 16 * 4 * 4, "two levels must give 16 subbands at h/4");
+        let (r, rh, rw) = haar_idwt2(&x16, h2, w2);
+        assert_eq!((rh, rw), (h, w));
+        for i in 0..x.len() {
+            assert!((r[i] - x[i]).abs() < 1e-3, "idwt2∘dwt2 != id at {i}: {} vs {}", r[i], x[i]);
+        }
+    }
+
+    #[test]
+    fn planes_stack_to_48_channels() {
+        let (h, w) = (8usize, 8usize);
+        let (y, u, v) = (ramp(h * w), ramp(h * w), ramp(h * w));
+        let (y16, _, _, yuv48, h2, w2) = icci_dwt_planes(&y, &u, &v, h, w);
+        assert_eq!((h2, w2), (2, 2));
+        assert_eq!(yuv48.len(), 48 * h2 * w2, "network input must be 48 channels");
+        // the Y block is the leading 16 channels
+        assert_eq!(&yuv48[..y16.len()], &y16[..]);
+    }
+}

wai/wai-rs/src/lib.rs

@@ -83,6 +83,10 @@ pub mod jpegai_mans;
 #[cfg(feature = "neural_int")]
 pub mod jpegai_pipeline;
 
+// JPEG-AI eICCI post-filter: the Haar wavelet domain its CNN operates in.
+#[cfg(feature = "neural_int")]
+pub mod jpegai_icci;
+
 // Deterministic integer generative synthesis (`wai.generate.*`) — the
 // `intent=create` analog of the integer-exact codecs: a generator whose entire
 // pipeline is integer/fixed-point, pinned to (model, prompt, seed, steps), so the