wai/wai-rs/src/jpegai_icci.rs

@@ -112,6 +112,136 @@ pub fn icci_dwt_planes(y: &[f32], u: &[f32], v: &[f32], h: usize, w: usize)
     (y16, u16, v16, yuv48, h2, w2)
 }
 
+// ---- network body ----
+//
+// The eICCI CNN is float in the reference, so we run it in f32 rather than quantising:
+// IEEE-754 add/mul with a fixed accumulation order is deterministic across machines, so
+// a faithful float port is byte-identical *and* avoids any PTQ error. BatchNorm is
+// folded into the preceding convolution, which is exact at inference:
+//
+//   BN(conv(x)) = γ·(w·x + b − μ)/√(σ²+ε) + β
+//               = (w · γ/√(σ²+ε))·x + ((b − μ)·γ/√(σ²+ε) + β)
+
+/// A convolution with its BatchNorm already folded in. Weights are `[cout, cin, kh, kw]`
+/// row-major; `bias` is per output channel. Padding is `(ph, pw)`, stride 1.
+#[derive(Clone)]
+pub struct Conv {
+    pub w: Vec<f32>,
+    pub b: Vec<f32>,
+    pub cin: usize,
+    pub cout: usize,
+    pub kh: usize,
+    pub kw: usize,
+    pub ph: usize,
+    pub pw: usize,
+}
+
+impl Conv {
+    /// Fold a BatchNorm into a convolution's weights/bias (exact at inference).
+    /// `gamma`/`beta`/`mean`/`var` are per output channel.
+    pub fn fold_bn(mut self, gamma: &[f32], beta: &[f32], mean: &[f32], var: &[f32], eps: f32) -> Self {
+        assert_eq!(gamma.len(), self.cout, "BN params must be per output channel");
+        let tap = self.cin * self.kh * self.kw;
+        for o in 0..self.cout {
+            let s = gamma[o] / (var[o] + eps).sqrt();
+            for t in 0..tap {
+                self.w[o * tap + t] *= s;
+            }
+            self.b[o] = (self.b[o] - mean[o]) * s + beta[o];
+        }
+        self
+    }
+
+    /// Stride-1 zero-padded convolution: `[cin, h, w]` → `[cout, h, w]`.
+    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
+        assert_eq!(x.len(), self.cin * h * w, "conv input must be [cin, h, w]");
+        let mut out = vec![0f32; self.cout * h * w];
+        for o in 0..self.cout {
+            for oy in 0..h {
+                for ox in 0..w {
+                    let mut acc = self.b[o];
+                    for i in 0..self.cin {
+                        for ky in 0..self.kh {
+                            let iy = oy + ky;
+                            if iy < self.ph || iy - self.ph >= h { continue; }
+                            let sy = iy - self.ph;
+                            for kx in 0..self.kw {
+                                let ix = ox + kx;
+                                if ix < self.pw || ix - self.pw >= w { continue; }
+                                let sx = ix - self.pw;
+                                acc += self.w[((o * self.cin + i) * self.kh + ky) * self.kw + kx]
+                                     * x[(i * h + sy) * w + sx];
+                            }
+                        }
+                    }
+                    out[(o * h + oy) * w + ox] = acc;
+                }
+            }
+        }
+        out
+    }
+}
+
+/// One `ResidualBlock_BN_RectKernel`: separable rectangular kernels with a scaled residual.
+///
+/// ```text
+/// out = x + scale · BN2( conv2_{3×1}( relu( BN1( conv1_{1×3}(x) ) ) ) )
+/// ```
+///
+/// Both BatchNorms are folded into their convolutions, so this is two convs + a ReLU.
+#[derive(Clone)]
+pub struct ResBlock {
+    pub conv1: Conv, // [1,3] kernel, pad [0,1]
+    pub conv2: Conv, // [3,1] kernel, pad [1,0]
+    pub scale: f32,
+}
+
+impl ResBlock {
+    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> Vec<f32> {
+        let mut t = self.conv1.forward(x, h, w);
+        for v in t.iter_mut() { if *v < 0.0 { *v = 0.0; } }
+        let t = self.conv2.forward(&t, h, w);
+        x.iter().zip(&t).map(|(&a, &b)| a + b * self.scale).collect()
+    }
+}
+
+/// One plane-path of the eICCI network: `conv_first` + BN + ReLU → N residual blocks →
+/// `conv_last`, then a scaled residual add onto the plane's own subbands.
+///
+/// ```text
+/// fea      = relu( BN( conv_first(yuv48) ) )      // BN folded into conv_first
+/// fea      = resblocks(fea)
+/// out16    = conv_last(fea) · scale + subbands16
+/// plane    = idwt2(out16)
+/// ```
+#[derive(Clone)]
+pub struct IcciPath {
+    pub conv_first: Conv, // 48 → nf, 3×3, pad 1 (BN folded)
+    pub blocks: Vec<ResBlock>,
+    pub conv_last: Conv, // nf → 16, 3×3, pad 1
+    pub scale: f32,
+}
+
+impl IcciPath {
+    /// Run the path and return the refined 16 subbands (before the inverse transform).
+    pub fn refine_subbands(&self, yuv48: &[f32], subbands16: &[f32], h: usize, w: usize) -> Vec<f32> {
+        let mut fea = self.conv_first.forward(yuv48, h, w);
+        for v in fea.iter_mut() { if *v < 0.0 { *v = 0.0; } }
+        for blk in &self.blocks {
+            fea = blk.forward(&fea, h, w);
+        }
+        let res = self.conv_last.forward(&fea, h, w);
+        assert_eq!(res.len(), subbands16.len(), "conv_last must emit the 16 subbands");
+        res.iter().zip(subbands16).map(|(&r, &s)| r * self.scale + s).collect()
+    }
+
+    /// Full plane: refine the subbands then invert the two-level transform.
+    pub fn forward_plane(&self, yuv48: &[f32], subbands16: &[f32], h: usize, w: usize) -> Vec<f32> {
+        let refined = self.refine_subbands(yuv48, subbands16, h, w);
+        haar_idwt2(&refined, h, w).0
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -168,4 +298,62 @@ mod tests {
         // the Y block is the leading 16 channels
         assert_eq!(&yuv48[..y16.len()], &y16[..]);
     }
+
+    fn ident_conv(c: usize, kh: usize, kw: usize, ph: usize, pw: usize) -> Conv {
+        // identity: centre tap 1 on the matching in/out channel, everything else 0
+        let mut w = vec![0f32; c * c * kh * kw];
+        for o in 0..c {
+            w[((o * c + o) * kh + kh / 2) * kw + kw / 2] = 1.0;
+        }
+        Conv { w, b: vec![0f32; c], cin: c, cout: c, kh, kw, ph, pw }
+    }
+
+    #[test]
+    fn conv_identity_and_zero_padding() {
+        let x = vec![1.0f32, 2.0, 3.0, 4.0];
+        // 1×3 identity with pad [0,1] reproduces the input
+        let o = ident_conv(1, 1, 3, 0, 1).forward(&x, 2, 2);
+        assert_eq!(o, x, "identity conv must be a no-op");
+
+        // a pure left-shift tap reads the zero pad at the left edge
+        let mut c = Conv { w: vec![0.0, 0.0, 1.0], b: vec![0.0], cin: 1, cout: 1, kh: 1, kw: 3, ph: 0, pw: 1 };
+        c.b[0] = 0.0;
+        let o = c.forward(&[5.0, 7.0], 1, 2);
+        // tap at kx=2 reads x[ox+2-1] = x[ox+1]: out[0]=x[1]=7, out[1]=pad=0
+        assert_eq!(o, vec![7.0, 0.0], "zero padding at the right edge");
+    }
+
+    #[test]
+    fn fold_bn_matches_explicit_batchnorm() {
+        // conv = identity, BN with known params → folded conv must equal BN(x) directly
+        let (gamma, beta, mean, var, eps) = ([2.0f32], [0.5f32], [1.0f32], [4.0f32], 1e-5f32);
+        let folded = ident_conv(1, 1, 1, 0, 0).fold_bn(&gamma, &beta, &mean, &var, eps);
+        let x = [3.0f32, -1.0];
+        let got = folded.forward(&x, 1, 2);
+        for i in 0..x.len() {
+            let want = gamma[0] * (x[i] - mean[0]) / (var[0] + eps).sqrt() + beta[0];
+            assert!((got[i] - want).abs() < 1e-5, "fold != BN: {} vs {want}", got[i]);
+        }
+    }
+
+    #[test]
+    fn resblock_is_identity_when_scale_is_zero() {
+        let c = 2usize;
+        let (h, w) = (3usize, 3usize);
+        let x = ramp(c * h * w);
+        let blk = ResBlock {
+            conv1: ident_conv(c, 1, 3, 0, 1),
+            conv2: ident_conv(c, 3, 1, 1, 0),
+            scale: 0.0,
+        };
+        assert_eq!(blk.forward(&x, h, w), x, "scale=0 must leave the identity path alone");
+
+        // scale=1 with identity convs adds relu(x) back onto x
+        let blk1 = ResBlock { scale: 1.0, ..blk };
+        let got = blk1.forward(&x, h, w);
+        for i in 0..x.len() {
+            let want = x[i] + if x[i] > 0.0 { x[i] } else { 0.0 };
+            assert!((got[i] - want).abs() < 1e-5, "at {i}: {} vs {want}", got[i]);
+        }
+    }
 }