diff --git a/wai/wai-rs/src/jpegai_icci.rs b/wai/wai-rs/src/jpegai_icci.rs new file mode 100644 index 000000000..f00fe2d53 --- /dev/null +++ b/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, 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, 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, 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, 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, Vec, Vec, Vec, 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 { + (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[..]); + } +} diff --git a/wai/wai-rs/src/lib.rs b/wai/wai-rs/src/lib.rs index 287ff1090..b3897be85 100644 --- a/wai/wai-rs/src/lib.rs +++ b/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