Merge pull request #2351 from vangel76/point-two-stage-search

Point tracker: refine in a 3x3 neighbourhood when the motion was small

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<isize>,
+    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<Taps> {
+        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<f64> {
+        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<Taps>,
+}
+
+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<Plane>,
-    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;
@@ -298,6 +399,12 @@ impl Crop {
     }
 }
 
+/// Largest distance (in top-level pixels, per axis) between the best coarse match and the
+/// motion prediction for which the fine stage searches only ±1 pixel per level.
+const SMALL_MOTION: f64 = 1.0;
+/// Lowest coarse NCC score of the best match for which the fine stage may search only ±1.
+const RELIABLE_SCORE: f64 = 0.8;
+
 /// Per-point tracking state.
 #[derive(Clone, Debug)]
 struct PointState {
@@ -372,13 +479,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();
@@ -396,7 +505,10 @@ impl PointState {
                 break;
             }
         }
-        // Fine: refine each seed down the pyramid in a 5 × 5 neighbourhood.
+        // Fine: refine each seed down the pyramid. A good motion prediction plus a reliable coarse
+        // match only needs a 3 × 3 neighbourhood per level; otherwise keep 5 × 5.
+        let near = cands[0].0 >= RELIABLE_SCORE && seeds.first().is_some_and(|c| (c[0] - cp[0]).abs() <= SMALL_MOTION && (c[1] - cp[1]).abs() <= SMALL_MOTION);
+        let rr: i32 = if near { 1 } else { 2 };
         let mut best = (f64::MIN, sc);
         for seed in seeds {
             let mut c = seed;
@@ -405,14 +517,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);
-                for j in -2..=2 {
-                    for i in -2..=2 {
+                let scorer = Scorer::new(pl, tp, c, shape);
+                let mut lb = (scorer.at(0, 0), c);
+                for j in -rr..=rr {
+                    for i in -rr..=rr {
                         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 +544,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 +620,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 +700,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;

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 });
+                }
+            }
+        }
+    }
 }

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| {
@@ -114,6 +163,27 @@ fn recovers_translation_path_subpixel() {
     assert!(worst < 0.25, "worst error {worst} px");
 }
 
+#[test]
+fn recovers_after_sudden_jumps_the_prediction_misses() {
+    // Alternating large steps defeat the velocity prediction, so the fine stage must keep its
+    // wider neighbourhood; smooth frames in between use the narrow one.
+    let path = |f: usize| {
+        let t = f as f64;
+        let jump = [0.0, 7.0, -5.0, 9.0, -8.0, 6.0, -7.0, 8.0];
+        let j = jump[f % jump.len()];
+        Pose { c: [120.0 + 1.5 * t + j, 100.0 + 1.1 * t - 0.6 * j], rot: 0.0, scale: 1.0 }
+    };
+    let f0 = frame(360, 260, path(0), 40.0, 0.02, 0, None);
+    let mut tr = Tracker::new(TrackOptions::default(), &[spec(path(0).c, 31.0, 91.0)], &Frame::new(&f0));
+    for f in 1..24 {
+        let img = frame(360, 260, path(f), 40.0, 0.02, f as u32, None);
+        let r = tr.step(&Frame::new(&img))[0];
+        let e = dist(r.center, path(f).c);
+        assert!(e < 0.3, "frame {f}: error {e} px");
+        assert_eq!(r.status, Status::Tracked);
+    }
+}
+
 #[test]
 fn layer_offset_and_channels() {
     let pose = |f: usize| Pose { c: [120.0 + 3.3 * f as f64, 100.0 - 1.7 * f as f64], rot: 0.0, scale: 1.0 };