AI Flow: Reverse Engineering

//! Image arithmetic the measurements are built on: float planes, masks, separable blurs, morphology by
//! distance, connected components and boundary tracing. Plain Rust, no image-processing crate, so every
//! number the flow reports can be traced to code in this repository.
use image::{Rgb, RgbImage};

#[derive(Clone)]
pub struct Plane { pub w: usize, pub h: usize, pub d: Vec<f32> }

#[derive(Clone)]
pub struct Mask { pub w: usize, pub h: usize, pub d: Vec<bool> }

impl Plane {
    pub fn new(w: usize, h: usize, v: f32) -> Plane { Plane { w, h, d: vec![v; w * h] } }
    #[inline] pub fn at(&self, x: usize, y: usize) -> f32 { self.d[y * self.w + x] }
    /// Bilinear sample, clamped at the edges.
    pub fn sample(&self, x: f32, y: f32) -> f32 {
        let x = x.clamp(0.0, (self.w - 1) as f32); let y = y.clamp(0.0, (self.h - 1) as f32);
        let (x0, y0) = (x.floor() as usize, y.floor() as usize);
        let (x1, y1) = ((x0 + 1).min(self.w - 1), (y0 + 1).min(self.h - 1));
        let (fx, fy) = (x - x0 as f32, y - y0 as f32);
        let a = self.at(x0, y0) * (1.0 - fx) + self.at(x1, y0) * fx;
        let b = self.at(x0, y1) * (1.0 - fx) + self.at(x1, y1) * fx;
        a * (1.0 - fy) + b * fy
    }
    pub fn threshold(&self, t: f32) -> Mask { Mask { w: self.w, h: self.h, d: self.d.iter().map(|v| *v > t).collect() } }
    pub fn mean_in(&self, m: &Mask) -> f32 {
        let (mut s, mut n) = (0.0f64, 0usize);
        for i in 0..self.d.len() { if m.d[i] { s += self.d[i] as f64; n += 1; } }
        if n == 0 { 0.0 } else { (s / n as f64) as f32 }
    }
}

impl Mask {
    pub fn new(w: usize, h: usize) -> Mask { Mask { w, h, d: vec![false; w * h] } }
    #[inline] pub fn at(&self, x: usize, y: usize) -> bool { self.d[y * self.w + x] }
    #[inline] pub fn set(&mut self, x: usize, y: usize, v: bool) { self.d[y * self.w + x] = v; }
    pub fn count(&self) -> usize { self.d.iter().filter(|v| **v).count() }
    pub fn and(&self, o: &Mask) -> Mask { Mask { w: self.w, h: self.h, d: self.d.iter().zip(&o.d).map(|(a, b)| *a && *b).collect() } }
    pub fn or(&self, o: &Mask) -> Mask { Mask { w: self.w, h: self.h, d: self.d.iter().zip(&o.d).map(|(a, b)| *a || *b).collect() } }
    pub fn not(&self) -> Mask { Mask { w: self.w, h: self.h, d: self.d.iter().map(|a| !*a).collect() } }
    pub fn andnot(&self, o: &Mask) -> Mask { Mask { w: self.w, h: self.h, d: self.d.iter().zip(&o.d).map(|(a, b)| *a && !*b).collect() } }

    /// Fill every background region that does not touch the image border.
    pub fn fill_holes(&self) -> Mask {
        let (w, h) = (self.w, self.h);
        let mut outside = vec![false; w * h];
        let mut st: Vec<usize> = vec![];
        for x in 0..w { for y in [0, h - 1] { let i = y * w + x; if !self.d[i] && !outside[i] { outside[i] = true; st.push(i); } } }
        for y in 0..h { for x in [0, w - 1] { let i = y * w + x; if !self.d[i] && !outside[i] { outside[i] = true; st.push(i); } } }
        while let Some(i) = st.pop() {
            let (x, y) = (i % w, i / w);
            let mut push = |j: usize| { if !self.d[j] && !outside[j] { outside[j] = true; st.push(j); } };
            if x > 0 { push(i - 1); } if x + 1 < w { push(i + 1); } if y > 0 { push(i - w); } if y + 1 < h { push(i + w); }
        }
        Mask { w, h, d: outside.iter().map(|o| !*o).collect() }
    }

    /// Fill enclosed background regions no larger than `max_area` pixels (markings inside a part, not the board inside a frame).
    pub fn fill_holes_upto(&self, max_area: usize) -> Mask {
        let holes = self.not();
        let (lab, cs) = components(&holes);
        let mut small = vec![false; cs.len() + 1];
        for c in &cs { small[c.label as usize] = c.area <= max_area && c.x0 > 0 && c.y0 > 0 && c.x1 + 1 < self.w && c.y1 + 1 < self.h; }
        Mask { w: self.w, h: self.h, d: self.d.iter().zip(&lab).map(|(s, l)| *s || small[*l as usize]).collect() }
    }

    /// Euclidean-ish distance (chamfer 3-4, in pixels) from every pixel to the nearest set pixel.
    pub fn distance_to_set(&self) -> Plane {
        let (w, h) = (self.w, self.h);
        let inf = 1e9f32;
        let mut d: Vec<f32> = self.d.iter().map(|v| if *v { 0.0 } else { inf }).collect();
        for y in 0..h { for x in 0..w {
            let i = y * w + x; let mut v = d[i];
            if x > 0 { v = v.min(d[i - 1] + 3.0); }
            if y > 0 { v = v.min(d[i - w] + 3.0); if x > 0 { v = v.min(d[i - w - 1] + 4.0); } if x + 1 < w { v = v.min(d[i - w + 1] + 4.0); } }
            d[i] = v;
        } }
        for y in (0..h).rev() { for x in (0..w).rev() {
            let i = y * w + x; let mut v = d[i];
            if x + 1 < w { v = v.min(d[i + 1] + 3.0); }
            if y + 1 < h { v = v.min(d[i + w] + 3.0); if x + 1 < w { v = v.min(d[i + w + 1] + 4.0); } if x > 0 { v = v.min(d[i + w - 1] + 4.0); } }
            d[i] = v;
        } }
        Plane { w, h, d: d.into_iter().map(|v| v / 3.0).collect() }
    }
    /// Grow by `r` pixels (disk).
    pub fn dilate(&self, r: f32) -> Mask { let d = self.distance_to_set(); Mask { w: self.w, h: self.h, d: d.d.iter().map(|v| *v <= r).collect() } }
    /// Shrink by `r` pixels (disk).
    pub fn erode(&self, r: f32) -> Mask { self.not().dilate(r).not() }
    pub fn open(&self, r: f32) -> Mask { self.erode(r).dilate(r) }
    pub fn close(&self, r: f32) -> Mask { self.dilate(r).erode(r) }
}

/// Separable box blur, `passes` times (three passes approximate a Gaussian of sigma ~ r).
pub fn blur(p: &Plane, r: usize, passes: usize) -> Plane {
    let (w, h) = (p.w, p.h);
    let mut a = p.d.clone();
    let mut b = vec![0f32; w * h];
    for _ in 0..passes {
        for y in 0..h {
            let row = &a[y * w..(y + 1) * w];
            let mut acc = 0f32; let mut n = 0f32;
            for x in 0..=r.min(w - 1) { acc += row[x]; n += 1.0; }
            for x in 0..w {
                b[y * w + x] = acc / n;
                if x + r + 1 < w { acc += row[x + r + 1]; n += 1.0; }
                if x >= r { acc -= row[x - r]; n -= 1.0; }
            }
        }
        for x in 0..w {
            let mut acc = 0f32; let mut n = 0f32;
            for y in 0..=r.min(h - 1) { acc += b[y * w + x]; n += 1.0; }
            for y in 0..h {
                a[y * w + x] = acc / n;
                if y + r + 1 < h { acc += b[(y + r + 1) * w + x]; n += 1.0; }
                if y >= r { acc -= b[(y - r) * w + x]; n -= 1.0; }
            }
        }
    }
    Plane { w, h, d: a }
}

/// Hue (degrees), saturation and value of an RGB pixel.
#[inline]
pub fn hsv(p: &Rgb<u8>) -> (f32, f32, f32) {
    let (r, g, b) = (p[0] as f32 / 255.0, p[1] as f32 / 255.0, p[2] as f32 / 255.0);
    let mx = r.max(g).max(b); let mn = r.min(g).min(b); let c = mx - mn;
    let h = if c < 1e-6 { 0.0 } else if mx == r { 60.0 * (((g - b) / c).rem_euclid(6.0)) } else if mx == g { 60.0 * ((b - r) / c + 2.0) } else { 60.0 * ((r - g) / c + 4.0) };
    (h, if mx > 0.0 { c / mx } else { 0.0 }, mx)
}

pub fn luma(img: &RgbImage) -> Plane {
    Plane { w: img.width() as usize, h: img.height() as usize, d: img.pixels().map(|p| 0.299 * p[0] as f32 + 0.587 * p[1] as f32 + 0.114 * p[2] as f32).collect() }
}

pub fn mask_from(img: &RgbImage, f: impl Fn(&Rgb<u8>) -> bool) -> Mask {
    Mask { w: img.width() as usize, h: img.height() as usize, d: img.pixels().map(|p| f(p)).collect() }
}

pub struct Component { pub label: u32, pub area: usize, pub x0: usize, pub y0: usize, pub x1: usize, pub y1: usize, pub cx: f64, pub cy: f64 }

/// 8-connected components of a mask: the label image (0 = background) and each component's stats.
pub fn components(m: &Mask) -> (Vec<u32>, Vec<Component>) {
    let (w, h) = (m.w, m.h);
    let mut lab = vec![0u32; w * h];
    let mut out = vec![];
    let mut st = vec![];
    for s in 0..w * h {
        if !m.d[s] || lab[s] != 0 { continue; }
        let l = out.len() as u32 + 1;
        lab[s] = l; st.push(s);
        let (mut area, mut x0, mut y0, mut x1, mut y1, mut sx, mut sy) = (0usize, usize::MAX, usize::MAX, 0usize, 0usize, 0f64, 0f64);
        while let Some(i) = st.pop() {
            let (x, y) = (i % w, i / w);
            area += 1; x0 = x0.min(x); y0 = y0.min(y); x1 = x1.max(x); y1 = y1.max(y); sx += x as f64; sy += y as f64;
            for dy in -1i64..=1 { for dx in -1i64..=1 {
                let (nx, ny) = (x as i64 + dx, y as i64 + dy);
                if nx < 0 || ny < 0 || nx >= w as i64 || ny >= h as i64 { continue; }
                let j = ny as usize * w + nx as usize;
                if m.d[j] && lab[j] == 0 { lab[j] = l; st.push(j); }
            } }
        }
        out.push(Component { label: l, area, x0, y0, x1, y1, cx: sx / area as f64, cy: sy / area as f64 });
    }
    (lab, out)
}

/// The outer boundary of one labelled component, through the middles of its edge pixels, in order (Moore neighbour tracing). A polygon through these is half a pixel inside the component all round.
pub fn trace_boundary(lab: &[u32], w: usize, h: usize, c: &Component) -> Vec<(f64, f64)> {
    let inside = |x: i64, y: i64| x >= 0 && y >= 0 && x < w as i64 && y < h as i64 && lab[y as usize * w + x as usize] == c.label;
    // start: the first pixel of the component in raster order
    let mut start = (0i64, 0i64);
    'f: for y in c.y0..=c.y1 { for x in c.x0..=c.x1 { if lab[y * w + x] == c.label { start = (x as i64, y as i64); break 'f; } } }
    let dirs = [(1i64, 0i64), (1, 1), (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1)];
    let mut pts = vec![(start.0 as f64 + 0.5, start.1 as f64 + 0.5)];
    let (mut cur, mut d) = (start, 6usize);
    for _ in 0..(c.area * 4 + 16) {
        let mut found = false;
        for k in 0..8 {
            let nd = (d + 6 + k) % 8; // start looking back-left of the last direction
            let n = (cur.0 + dirs[nd].0, cur.1 + dirs[nd].1);
            if inside(n.0, n.1) { cur = n; d = nd; found = true; break; }
        }
        if !found || cur == start { break; }
        pts.push((cur.0 as f64 + 0.5, cur.1 as f64 + 0.5));
    }
    pts
}

pub fn to_rgb_plane(p: &Plane) -> RgbImage {
    let (mn, mx) = p.d.iter().fold((f32::MAX, f32::MIN), |(a, b), v| (a.min(*v), b.max(*v)));
    let s = if mx > mn { 255.0 / (mx - mn) } else { 1.0 };
    RgbImage::from_fn(p.w as u32, p.h as u32, |x, y| { let v = ((p.at(x as usize, y as usize) - mn) * s) as u8; Rgb([v, v, v]) })
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn morphology_and_components() {
        let mut m = Mask::new(40, 20);
        for y in 5..15 { for x in 5..15 { m.set(x, y, true); } }
        for y in 5..15 { for x in 25..35 { m.set(x, y, true); } }
        m.set(10, 10, false); // a hole
        assert!(m.fill_holes().at(10, 10));
        let (_, cs) = components(&m);
        assert_eq!(cs.len(), 2);
        // a 10 x 10 square eroded by 2 px keeps its 6 x 6 core
        let mut sq = Mask::new(20, 20);
        for y in 5..15 { for x in 5..15 { sq.set(x, y, true); } }
        assert_eq!(sq.erode(2.0).count(), 36);
        let opened = sq.erode(2.0).dilate(2.0).count(); // opening rounds the four corners only
        assert!((84..100).contains(&opened), "{opened}");
        let b = trace_boundary(&components(&m).0, 40, 20, &components(&m).1[0]);
        assert!(b.len() >= 36 && b.len() <= 40, "{}", b.len());
    }
}