diff --git a/crates/track/src/lib.rs b/crates/track/src/lib.rs index fe36bc096..898019726 100644 --- a/crates/track/src/lib.rs +++ b/crates/track/src/lib.rs @@ -201,7 +201,7 @@ fn ncc(plane: &Plane, p: &Patch, c: [f64; 2], a: [f64; 2]) -> f64 { } let ch = p.ch; let (mut s, mut s2, mut st) = (0.0f64, 0.0f64, 0.0f64); - let identity = (a[0] - 1.0).abs() < 1e-9 && a[1].abs() < 1e-9; + let identity = is_identity(a); let fx = c[0] - 0.5; let fy = c[1] - 0.5; let (bx, by) = (fx.floor(), fy.floor()); @@ -256,12 +256,114 @@ fn ncc(plane: &Plane, p: &Patch, c: [f64; 2], a: [f64; 2]) -> f64 { (st / (var * p.ss).sqrt()).clamp(-1.0, 1.0) } +fn is_identity(a: [f64; 2]) -> bool { + (a[0] - 1.0).abs() < 1e-9 && a[1].abs() < 1e-9 +} + +/// The bilinear taps of every template pixel at `c + A u`. Candidates `c + (i, j)` with integer +/// `(i, j)` give each pixel the same fractional offsets, so one table serves them all and a +/// candidate is integer-indexed lerps (the arithmetic of `Plane::sample`) with one bounds check. +struct Taps { + /// Offset (in samples) of the top-left tap per template pixel from the one at `base`. + off: Vec, + w: Vec<[f32; 2]>, + base: isize, + /// Tap bounding box (x0, y0, x1, y1) in plane pixels at the reference, `+1` taps included. + bb: [isize; 4], +} + +impl Taps { + fn new(plane: &Plane, p: &Patch, c: [f64; 2], a: [f64; 2]) -> Option { + let n = ((2 * p.r[0] + 1) * (2 * p.r[1] + 1)) as usize; + let (mut off, mut w) = (Vec::with_capacity(n), Vec::with_capacity(n)); + let (rx, ry) = ((c[0] - 0.5).floor(), (c[1] - 0.5).floor()); + if !(rx.abs() < 1e9 && ry.abs() < 1e9) { + return None; + } + let (rx, ry) = (rx as isize, ry as isize); + let mut bb = [isize::MAX, isize::MAX, isize::MIN, isize::MIN]; + for gy in -p.r[1]..=p.r[1] { + for gx in -p.r[0]..=p.r[0] { + let (ux, uy) = (gx as f64, gy as f64); + let fx = c[0] + a[0] * ux - a[1] * uy - 0.5; + let fy = c[1] + a[1] * ux + a[0] * uy - 0.5; + let (x0, y0) = (fx.floor(), fy.floor()); + if !(x0.abs() < 1e9 && y0.abs() < 1e9) { + return None; + } + let (dx, dy) = (x0 as isize - rx, y0 as isize - ry); + bb = [bb[0].min(dx), bb[1].min(dy), bb[2].max(dx + 1), bb[3].max(dy + 1)]; + off.push((dy * plane.w as isize + dx) * plane.ch as isize); + w.push([(fx - x0) as f32, (fy - y0) as f32]); + } + } + let base = (ry * plane.w as isize + rx) * plane.ch as isize; + Some(Taps { off, w, base, bb: [bb[0] + rx, bb[1] + ry, bb[2] + rx, bb[3] + ry] }) + } + + /// The correlation at `c + (i, j)`, or `None` when a tap leaves the plane. + fn ncc(&self, plane: &Plane, p: &Patch, i: isize, j: isize) -> Option { + if self.bb[0] + i < 0 || self.bb[1] + j < 0 || self.bb[2] + i >= plane.w as isize || self.bb[3] + j >= plane.h as isize { + return None; + } + let ch = p.ch; + let row = plane.w * ch; + let b = self.base + (j * plane.w as isize + i) * ch as isize; + let d = &plane.data; + let (mut s, mut s2, mut st) = (0.0f64, 0.0f64, 0.0f64); + let mut k = 0; + for (o, t) in self.off.iter().zip(&self.w) { + let i00 = (b + o) as usize; + let (i10, i01) = (i00 + ch, i00 + row); + for cc in 0..ch { + let top = d[i00 + cc] + (d[i10 + cc] - d[i00 + cc]) * t[0]; + let bot = d[i01 + cc] + (d[i01 + ch + cc] - d[i01 + cc]) * t[0]; + let x = (top + (bot - top) * t[1]) as f64; + s += x; + s2 += x * x; + st += x * p.zm[k] as f64; + k += 1; + } + } + let n = p.zm.len() as f64; + let var = s2 - s * s / n; + if var <= 1e-12 { + return Some(0.0); + } + Some((st / (var * p.ss).sqrt()).clamp(-1.0, 1.0)) + } +} + +/// `ncc` of one template at integer offsets from a reference centre, through a [`Taps`] table +/// when the shape is not the identity (the identity has its own fast path in `ncc`). +struct Scorer<'a> { + plane: &'a Plane, + p: &'a Patch, + c: [f64; 2], + a: [f64; 2], + taps: Option, +} + +impl<'a> Scorer<'a> { + fn new(plane: &'a Plane, p: &'a Patch, c: [f64; 2], a: [f64; 2]) -> Scorer<'a> { + let taps = if is_identity(a) || p.ss <= 1e-12 { None } else { Taps::new(plane, p, c, a) }; + Scorer { plane, p, c, a, taps } + } + + /// The correlation at the reference centre plus `(i, j)` pixels. + fn at(&self, i: i32, j: i32) -> f64 { + self.taps + .as_ref() + .and_then(|t| t.ncc(self.plane, self.p, i as isize, j as isize)) + .unwrap_or_else(|| ncc(self.plane, self.p, [self.c[0] + i as f64, self.c[1] + j as f64], self.a)) + } +} + /// The pre-processed pixels around one point: a pyramid of a crop. Plane coordinates are image /// pixels minus `origin`, divided by `2^level`. struct Crop { origin: [f64; 2], levels: Vec, - grad0: (Plane, Plane), } impl Crop { @@ -280,13 +382,12 @@ impl Crop { if opts.enhance { p = p.enhance(); } - let grad0 = p.gradients(); let mut lv = vec![p]; for _ in 0..levels { let d = lv.last().map(Plane::downsample).unwrap_or_default(); lv.push(d); } - Crop { origin: [x0, y0], levels: lv, grad0 } + Crop { origin: [x0, y0], levels: lv } } fn to_level(&self, p: [f64; 2], k: usize) -> [f64; 2] { let s = (1u32 << k) as f64; @@ -372,13 +473,15 @@ impl PointState { let (rx, ry) = ((range[0] / s).ceil() as i32, (range[1] / s).ceil() as i32); let plane = &crop.levels[top]; let tpl = &self.templates[top]; + let scorer = Scorer::new(plane, tpl, ctop, self.shape); let shape = self.shape; let mut cands: Vec<(f64, [f64; 2])> = (-ry..=ry) .into_par_iter() .flat_map_iter(|j| { + let scorer = &scorer; (-rx..=rx).map(move |i| { let c = [ctop[0] + i as f64, ctop[1] + j as f64]; - (ncc(plane, tpl, c, shape), c) + (scorer.at(i, j), c) }) }) .collect(); @@ -405,14 +508,15 @@ impl PointState { c = [c[0] * 2.0, c[1] * 2.0]; let (pl, tp) = (&crop.levels[k], &self.templates[k]); // Start from the projected match; neighbours must beat it (no drift on ties). - let mut lb = (ncc(pl, tp, c, shape), c); + let scorer = Scorer::new(pl, tp, c, shape); + let mut lb = (scorer.at(0, 0), c); for j in -2..=2 { for i in -2..=2 { if i == 0 && j == 0 { continue; } let q = [c[0] + i as f64, c[1] + j as f64]; - let v = ncc(pl, tp, q, shape); + let v = scorer.at(i, j); if v > lb.0 + 1e-9 { lb = (v, q); } @@ -431,15 +535,16 @@ impl PointState { // Sub-pixel (and shape) refinement. let (mut c, mut a) = (c0, shape); if opts.subpixel { + let mut grad = Grad::around(&crop.levels[0], c0, shape, self.templates[0].r); let mut ok = false; if opts.track_shape - && let Some((cr, ar)) = lk_refine(&crop, &self.templates[0], c0, shape, true) + && let Some((cr, ar)) = lk_refine(&crop, &mut grad, &self.templates[0], c0, shape, true) { (c, a) = (cr, ar); ok = true; } if !ok { - if let Some((cr, _)) = lk_refine(&crop, &self.templates[0], c0, shape, false) { + if let Some((cr, _)) = lk_refine(&crop, &mut grad, &self.templates[0], c0, shape, false) { c = cr; } else { c = parabolic(&crop.levels[0], &self.templates[0], c0, shape); @@ -506,12 +611,52 @@ fn parabolic(plane: &Plane, p: &Patch, c: [f64; 2], a: [f64; 2]) -> [f64; 2] { [c[0] + off(f(-1.0, 0.0), f(1.0, 0.0)), c[1] + off(f(0.0, -1.0), f(0.0, 1.0))] } +/// Level-0 gradients of a crop, computed only where `lk_refine` reads them: valid on the pixel +/// rectangle `valid` (x0, y0, x1, y1, inclusive) and zero elsewhere. +struct Grad { + gx: Plane, + gy: Plane, + valid: [isize; 4], +} + +impl Grad { + /// Extra pixels around the first window of taps, for the drift of the refinement. A window that + /// still runs out is replaced by the full gradients. + const MARGIN: isize = 6; + + fn full(plane: &Plane) -> Grad { + let (gx, gy) = plane.gradients(); + Grad { gx, gy, valid: [isize::MIN, isize::MIN, isize::MAX, isize::MAX] } + } + + fn around(plane: &Plane, c: [f64; 2], a: [f64; 2], r: [i32; 2]) -> Grad { + let [x0, y0, x1, y1] = Grad::taps(c, a, r); + let m = Grad::MARGIN; + let clip = |v: isize, hi: usize| v.clamp(0, (hi as isize - 1).max(0)) as usize; + let rect = [clip(x0 - m, plane.w), clip(y0 - m, plane.h), clip(x1 + m, plane.w), clip(y1 + m, plane.h)]; + let (gx, gy) = plane.gradients_in(rect); + Grad { gx, gy, valid: rect.map(|v| v as isize) } + } + + /// The pixels read by sampling the template window at `c` with shape `a` (bilinear taps). + fn taps(c: [f64; 2], a: [f64; 2], r: [i32; 2]) -> [isize; 4] { + let ex = (a[0].abs() * r[0] as f64 + a[1].abs() * r[1] as f64) * (1.0 + 1e-9) + 1e-9; + let ey = (a[1].abs() * r[0] as f64 + a[0].abs() * r[1] as f64) * (1.0 + 1e-9) + 1e-9; + let px = |v: f64| v.clamp(-1e9, 1e9).floor() as isize; + [px(c[0] - ex - 0.5), px(c[1] - ey - 0.5), px(c[0] + ex - 0.5) + 1, px(c[1] + ey - 0.5) + 1] + } + + fn covers(&self, c: [f64; 2], a: [f64; 2], r: [i32; 2]) -> bool { + let t = Grad::taps(c, a, r); + t[0] >= self.valid[0] && t[1] >= self.valid[1] && t[2] <= self.valid[2] && t[3] <= self.valid[3] + } +} + /// Forward-additive Lucas–Kanade (Gauss–Newton) alignment of the level-0 template with gain and /// bias: minimise Σ (I(c + A u) − g T(u) − b)² over the centre `c` and, with `shape`, the /// similarity `A`. Returns `None` when it diverges. -fn lk_refine(crop: &Crop, p: &Patch, c0: [f64; 2], a0: [f64; 2], shape: bool) -> Option<([f64; 2], [f64; 2])> { +fn lk_refine(crop: &Crop, grad: &mut Grad, p: &Patch, c0: [f64; 2], a0: [f64; 2], shape: bool) -> Option<([f64; 2], [f64; 2])> { let plane = &crop.levels[0]; - let (gxp, gyp) = &crop.grad0; let ch = p.ch; let (mut c, mut a) = (c0, a0); let mut g = 1.0f64; @@ -546,6 +691,10 @@ fn lk_refine(crop: &Crop, p: &Patch, c0: [f64; 2], a0: [f64; 2], shape: bool) -> let np = if shape { 6 } else { 4 }; let (mut v, mut vx, mut vy) = ([0.0f32; 3], [0.0f32; 3], [0.0f32; 3]); for _ in 0..30 { + if !grad.covers(c, a, p.r) { + *grad = Grad::full(plane); + } + let (gxp, gyp) = (&grad.gx, &grad.gy); let mut h = [[0.0f64; 6]; 6]; let mut rhs = [0.0f64; 6]; let mut k = 0; diff --git a/crates/track/src/plane.rs b/crates/track/src/plane.rs index bd2d987d6..ff15d08b2 100644 --- a/crates/track/src/plane.rs +++ b/crates/track/src/plane.rs @@ -191,15 +191,32 @@ impl Plane { /// Central-difference gradients `(d/dx, d/dy)`. pub fn gradients(&self) -> (Plane, Plane) { + self.gradients_in([0, 0, self.w.saturating_sub(1), self.h.saturating_sub(1)]) + } + + /// The gradients of the pixels in `[x0, y0, x1, y1]` (inclusive, within the plane); all other + /// pixels of the returned planes stay zero. + pub fn gradients_in(&self, [x0, y0, x1, y1]: [usize; 4]) -> (Plane, Plane) { let mut gx = Plane::new(self.w, self.h, self.ch); let mut gy = Plane::new(self.w, self.h, self.ch); - for y in 0..self.h { - for x in 0..self.w { - for c in 0..self.ch { - let (xi, yi) = (x as isize, y as isize); - let i = self.idx(x, y) + c; - gx.data[i] = 0.5 * (self.at(xi + 1, yi, c) - self.at(xi - 1, yi, c)); - gy.data[i] = 0.5 * (self.at(xi, yi + 1, c) - self.at(xi, yi - 1, c)); + if self.w == 0 || self.h == 0 { + return (gx, gy); + } + let (ch, row) = (self.ch, self.w * self.ch); + for y in y0..=y1.min(self.h - 1) { + for x in x0..=x1.min(self.w - 1) { + let i = self.idx(x, y); + if x > 0 && y > 0 && x + 1 < self.w && y + 1 < self.h { + for c in i..i + ch { + gx.data[c] = 0.5 * (self.data[c + ch] - self.data[c - ch]); + gy.data[c] = 0.5 * (self.data[c + row] - self.data[c - row]); + } + } else { + for c in 0..ch { + let (xi, yi) = (x as isize, y as isize); + gx.data[i + c] = 0.5 * (self.at(xi + 1, yi, c) - self.at(xi - 1, yi, c)); + gy.data[i + c] = 0.5 * (self.at(xi, yi + 1, c) - self.at(xi, yi - 1, c)); + } } } } @@ -240,4 +257,27 @@ mod tests { assert_eq!(Plane::from_image(&img, Channel::Saturation, 0, 0, 2, 2).data[0], 1.0); assert_eq!(Channel::from_name("luminance"), Some(Channel::Luminance)); } + + #[test] + fn gradient_window_and_interior_path_match_the_clamped_definition() { + let mut p = Plane::new(11, 9, 3); + for (i, v) in p.data.iter_mut().enumerate() { + *v = ((i * 7919) % 101) as f32 / 37.0; + } + let (gx, gy) = p.gradients(); + let rect = [2, 1, 8, 6]; + let (wx, wy) = p.gradients_in(rect); + for y in 0..p.h { + for x in 0..p.w { + for c in 0..3 { + let (xi, yi, i) = (x as isize, y as isize, p.idx(x, y) + c); + assert_eq!(gx.data[i].to_bits(), (0.5 * (p.at(xi + 1, yi, c) - p.at(xi - 1, yi, c))).to_bits()); + assert_eq!(gy.data[i].to_bits(), (0.5 * (p.at(xi, yi + 1, c) - p.at(xi, yi - 1, c))).to_bits()); + let inside = (rect[0]..=rect[2]).contains(&x) && (rect[1]..=rect[3]).contains(&y); + assert_eq!(wx.data[i], if inside { gx.data[i] } else { 0.0 }); + assert_eq!(wy.data[i], if inside { gy.data[i] } else { 0.0 }); + } + } + } + } } diff --git a/crates/track/src/tests.rs b/crates/track/src/tests.rs index 886478d1e..e97452ddc 100644 --- a/crates/track/src/tests.rs +++ b/crates/track/src/tests.rs @@ -93,6 +93,55 @@ fn dist(a: [f64; 2], b: [f64; 2]) -> f64 { ((a[0] - b[0]).powi(2) + (a[1] - b[1]).powi(2)).sqrt() } +fn bits(v: &Option<([f64; 2], [f64; 2])>) -> Option<[u64; 4]> { + v.map(|(c, a)| [c[0].to_bits(), c[1].to_bits(), a[0].to_bits(), a[1].to_bits()]) +} + +#[test] +fn scorer_matches_ncc_for_integer_offsets() { + let pose = Pose { c: [60.0, 50.0], rot: 0.0, scale: 1.0 }; + let img = frame(120, 100, pose, 20.0, 0.02, 1, None); + let crop = Crop::new(&Frame::new(&img), &TrackOptions::default(), [60.0, 50.0], [40.0, 30.0], 1); + for (lvl, c, a) in [(0, [30.3, 25.9], [0.97, 0.21]), (0, [12.0, 9.5], [1.2, -0.3]), (1, [20.4, 14.2], [0.8, 0.0]), (0, [30.3, 25.9], [1.0, 0.0])] { + let plane = &crop.levels[lvl]; + let r = [7, 5]; + let patch = Patch::sample(plane, c, a, r); + let sc = Scorer::new(plane, &patch, [c[0] + 0.3, c[1] - 0.45], a); + for j in -9..=9 { + for i in -9..=9 { + let q = [c[0] + 0.3 + i as f64, c[1] - 0.45 + j as f64]; + assert_eq!(sc.at(i, j).to_bits(), ncc(plane, &patch, q, a).to_bits(), "level {lvl} {a:?} ({i}, {j})"); + } + } + } +} + +#[test] +fn lk_refine_with_a_gradient_window_equals_the_full_gradients() { + let pose = Pose { c: [100.0, 90.0], rot: 5.0, scale: 1.05 }; + let (img0, img1) = (frame(200, 180, pose, 30.0, 0.02, 1, None), frame(200, 180, Pose { c: [103.2, 88.1], rot: 6.5, scale: 1.08 }, 30.0, 0.02, 2, None)); + let (opts, r) = (TrackOptions::default(), [24, 24]); + let (c0, a0) = ([100.0, 90.0], [1.0, 0.0]); + let tpl = Crop::new(&Frame::new(&img0), &opts, c0, [60.0, 60.0], 0); + let crop = Crop::new(&Frame::new(&img1), &opts, c0, [60.0, 60.0], 0); + let plane = &crop.levels[0]; + let patch = Patch::sample(&tpl.levels[0], tpl.to_level(c0, 0), a0, r); + let mid = crop.to_level(c0, 0); + // Starts within the window, and far enough off that the iterations leave it (even near the edge). + for (d, edge) in [([0.0, 0.0], false), ([1.3, -0.8], false), ([4.0, 5.0], false), ([9.0, -9.0], false), ([0.0, 0.0], true)] { + let start = if edge { [4.0, 4.0] } else { [mid[0] + d[0], mid[1] + d[1]] }; + for shape in [true, false] { + let (mut full, mut win) = (Grad::full(plane), Grad::around(plane, start, a0, r)); + let x = lk_refine(&crop, &mut full, &patch, start, a0, shape); + let y = lk_refine(&crop, &mut win, &patch, start, a0, shape); + assert_eq!(bits(&x), bits(&y), "{d:?} edge {edge} shape {shape}"); + // A window that misses the taps altogether is replaced by the full gradients. + let mut off = Grad::around(plane, [start[0] + 15.0, start[1] + 15.0], a0, r); + assert_eq!(bits(&x), bits(&lk_refine(&crop, &mut off, &patch, start, a0, shape)), "{d:?} edge {edge} shape {shape}: missed window"); + } + } +} + #[test] fn recovers_translation_path_subpixel() { let path = |f: usize| {