diff --git a/crates/merge/src/align.rs b/crates/merge/src/align.rs index 34a606155..be103c4e5 100644 --- a/crates/merge/src/align.rs +++ b/crates/merge/src/align.rs @@ -5,7 +5,6 @@ use lightcraft_geom::{Homography, Point}; use lightcraft_raster::resample::{Filter, fit}; use lightcraft_raster::{Plane, Rgb32f}; -use rayon::prelude::*; use crate::features::{self, Features}; use crate::ransac::{self, Model}; @@ -129,23 +128,10 @@ pub fn align_pair(ref_feat: &Features, feat: &Features, scale: f64, full_w: usiz /// Resample `img` into the reference geometry: `out(p) = img(h(p))` (Lanczos-3). Returns the warped /// image and a validity mask (false where `h(p)` falls outside `img`). pub fn warp(img: &Rgb32f, h: &Homography) -> (Rgb32f, Vec) { - let (w, ht) = (img.width, img.height); if *h == Homography::IDENTITY { - return (img.clone(), vec![true; w * ht]); + return (img.clone(), vec![true; img.width * img.height]); } - let mut out = Rgb32f::new(w, ht); - let mut valid = vec![false; w * ht]; - out.data.par_chunks_mut(w).zip(valid.par_chunks_mut(w)).enumerate().for_each(|(y, (row, vrow))| { - for x in 0..w { - let p = h.apply(Point::new(x as f64 + 0.5, y as f64 + 0.5)); - let (sx, sy) = (p.x as f32, p.y as f32); - if sx >= 0.0 && sy >= 0.0 && sx <= w as f32 && sy <= ht as f32 { - row[x] = img.sample_lanczos3(sx, sy); - vrow[x] = true; - } - } - }); - (out, valid) + crate::flow::warp(img, h, &crate::flow::Flow::default()) } #[cfg(test)] diff --git a/crates/merge/src/flow.rs b/crates/merge/src/flow.rs new file mode 100644 index 000000000..a8a9a6bf9 --- /dev/null +++ b/crates/merge/src/flow.rs @@ -0,0 +1,239 @@ +//! Local alignment: what one motion model for the whole frame leaves. +//! +//! A homography is exact only for a camera that turned about its lens. Handheld, the camera also +//! moves, and things at different distances shift differently (parallax): on a real bracket (Sony, +//! 33 MP, gravel a few metres away and a building a kilometre off) the fit, dominated by the +//! foreground's texture, left the building 2.3 px off, and the merge doubled its inscription. Lens +//! distortion and rolling shutter leave local shifts too. +//! +//! After the global fit ([`crate::align`], [`crate::refine`]) the frame is compared with the +//! reference tile by tile: +//! +//! 1. **Measure.** On a grid of about [`GRID`] tiles along the long edge, the reference's patch at +//! the tile's centre is found in the frame, through the global fit, within [`RADIUS`] pixels +//! ([`crate::refine::match_patch`]: masked correlation, sub-pixel peak, ambiguous and textureless +//! patches dropped). +//! 2. **Clean.** A shift more than [`OUTLIER`] px from the median of its measured neighbours is +//! dropped (a person walking, a mismatch), as is one with no measured neighbour to check it by. Tiles without a measurement take the mean of their +//! measured neighbours, spreading outwards until every tile has one. +//! 3. **Smooth** with two 3 × 3 box passes, so the field bends gently. +//! 4. **Warp.** Each output pixel `p` samples the frame at `h(p + flow(p))` (the flow interpolated +//! bilinearly between tile centres): one Lanczos-3 resampling for both corrections. + +use lightcraft_geom::{Homography, Point}; +use lightcraft_raster::{Plane, Rgb32f}; +use rayon::prelude::*; + +use crate::refine::{self, PATCH}; + +/// Tiles along the long edge. +pub const GRID: usize = 96; +/// Largest local shift measured (pixels); farther ones are filled from neighbours. +pub const RADIUS: usize = 5; +/// A tile this far (pixels) from the median of its neighbours is an outlier. +pub const OUTLIER: f32 = 0.75; + +/// A smooth field of local shifts in reference pixels: one value per tile, tile `(i, j)`'s at +/// `((i + ½)·step, (j + ½)·step)`, so pixel centres of `dx` / `dy` line up with the tiles'. +#[derive(Clone, Debug, PartialEq)] +pub struct Flow { + pub step: usize, + pub dx: Plane, + pub dy: Plane, + /// Tiles that were measured (before filling). + pub measured: usize, +} + +impl Default for Flow { + /// No shift anywhere. + fn default() -> Self { + Flow { step: 0, dx: Plane::new(0, 0), dy: Plane::new(0, 0), measured: 0 } + } +} + +impl Flow { + /// The shift at reference coordinates (pixel centres at +0.5), bilinear between tile centres; + /// zero for an empty flow. + pub fn at(&self, x: f64, y: f64) -> (f64, f64) { + let (fx, fy) = ((x / self.step as f64) as f32, (y / self.step as f64) as f32); + if self.dx.data.is_empty() || !fx.is_finite() || !fy.is_finite() { + return (0.0, 0.0); + } + (crate::features::bilinear(&self.dx, fx, fy) as f64, crate::features::bilinear(&self.dy, fx, fy) as f64) + } + + /// The largest shift in the field (pixels). + pub fn max(&self) -> f32 { + self.dx.data.iter().zip(&self.dy.data).map(|(x, y)| x.hypot(*y)).fold(0.0, f32::max) + } +} + +/// The local shifts between `reference` and `frame` (both from [`refine::luma`]) once `h` (reference +/// → frame pixels) is applied, within `radius` pixels. +pub fn estimate(reference: &Plane, frame: &Plane, h: &Homography, radius: usize) -> Flow { + let (w, ht) = (reference.width, reference.height); + // an even step, so a patch centred on a tile has a whole-pixel corner + let step = ((w.max(ht) / GRID).max(PATCH) + 1) & !1; + let (gw, gh) = (w / step, ht / step); + if gw == 0 || gh == 0 || frame.data.is_empty() { + return Flow::default(); + } + let r = radius.max(1); + let raw: Vec> = (0..gw * gh) + .into_par_iter() + .map(|t| { + let (i, j) = (t % gw, t / gw); + let (sx, sy) = (i * step + (step - PATCH) / 2, j * step + (step - PATCH) / 2); + if sx < r + 1 || sy < r + 1 || sx + PATCH + r + 1 > w || sy + PATCH + r + 1 > ht { + return None; + } + refine::match_patch(reference, (sx, sy), frame, h, r).map(|m| [m.offset.x as f32, m.offset.y as f32]) + }) + .collect(); + let measured = raw.iter().filter(|v| v.is_some()).count(); + let at = |d: &[Option<[f32; 2]>], k: usize| d.get(k).copied().flatten(); + // 2. outliers against the median of their measured neighbours + let mut d: Vec> = (0..gw * gh) + .map(|t| { + let v = at(&raw, t)?; + let n = neighbours(|k| at(&raw, k), gw, gh, t); + let med = [median(n.iter().map(|v| v[0]).collect()), median(n.iter().map(|v| v[1]).collect())]; + // a tile with no measured neighbour can't be checked, and the fill would spread it: dropped + (!n.is_empty() && (v[0] - med[0]).hypot(v[1] - med[1]) <= OUTLIER).then_some(v) + }) + .collect(); + // fill outwards from measured tiles + for _ in 0..gw + gh { + if d.iter().all(Option::is_some) || d.iter().all(Option::is_none) { + break; + } + let prev = d.clone(); + for (t, slot) in d.iter_mut().enumerate().filter(|(_, s)| s.is_none()) { + *slot = mean(&neighbours(|k| at(&prev, k), gw, gh, t)); + } + } + let mut f: Vec<[f32; 2]> = d.into_iter().map(|v| v.unwrap_or([0.0; 2])).collect(); + // 3. smooth (each tile with its neighbours) + for _ in 0..2 { + let prev = f.clone(); + for (t, v) in f.iter_mut().enumerate() { + let mut n = neighbours(|k| prev.get(k).copied(), gw, gh, t); + n.push(*v); + *v = mean(&n).unwrap_or(*v); + } + } + let plane = |c: usize| Plane::from_fn(gw, gh, |i, j| f.get(j * gw + i).map_or(0.0, |v| v[c])); + Flow { step, dx: plane(0), dy: plane(1), measured } +} + +/// The values `get` has among the 8 neighbours of tile `t`. +fn neighbours(get: impl Fn(usize) -> Option<[f32; 2]>, gw: usize, gh: usize, t: usize) -> Vec<[f32; 2]> { + let (i, j) = (t % gw, t / gw); + let mut out = Vec::with_capacity(8); + for b in j.saturating_sub(1)..=(j + 1).min(gh - 1) { + for a in i.saturating_sub(1)..=(i + 1).min(gw - 1) { + if (a, b) != (i, j) + && let Some(v) = get(b * gw + a) + { + out.push(v); + } + } + } + out +} + +fn mean(n: &[[f32; 2]]) -> Option<[f32; 2]> { + let k = n.len() as f32; + (!n.is_empty()).then(|| [n.iter().map(|v| v[0]).sum::() / k, n.iter().map(|v| v[1]).sum::() / k]) +} + +fn median(mut v: Vec) -> f32 { + v.sort_by(|a, b| a.total_cmp(b)); + v.get(v.len() / 2).copied().unwrap_or(0.0) +} + +/// `img` resampled into the reference's geometry through `h` after the local shifts: +/// `out(p) = img(h(p + flow(p)))` (Lanczos-3), and where that falls inside `img`. +pub fn warp(img: &Rgb32f, h: &Homography, flow: &Flow) -> (Rgb32f, Vec) { + let (w, ht) = (img.width, img.height); + let mut out = Rgb32f::new(w, ht); + let mut valid = vec![false; w * ht]; + if w == 0 || ht == 0 { + return (out, valid); + } + out.data.par_chunks_mut(w).zip(valid.par_chunks_mut(w)).enumerate().for_each(|(y, (row, vrow))| { + for x in 0..w { + let (px, py) = (x as f64 + 0.5, y as f64 + 0.5); + let (dx, dy) = flow.at(px, py); + let q = h.apply(Point::new(px + dx, py + dy)); + let (sx, sy) = (q.x as f32, q.y as f32); + if sx >= 0.0 && sy >= 0.0 && sx <= w as f32 && sy <= ht as f32 { + row[x] = img.sample_lanczos3(sx, sy); + vrow[x] = true; + } + } + }); + (out, valid) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::refine::tests::tex; + + /// Parallax: the top of the frame (far) shifted 2.2 px against the bottom (near), with a + /// smooth band between. The flow finds both, and the warp lines the frame up everywhere. + #[test] + fn finds_and_undoes_a_local_shift() { + let (w, h) = (1200usize, 800usize); + let shift = |y: f64| { + let t = ((y - 300.0) / 200.0).clamp(0.0, 1.0); + let s = t * t * (3.0 - 2.0 * t); + ((1.0 - s) * 2.2, (1.0 - s) * -1.1) + }; + let reference = Plane::from_fn(w, h, |x, y| tex(x as f64 + 0.5, y as f64 + 0.5)); + // the frame's pixel q shows the scene at q + shift(q.y) + let frame = Rgb32f::from_fn(w, h, |x, y| { + let (sx, sy) = shift(y as f64 + 0.5); + [tex(x as f64 + 0.5 + sx, y as f64 + 0.5 + sy); 3] + }); + let fl = refine::luma(&frame, 1.0, 10.0); + let flow = estimate(&reference, &fl, &Homography::IDENTITY, RADIUS); + let tiles = flow.dx.width * flow.dx.height; + assert!(flow.measured > tiles / 2, "{} of {tiles}", flow.measured); + // reference p matches frame p + d with d = −shift + let (tx, ty) = flow.at(600.0, 100.0); + assert!((tx + 2.2).abs() < 0.15 && (ty - 1.1).abs() < 0.15, "top {tx} {ty}"); + let (bx, by) = flow.at(600.0, 700.0); + assert!(bx.abs() < 0.15 && by.abs() < 0.15, "bottom {bx} {by}"); + let (out, valid) = warp(&frame, &Homography::IDENTITY, &flow); + let error = |img: &Rgb32f| { + let (mut e, mut n) = (0f64, 0usize); + for y in (40..h - 40).step_by(3) { + for x in (40..w - 40).step_by(3) { + if valid[y * w + x] { + e += (img.get(x, y)[0] - reference.get(x, y)).abs() as f64; + n += 1; + } + } + } + e / n as f64 + }; + let (before, after) = (error(&frame), error(&out)); + assert!(after < before / 6.0, "mean error {before} unwarped, {after} warped"); + } + + #[test] + fn empty_and_flat_inputs_give_no_shift() { + let flat = Plane::from_fn(400, 300, |_, _| 0.5); + let f = estimate(&flat, &flat, &Homography::IDENTITY, RADIUS); + assert_eq!(f.measured, 0); + assert_eq!(f.max(), 0.0); + assert_eq!(estimate(&Plane::new(0, 0), &Plane::new(0, 0), &Homography::IDENTITY, 3), Flow::default()); + assert_eq!(estimate(&flat, &Plane::new(0, 0), &Homography::IDENTITY, 3), Flow::default()); + assert_eq!(Flow::default().at(5.0, f64::NAN), (0.0, 0.0)); + assert_eq!(f.at(f64::INFINITY, 3.0), (0.0, 0.0)); + let (img, v) = warp(&Rgb32f::new(0, 0), &Homography::IDENTITY, &Flow::default()); + assert!(img.data.is_empty() && v.is_empty()); + } +} diff --git a/crates/merge/src/lib.rs b/crates/merge/src/lib.rs index 4b4e206e6..b5ccd4bb8 100644 --- a/crates/merge/src/lib.rs +++ b/crates/merge/src/lib.rs @@ -15,12 +15,14 @@ pub mod align; pub mod features; +pub mod flow; pub mod frame; pub mod hdr; pub mod linalg; pub mod output; pub mod pano; pub mod ransac; +pub mod refine; pub mod synth; pub use frame::{Frame, FrameColor, RawColour, load_frame, load_frame_with}; diff --git a/crates/merge/src/refine.rs b/crates/merge/src/refine.rs new file mode 100644 index 000000000..e9044bc8b --- /dev/null +++ b/crates/merge/src/refine.rs @@ -0,0 +1,359 @@ +//! Full-resolution refinement of an alignment found on small versions of the frames. +//! +//! [`crate::align`] registers frames on versions of at most [`crate::align::WORK_EDGE`] pixels. +//! For a 24–60 MP photo that is 5–8× smaller, and a tenth of a work pixel off is most of a full +//! pixel off: merged frames that far apart blur the result. Here the motion is fitted again at +//! full size: +//! +//! 1. **Sites.** The reference is cut into a grid of cells (≈ [`CELLS`] along the long edge); in +//! each, the [`PATCH`]-pixel square with the most two-dimensional texture (the smaller +//! eigenvalue of its gradient structure tensor, so a straight edge, which only fixes one +//! direction, doesn't count) is a site. +//! 2. **Matching.** Each site's patch is compared with the frame resampled through the current +//! alignment, by normalised cross-correlation (immune to the frames' different brightness and +//! contrast) over integer offsets within a search radius, counting only pixels usable in both +//! (clipped ones aren't); the peak is refined to a fraction of a pixel by a parabola through its +//! neighbours on each axis. Patches less than half usable, with too little texture, a weak peak, a peak on the search border or a second peak almost +//! as strong (repeating texture) are dropped. +//! 3. **Fit.** The matches are fitted with RANSAC at [`THRESH`] pixels, choosing the simplest of +//! translation / similarity / homography that explains (nearly) as many as the most general +//! one, as [`crate::align::align_pair`] does. A second, narrower pass starts from that fit. +//! +//! When too few patches match, or the fit isn't a plausible handheld motion, the alignment found +//! on the small versions stays: refining never makes an alignment worse than not refining. + +use lightcraft_geom::{Homography, Point}; +use lightcraft_raster::{Plane, Rgb32f}; +use rayon::prelude::*; + +use crate::align::{Alignment, Models, choose_model, luma as lum}; +use crate::ransac::Model; + +/// Side of a matched patch (pixels). +pub const PATCH: usize = 48; +/// Cells along the long edge (sites are picked one per cell). +pub const CELLS: usize = 24; +/// RANSAC inlier threshold at full resolution (pixels). +pub const THRESH: f64 = 0.75; +/// Weakest correlation peak a match may have. +const MIN_SCORE: f32 = 0.8; +/// The best peak away from the best one must be at least this much weaker. +const UNIQUE: f32 = 0.03; +/// Least standard deviation (of the gamma-encoded luminance) a patch needs to be matched. +const MIN_STD: f32 = 0.01; +/// Fewest matches a fit is made from. +const MIN_MATCHES: usize = 12; +/// Share of the matches the fit must explain. +const MIN_SHARE: f64 = 0.4; + +/// A frame's luminance for matching: `(luminance · gain)^(1/2.2)` at full size, NaN where the +/// frame is clipped (any channel at or above 95 % of `clip`) or `gain` takes it past 1, so such +/// pixels never take part. `gain` brings frames of a bracket to the same brightness. +pub fn luma(img: &Rgb32f, gain: f32, clip: f32) -> Plane { + let g = if gain.is_finite() && gain > 0.0 { gain } else { 1.0 }; + img.map(|c| { + if c[0].max(c[1]).max(c[2]) >= 0.95 * clip { + return f32::NAN; + } + let y = lum(c) * g; + if y.is_finite() && y < 1.0 { y.max(0.0).powf(1.0 / 2.2) } else { f32::NAN } + }) +} + +/// The search radius (full-resolution pixels) for an alignment found at `scale` (work pixels per +/// full pixel): about two work pixels. +pub fn radius_for(scale: f64) -> usize { + if scale.is_finite() && scale > 0.0 { ((2.0 / scale).ceil() as usize + 1).clamp(3, 16) } else { 3 } +} + +/// A patch of the reference found in a frame. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Match { + /// Centre of the patch in the reference (pixel centres at +0.5). + pub reference: Point, + /// Where it is in the frame. + pub frame: Point, + /// The offset from where the alignment put it, in reference pixels. + pub offset: Point, + /// Correlation at the peak (−1..1). + pub score: f32, +} + +/// Top-left corners of the patches to match: one per cell, where the reference has the most +/// two-dimensional texture, `margin` pixels clear of the edges. +pub fn sites(reference: &Plane, margin: usize) -> Vec<(usize, usize)> { + let (w, h) = (reference.width, reference.height); + let need = PATCH + 2 * margin; + if w < need || h < need { + return Vec::new(); + } + let long = w.max(h); + let cell = long.div_ceil(CELLS).max(PATCH); + let (nx, ny) = ((w - 2 * margin) / cell, (h - 2 * margin) / cell); + let (nx, ny) = (nx.max(1), ny.max(1)); + let (cw, chh) = ((w - 2 * margin) / nx, (h - 2 * margin) / ny); + (0..nx * ny) + .into_par_iter() + .filter_map(|c| { + let (cx, cy) = (margin + (c % nx) * cw, margin + (c / nx) * chh); + let span_x = cw.saturating_sub(PATCH); + let span_y = chh.saturating_sub(PATCH); + let mut best: Option<((usize, usize), f32)> = None; + for j in 0..3 { + for i in 0..3 { + let (x, y) = (cx + span_x * i / 2, cy + span_y * j / 2); + if x + PATCH > w - margin || y + PATCH > h - margin { + continue; + } + if let Some(t) = texture(reference, x, y) + && best.is_none_or(|(_, b)| t > b) + { + best = Some(((x, y), t)); + } + } + } + best.filter(|(_, t)| *t > 1e-6).map(|(p, _)| p) + }) + .collect() +} + +/// The smaller eigenvalue of the patch's mean gradient structure tensor, over its usable pixels; +/// `None` when fewer than half of them are. +fn texture(p: &Plane, x0: usize, y0: usize) -> Option { + let w = p.width; + let at = |x: usize, y: usize| p.data.get(y * w + x).copied().unwrap_or(f32::NAN); + let (mut sxx, mut syy, mut sxy, mut n) = (0f64, 0f64, 0f64, 0usize); + for y in y0 + 1..y0 + PATCH - 1 { + for x in x0 + 1..x0 + PATCH - 1 { + let gx = (at(x + 1, y) - at(x - 1, y)) as f64; + let gy = (at(x, y + 1) - at(x, y - 1)) as f64; + if gx.is_finite() && gy.is_finite() { + sxx += gx * gx; + syy += gy * gy; + sxy += gx * gy; + n += 1; + } + } + } + if n * 2 < (PATCH - 2) * (PATCH - 2) { + return None; + } + let n = n as f64; + let (a, b, c) = (sxx / n, syy / n, sxy / n); + let tr = a + b; + let det = a * b - c * c; + Some((tr / 2.0 - ((tr * tr / 4.0 - det).max(0.0)).sqrt()) as f32) +} + +/// Find the reference patch at `site` (top-left) in `frame`, which `h` maps reference pixels into, +/// searching offsets up to `radius` pixels. +pub fn match_patch(reference: &Plane, site: (usize, usize), frame: &Plane, h: &Homography, radius: usize) -> Option { + let (sx, sy) = site; + let (rw, rh) = (reference.width, reference.height); + if sx + PATCH > rw || sy + PATCH > rh || frame.width == 0 || frame.height == 0 { + return None; + } + let r = radius.max(1); + let s = PATCH + 2 * r; + let mut pr: Vec = Vec::with_capacity(PATCH * PATCH); + for y in sy..sy + PATCH { + for x in sx..sx + PATCH { + pr.push(reference.data.get(y * rw + x).copied().unwrap_or(f32::NAN)); + } + } + // the frame around the patch, through `h`; NaN where it can't be used + let (fw, fh) = (frame.width as f64, frame.height as f64); + let mut win = vec![f32::NAN; s * s]; + for j in 0..s { + for i in 0..s { + let q = h.apply(Point::new((sx + i) as f64 - r as f64 + 0.5, (sy + j) as f64 - r as f64 + 0.5)); + if !(q.x >= 0.5 && q.y >= 0.5 && q.x <= fw - 0.5 && q.y <= fh - 0.5) { + return None; + } + win[j * s + i] = crate::features::bilinear(frame, q.x as f32, q.y as f32); + } + } + let d = 2 * r + 1; + // offsets that can't be scored (too few usable pixels) stay at −∞ + let mut score = vec![f32::NEG_INFINITY; d * d]; + for dy in 0..d { + for dx in 0..d { + if let Some(c) = correlation(&pr, |u, v| win.get((v + dy) * s + dx + u).copied().unwrap_or(f32::NAN)) { + score[dy * d + dx] = c; + } + } + } + let (bi, &best) = score.iter().enumerate().max_by(|a, b| a.1.total_cmp(b.1))?; + let (bx, by) = (bi % d, bi / d); + if best < MIN_SCORE || bx == 0 || by == 0 || bx == d - 1 || by == d - 1 { + return None; + } + // a second peak nearly as strong: repeating texture, the match is ambiguous + let second = + score.iter().enumerate().filter(|(k, _)| (k % d).abs_diff(bx) > 1 || (k / d).abs_diff(by) > 1).map(|(_, v)| *v).fold(-1f32, f32::max); + if second > best - UNIQUE { + return None; + } + let at = |x: usize, y: usize| score[y * d + x]; + // the sub-pixel parabola needs the peak's four neighbours: one that couldn't be scored would pull it away + if [at(bx - 1, by), at(bx + 1, by), at(bx, by - 1), at(bx, by + 1)].iter().any(|v| !v.is_finite()) { + return None; + } + let sub = |m: f32, c: f32, p: f32| { + let den = m - 2.0 * c + p; + if den < -1e-9 { (0.5 * (m - p) / den).clamp(-0.5, 0.5) } else { 0.0 } + }; + let ox = bx as f64 - r as f64 + sub(at(bx - 1, by), best, at(bx + 1, by)) as f64; + let oy = by as f64 - r as f64 + sub(at(bx, by - 1), best, at(bx, by + 1)) as f64; + let c = Point::new(sx as f64 + PATCH as f64 / 2.0, sy as f64 + PATCH as f64 / 2.0); + Some(Match { reference: c, frame: h.apply(Point::new(c.x + ox, c.y + oy)), offset: Point::new(ox, oy), score: best }) +} + +/// Correlation of the reference patch `a` (row-major, [`PATCH`] wide) with `b(u, v)` over the +/// pixels usable in both; `None` with fewer than half of the patch usable, or too little texture in +/// either. +fn correlation(a: &[f32], b: impl Fn(usize, usize) -> f32) -> Option { + let (mut n, mut sa, mut sb, mut saa, mut sbb, mut sab) = (0usize, 0f64, 0f64, 0f64, 0f64, 0f64); + for v in 0..PATCH { + for u in 0..PATCH { + let (x, y) = (a.get(v * PATCH + u).copied().unwrap_or(f32::NAN), b(u, v)); + if x.is_finite() && y.is_finite() { + let (x, y) = (x as f64, y as f64); + n += 1; + sa += x; + sb += y; + saa += x * x; + sbb += y * y; + sab += x * y; + } + } + } + if n * 2 < PATCH * PATCH { + return None; + } + let nf = n as f64; + let (ma, mb) = (sa / nf, sb / nf); + let (va, vb) = (saa / nf - ma * ma, sbb / nf - mb * mb); + let min = (MIN_STD as f64).powi(2); + if !(va >= min && vb >= min * 0.25) { + return None; + } + let c = (sab / nf - ma * mb) / (va * vb).sqrt(); + c.is_finite().then_some(c as f32) +} + +/// Translation, similarity and homography: what a handheld frame may need. +pub const ALL_MODELS: [(Model, &str, usize); 3] = + [(Model::Translation, "translation", 6), (Model::Similarity, "similarity", 8), (Model::Homography, "homography", 12)]; + +/// Search radius (pixels) for a frame with no alignment to start from (the small versions found +/// none: a bright frame mostly clipped, too little texture left for features). +pub const RECOVER_RADIUS: usize = 16; + +/// Fit `models` (reference → frame, full resolution) to patches of `reference` matched in `frame` +/// around `start` within `radius` pixels, then again within 2 px of that fit; both planes from +/// [`luma`]. `start` is an alignment found on small versions (`radius` from [`radius_for`]), or the +/// identity with [`RECOVER_RADIUS`]. `None` when too few patches match or no fit is a plausible +/// handheld motion: the caller keeps what it had. +pub fn fit(reference: &Plane, frame: &Plane, start: &Homography, radius: usize, seed: u64, models: &Models) -> Option { + let (w, ht) = (reference.width as f64, reference.height as f64); + let r0 = radius.max(2); + // the sites for the widest search fit the narrower one too + let s = sites(reference, r0 + 2); + let mut cur: Option = None; + for r in [r0, 2] { + let h = cur.as_ref().map_or(*start, |a| a.h); + let m: Vec = s.par_iter().filter_map(|&p| match_patch(reference, p, frame, &h, r)).collect(); + if m.len() < MIN_MATCHES { + break; + } + let pa: Vec = m.iter().map(|m| m.reference).collect(); + let pb: Vec = m.iter().map(|m| m.frame).collect(); + let Some((name, f)) = choose_model(models, &pa, &pb, THRESH, MIN_SHARE, (w, ht), seed) else { break }; + cur = Some(Alignment { h: f.h, model: name, inliers: f.inliers.len(), rms: f.rms }); + } + cur +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + + /// A smooth random texture, sampled anywhere. + pub(crate) fn tex(x: f64, y: f64) -> f32 { + let mut v = 0.45; + for k in 0..9 { + let a = k as f64 * 2.399; + let f = 0.02 + 0.025 * k as f64; + v += 0.05 * ((x * a.cos() + y * a.sin()) * f * std::f64::consts::TAU + k as f64).sin(); + } + v as f32 + } + + fn plane(w: usize, h: usize, map: &Homography) -> Plane { + Plane::from_fn(w, h, |x, y| { + let p = map.apply(Point::new(x as f64 + 0.5, y as f64 + 0.5)); + tex(p.x, p.y) + }) + } + + fn corner_error(a: &Homography, b: &Homography, w: f64, h: f64) -> f64 { + [(0.0, 0.0), (w, 0.0), (0.0, h), (w, h)].iter().map(|&(x, y)| a.apply(Point::new(x, y)).dist(b.apply(Point::new(x, y)))).fold(0.0, f64::max) + } + + /// A frame rotated 0.2° and moved 1.37 px, given an alignment 2 px off: refined to a tenth of + /// a pixel at the corners. + #[test] + fn refines_a_rough_alignment_to_a_fraction_of_a_pixel() { + let (w, h) = (900usize, 600usize); + let reference = plane(w, h, &Homography::IDENTITY); + // the frame's pixel q shows the scene at s(q); reference → frame is s⁻¹ + let (c, sn) = (0.2f64.to_radians().cos(), 0.2f64.to_radians().sin()); + let s = Homography([c, -sn, 1.37, sn, c, -0.83, 0.0, 0.0, 1.0]); + let truth = s.inverse().unwrap(); + let frame = plane(w, h, &s); + let rough = Homography([truth.0[0], truth.0[1], truth.0[2] + 1.6, truth.0[3], truth.0[4], truth.0[5] - 1.2, 0.0, 0.0, 1.0]); + assert!(corner_error(&rough, &truth, w as f64, h as f64) > 1.5); + let r = fit(&reference, &frame, &rough, 4, 1, &ALL_MODELS).unwrap(); + let e = corner_error(&r.h, &truth, w as f64, h as f64); + assert!(e < 0.1, "corner error {e} px, model {}", r.model); + assert!(r.inliers >= MIN_MATCHES); + } + + /// A frame the small versions couldn't align, 5 px and 0.1° off: found from nothing. + #[test] + fn recovers_a_frame_with_no_alignment() { + let (w, h) = (900usize, 600usize); + let reference = plane(w, h, &Homography::IDENTITY); + let (c, sn) = (0.1f64.to_radians().cos(), 0.1f64.to_radians().sin()); + let s = Homography([c, -sn, -4.6, sn, c, 2.1, 0.0, 0.0, 1.0]); + let r = fit(&reference, &plane(w, h, &s), &Homography::IDENTITY, RECOVER_RADIUS, 1, &ALL_MODELS).unwrap(); + let e = corner_error(&r.h, &s.inverse().unwrap(), w as f64, h as f64); + assert!(e < 0.1, "corner error {e} px, model {}", r.model); + assert!(fit(&reference, &Plane::from_fn(w, h, |_, _| 0.5), &Homography::IDENTITY, RECOVER_RADIUS, 1, &ALL_MODELS).is_none()); + } + + /// Nothing to match (a flat frame, planes too small for a patch, empty ones): no fit. + #[test] + fn no_fit_when_nothing_matches() { + let h = Homography([1.0, 0.0, 0.7, 0.0, 1.0, 0.2, 0.0, 0.0, 1.0]); + let flat = Plane::from_fn(400, 300, |_, _| 0.5); + assert!(fit(&flat, &flat, &h, 4, 1, &ALL_MODELS).is_none()); + let tiny = Plane::from_fn(20, 20, |x, _| x as f32); + assert!(fit(&tiny, &tiny, &h, 4, 1, &ALL_MODELS).is_none()); + assert!(fit(&Plane::new(0, 0), &Plane::new(0, 0), &h, 4, 1, &ALL_MODELS).is_none()); + } + + #[test] + fn clipped_pixels_and_odd_gains_are_left_out() { + let img = Rgb32f::from_fn(4, 1, |x, _| [x as f32 * 0.33; 3]); + let l = luma(&img, 1.0, 1.0); + assert!(l.data[0] == 0.0 && l.data[3].is_nan(), "{:?}", l.data); + let l = luma(&img, 4.0, 1.0); + assert!(l.data[1].is_nan(), "0.33 × 4 is past 1: {:?}", l.data); + assert!(luma(&img, f32::NAN, 1.0).data[1].is_finite()); + assert_eq!(radius_for(0.125), 16); + assert_eq!(radius_for(1.0), 3); + assert_eq!(radius_for(f64::NAN), 3); + } +}