AI Flow: Reverse Engineering

//! Register a photo to the board frame through shared holes. Four corners cannot do it: a fitted corner is off by
//! several pixels and its error spreads over the frame as a warp (up to 0.58 mm between the Driver faces).
//! Through-holes appear on both faces and are round, so they pin the bottom photo to the top with a
//! quadratic map, which is what a small error in the homography looks like.
use crate::geom::{self, P};
use image::RgbImage;
use serde_json::{json, Value};
use std::path::Path;

/// Quadratic map: x' = c[0] + c[1] X + c[2] Y + c[3] X^2 + c[4] XY + c[5] Y^2, with X, Y scaled about the centre of the data.
#[derive(Clone, Debug)]
pub struct Quad { pub cx: [f64; 6], pub cy: [f64; 6], pub mid: P, pub half: P }

impl Quad {
    fn basis(&self, p: P) -> [f64; 6] { let (x, y) = ((p.0 - self.mid.0) / self.half.0, (p.1 - self.mid.1) / self.half.1); [1.0, x, y, x * x, x * y, y * y] }
    pub fn apply(&self, p: P) -> P { let b = self.basis(p); (b.iter().zip(&self.cx).map(|(a, c)| a * c).sum(), b.iter().zip(&self.cy).map(|(a, c)| a * c).sum()) }
    pub fn fit(pairs: &[(P, P)]) -> Option<Quad> {
        if pairs.len() < 8 { return None; }
        let (x0, x1) = pairs.iter().fold((f64::MAX, f64::MIN), |a, q| (a.0.min(q.0 .0), a.1.max(q.0 .0)));
        let (y0, y1) = pairs.iter().fold((f64::MAX, f64::MIN), |a, q| (a.0.min(q.0 .1), a.1.max(q.0 .1)));
        let mut q = Quad { cx: [0.0; 6], cy: [0.0; 6], mid: ((x0 + x1) / 2.0, (y0 + y1) / 2.0), half: (((x1 - x0) / 2.0).max(1.0), ((y1 - y0) / 2.0).max(1.0)) };
        let mut ata = [[0f64; 6]; 6]; let (mut atx, mut aty) = ([0f64; 6], [0f64; 6]);
        for (a, b) in pairs { let f = q.basis(*a); for i in 0..6 { for j in 0..6 { ata[i][j] += f[i] * f[j]; } atx[i] += f[i] * b.0; aty[i] += f[i] * b.1; } }
        q.cx = solve6(ata, atx)?; q.cy = solve6(ata, aty)?;
        Some(q)
    }
}

fn solve6(mut a: [[f64; 6]; 6], mut b: [f64; 6]) -> Option<[f64; 6]> {
    for c in 0..6 {
        let p = (c..6).max_by(|i, j| a[*i][c].abs().partial_cmp(&a[*j][c].abs()).unwrap())?;
        if a[p][c].abs() < 1e-12 { return None; }
        a.swap(c, p); b.swap(c, p);
        for r in 0..6 { if r != c { let f = a[r][c] / a[c][c]; for k in c..6 { a[r][k] -= f * a[c][k]; } b[r] -= f * b[c]; } }
    }
    let mut x = [0.0; 6]; for i in 0..6 { x[i] = b[i] / a[i][i]; } Some(x)
}

/// Pair holes of a photo with the top-frame holes (board px), nearest within `within`, each top hole used once at most.
pub fn pairs(top: &[P], other: &[P], within: f64) -> Vec<(P, P)> {
    let mut out = vec![];
    for o in other {
        if let Some(t) = top.iter().min_by(|a, b| geom::dist(**a, *o).partial_cmp(&geom::dist(**b, *o)).unwrap()) {
            if geom::dist(*t, *o) < within && !out.iter().any(|(q, _): &(P, P)| *q == *t) { out.push((*t, *o)); }
        }
    }
    out
}

/// What the registration achieved: the residual before and after, and the residual of each pair left out of the fit.
pub fn report(prs: &[(P, P)], q: &Quad, res: f64) -> Value {
    let rms = |v: &[f64]| (v.iter().map(|x| x * x).sum::<f64>() / v.len().max(1) as f64).sqrt();
    let before: Vec<f64> = { let n = prs.len() as f64; let (dx, dy) = prs.iter().fold((0.0, 0.0), |a, p| (a.0 + p.0 .0 - p.1 .0, a.1 + p.0 .1 - p.1 .1)); prs.iter().map(|p| geom::dist((p.1 .0 + dx / n, p.1 .1 + dy / n), p.0)).collect() };
    let after: Vec<f64> = prs.iter().map(|p| geom::dist(q.apply(p.0), p.1)).collect();
    let loo: Vec<f64> = (0..prs.len()).filter_map(|k| { let rest: Vec<(P, P)> = prs.iter().enumerate().filter(|(i, _)| *i != k).map(|(_, x)| *x).collect(); Quad::fit(&rest).map(|f| geom::dist(f.apply(prs[k].0), prs[k].1)) }).collect();
    let r3 = |x: f64| (x / res * 1000.0).round() / 1000.0;
    json!({"pairs": prs.len(), "translationOnlyRmsMm": r3(rms(&before)), "translationOnlyMaxMm": r3(before.iter().cloned().fold(0.0, f64::max)), "quadraticRmsMm": r3(rms(&after)), "leaveOneOutRmsMm": r3(rms(&loo)), "leaveOneOutMaxMm": r3(loo.iter().cloned().fold(0.0, f64::max))})
}

/// The photo drawn into the board frame through its corners, then through the quadratic map from the top frame to
/// where the corners alone put each point. `frame_margin` pixels are kept on every side (for reading the edge).
pub fn warp(path: &Path, corners: &[P; 4], dst: &[P; 4], q: Option<&Quad>, w: u32, h: u32, frame_margin: i64) -> Result<RgbImage, String> {
    let img = image::open(path).map_err(|e| e.to_string())?.to_rgb8();
    let hm = geom::homography(corners, dst).ok_or("degenerate corners")?;
    let inv = geom::invert(&hm).ok_or("not invertible")?;
    let (iw, ih) = (img.width() as usize, img.height() as usize);
    let raw = img.as_raw();
    let sample = |x: f64, y: f64, c: usize| -> f64 {
        let x = x.clamp(0.0, iw as f64 - 1.001); let y = y.clamp(0.0, ih as f64 - 1.001);
        let (x0, y0) = (x as usize, y as usize); let (fx, fy) = (x - x0 as f64, y - y0 as f64);
        let at = |a: usize, b: usize| raw[(b * iw + a) * 3 + c] as f64;
        (at(x0, y0) * (1.0 - fx) + at(x0 + 1, y0) * fx) * (1.0 - fy) + (at(x0, y0 + 1) * (1.0 - fx) + at(x0 + 1, y0 + 1) * fx) * fy
    };
    Ok(RgbImage::from_fn(w, h, |u, v| {
        let p = (u as f64 + 0.5 - frame_margin as f64, v as f64 + 0.5 - frame_margin as f64);
        let p = q.map(|m| m.apply(p)).unwrap_or(p);
        let s = geom::apply(&inv, p);
        image::Rgb([0, 1, 2].map(|c| sample(s.0 - 0.5, s.1 - 0.5, c).clamp(0.0, 255.0) as u8))
    }))
}

/// Where each patch of `primary` lies in `other`, both in the same frame to within a few pixels: patch matching on
/// luminance, sub-pixel by a parabola. Only photos are compared, never a layer. Returns (p, where p is in `other`).
pub fn patch_shifts(primary: &RgbImage, other: &RgbImage) -> Vec<(P, P)> {
    let (w, h) = (primary.width() as i64, primary.height() as i64);
    let luma = |i: &RgbImage| -> Vec<f32> { i.pixels().map(|p| 0.299 * p[0] as f32 + 0.587 * p[1] as f32 + 0.114 * p[2] as f32).collect() };
    let (a, b) = (luma(primary), luma(other));
    let at = |v: &[f32], x: i64, y: i64| v[(y * w + x) as usize];
    let (half, reach) = (30i64, 4i64);
    let mut out = vec![];
    let mut cy = half + reach + 1;
    while cy < h - half - reach - 1 {
        let mut cx = half + reach + 1;
        while cx < w - half - reach - 1 {
            let (mut mean, mut var) = (0f32, 0f32);
            for y in -half..=half { for x in -half..=half { mean += at(&a, cx + x, cy + y); } }
            mean /= ((2 * half + 1) * (2 * half + 1)) as f32;
            for y in -half..=half { for x in -half..=half { var += (at(&a, cx + x, cy + y) - mean).powi(2); } }
            var /= ((2 * half + 1) * (2 * half + 1)) as f32;
            // a patch with structure in it: from these photos, silk or a pad gives a variance far above grain
            if var > 900.0 {
                let ssd = |dx: i64, dy: i64| -> f32 { let mut s = 0f32; for y in (-half..=half).step_by(2) { for x in (-half..=half).step_by(2) { s += (at(&a, cx + x, cy + y) - at(&b, cx + x + dx, cy + y + dy)).powi(2); } } s };
                let mut best = (0i64, 0i64, f32::MAX);
                for dy in -reach..=reach { for dx in -reach..=reach { let v = ssd(dx, dy); if v < best.2 { best = (dx, dy, v); } } }
                if best.0.abs() < reach && best.1.abs() < reach {
                    let para = |m: f32, c: f32, p: f32| { let d = m - 2.0 * c + p; if d.abs() < 1e-3 { 0.0 } else { 0.5 * (m - p) / d } };
                    let fx = para(ssd(best.0 - 1, best.1), best.2, ssd(best.0 + 1, best.1)) as f64;
                    let fy = para(ssd(best.0, best.1 - 1), best.2, ssd(best.0, best.1 + 1)) as f64;
                    out.push(((cx as f64 + 0.5, cy as f64 + 0.5), (cx as f64 + 0.5 + best.0 as f64 + fx, cy as f64 + 0.5 + best.1 as f64 + fy)));
                }
            }
            cx += 100;
        }
        cy += 100;
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn a_quadratic_warp_is_recovered() {
        let truth = |p: P| (p.0 + 3.0 + 1e-3 * p.1 + 1e-5 * p.0 * p.1, p.1 - 2.0 + 2e-5 * p.0 * p.0);
        let prs: Vec<(P, P)> = (0..30).map(|i| { let p = ((i * 97 % 2900) as f64, (i * 211 % 2400) as f64); (p, truth(p)) }).collect();
        let q = Quad::fit(&prs).unwrap();
        for p in [(100.0, 2000.0), (2500.0, 300.0)] { assert!(geom::dist(q.apply(p), truth(p)) < 1e-6); }
    }
}