Skip to main content

optical_flow_lk/
lk.rs

1use image::GrayImage;
2
3use crate::pyramid::build_pyramid_into;
4use crate::utils::fast_gradients::compute_gradients_into;
5
6/// Default minimum-eigenvalue threshold used by [`calc_optical_flow`].
7///
8/// The minimum eigenvalue of the per-window spatial gradient matrix is
9/// normalized by the window area before comparison, so this threshold is
10/// independent of `window_size`. A window flatter than this is reported as
11/// [`TrackStatus::LowTexture`].
12pub const DEFAULT_MIN_EIGEN_THRESHOLD: f32 = 1e-3;
13
14/// Why a feature point ended up where it did after tracking.
15///
16/// See [`TrackResult`] for the coordinate convention.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum TrackStatus {
19    /// The iteration converged inside the image; the position is trustworthy.
20    Tracked,
21    /// The search window left the image (in the previous or the next frame).
22    OutOfBounds,
23    /// The iteration hit `max_iterations` without converging, or a step
24    /// exploded (non-finite or larger than the window).
25    Diverged,
26    /// The minimum eigenvalue of the spatial gradient matrix fell below the
27    /// configured threshold, i.e. the window is too flat to track reliably.
28    LowTexture,
29    /// The point was tracked forward, but re-tracking it back to the previous
30    /// frame did not return close enough to the original position. Only
31    /// produced by [`calc_optical_flow_fb`]. A strong occlusion/outlier signal.
32    FbInconsistent,
33}
34
35/// Default forward-backward round-trip threshold (pixels) for
36/// [`calc_optical_flow_fb`].
37pub const DEFAULT_FB_THRESHOLD: f32 = 0.7;
38
39/// Per-point result of [`calc_optical_flow_ex`].
40///
41/// # Coordinate convention
42/// Positions are in level-0 (full resolution) pixel coordinates. The origin is
43/// the center of the top-left pixel, x grows to the right and y downwards; this
44/// matches [`crate::good_features_to_track`] and bilinear sampling throughout
45/// the crate.
46#[derive(Debug, Clone, Copy, PartialEq)]
47pub struct TrackResult {
48    /// Tracked position in the next frame.
49    pub pos: (f32, f32),
50    /// Why tracking ended the way it did.
51    pub status: TrackStatus,
52    /// Mean absolute photometric residual over the window at the final
53    /// position, in 8-bit intensity units.
54    ///
55    /// Comparable across points regardless of `window_size`; intended for
56    /// downstream outlier gating. It is [`f32::INFINITY`] when no residual
57    /// could be measured (the window was out of bounds).
58    pub error: f32,
59}
60
61/// Compute optical flow using the pyramidal Lucas-Kanade method.
62///
63/// This is a thin wrapper over [`calc_optical_flow_ex`] that discards the
64/// per-point diagnostics.
65///
66/// # Arguments
67/// * `prev_pyramid` - Previous frame (pyramid of grayscale)
68/// * `curr_pyramid` - Next frame (pyramid of grayscale)
69/// * `prev_points` - Feature points to track (in prev frame)
70/// * `window_size` - Size of the search window (odd number)
71/// * `max_iterations` - Max iterations for correct points on each layer
72///
73/// # Returns
74/// Vector of points on next frame
75#[deprecated(
76    since = "0.3.0",
77    note = "use `calc_optical_flow_ex`, which also returns per-point status and error"
78)]
79pub fn calc_optical_flow(
80    prev_pyramid: &[GrayImage],
81    curr_pyramid: &[GrayImage],
82    prev_points: &[(f32, f32)],
83    window_size: usize,
84    max_iterations: usize,
85) -> Vec<(f32, f32)> {
86    calc_optical_flow_ex(
87        prev_pyramid,
88        curr_pyramid,
89        prev_points,
90        None,
91        window_size,
92        max_iterations,
93        DEFAULT_MIN_EIGEN_THRESHOLD,
94    )
95    .into_iter()
96    .map(|r| r.pos)
97    .collect()
98}
99
100/// Compute optical flow using the pyramidal Lucas-Kanade method, returning a
101/// per-point [`TrackResult`] with status and photometric error.
102///
103/// # Arguments
104/// * `prev_pyramid` - Previous frame (pyramid of grayscale)
105/// * `curr_pyramid` - Next frame (pyramid of grayscale)
106/// * `prev_points` - Feature points to track (in prev frame, level-0 coordinates)
107/// * `predicted` - Optional predicted positions in the next frame, one per
108///   `prev_point` (level-0 coordinates). When `Some`, the iteration is seeded
109///   with the predicted displacement at the coarsest pyramid level instead of
110///   zero, which markedly improves convergence under large inter-frame motion
111///   (e.g. gyroscope-predicted feature motion). `None` reproduces the classic
112///   zero-initialized behavior. Per-level propagation is unchanged.
113/// * `window_size` - Size of the search window (odd number)
114/// * `max_iterations` - Max iterations per pyramid level
115/// * `min_eigen_threshold` - Windows whose normalized minimum gradient
116///   eigenvalue is below this value are reported as [`TrackStatus::LowTexture`].
117///   See [`DEFAULT_MIN_EIGEN_THRESHOLD`].
118///
119/// # Panics
120/// Panics if `predicted` is `Some` but its length differs from `prev_points`.
121///
122/// # Returns
123/// One [`TrackResult`] per input point, in the same order.
124pub fn calc_optical_flow_ex(
125    prev_pyramid: &[GrayImage],
126    curr_pyramid: &[GrayImage],
127    prev_points: &[(f32, f32)],
128    predicted: Option<&[(f32, f32)]>,
129    window_size: usize,
130    max_iterations: usize,
131    min_eigen_threshold: f32,
132) -> Vec<TrackResult> {
133    let mut scratch = Scratch::default();
134    let mut out = Vec::new();
135    track_into(
136        prev_pyramid,
137        curr_pyramid,
138        prev_points,
139        predicted,
140        window_size,
141        max_iterations,
142        min_eigen_threshold,
143        &mut scratch,
144        &mut out,
145    );
146    out
147}
148
149/// Reusable per-call scratch buffers for the Lucas-Kanade loop. Owned by
150/// [`TrackerContext`] (or created transiently by the free functions) so the
151/// steady-state hot path performs no heap allocation.
152#[derive(Default)]
153struct Scratch {
154    offsets: Vec<(f32, f32)>,
155    prev_patch: Vec<f32>,
156    ix_patch: Vec<f32>,
157    iy_patch: Vec<f32>,
158    displacements: Vec<(f32, f32)>,
159    grad_x: Vec<i16>,
160    grad_y: Vec<i16>,
161}
162
163/// Core pyramidal Lucas-Kanade loop, writing one [`TrackResult`] per point into
164/// `out`. All temporaries live in `scratch`; given sufficient capacity this is
165/// allocation-free.
166#[allow(clippy::too_many_arguments)]
167fn track_into(
168    prev_pyramid: &[GrayImage],
169    curr_pyramid: &[GrayImage],
170    prev_points: &[(f32, f32)],
171    predicted: Option<&[(f32, f32)]>,
172    window_size: usize,
173    max_iterations: usize,
174    min_eigen_threshold: f32,
175    scratch: &mut Scratch,
176    out: &mut Vec<TrackResult>,
177) {
178    assert_eq!(prev_pyramid.len(), curr_pyramid.len());
179    assert!(
180        !prev_pyramid.is_empty(),
181        "pyramid must have at least 1 level"
182    );
183    assert!(window_size % 2 == 1, "Window size must be odd");
184    if let Some(predicted) = predicted {
185        assert_eq!(
186            predicted.len(),
187            prev_points.len(),
188            "predicted must have one entry per prev_point"
189        );
190    }
191
192    let n_levels = prev_pyramid.len();
193    let radius = window_size / 2;
194    let n_pixels = window_size * window_size;
195    let epsilon = 1e-3;
196    let det_epsilon = 1e-6;
197
198    let Scratch {
199        offsets,
200        prev_patch,
201        ix_patch,
202        iy_patch,
203        displacements,
204        grad_x: grad_x_buf,
205        grad_y: grad_y_buf,
206    } = scratch;
207
208    // Prepare reusable buffers. resize/clear+extend keep capacity, so none of
209    // this allocates once the buffers are warm.
210    build_window_offsets_into(radius, offsets);
211    prev_patch.resize(n_pixels, 0.0);
212    ix_patch.resize(n_pixels, 0.0);
213    iy_patch.resize(n_pixels, 0.0);
214
215    // Total displacement per point, accumulated coarse-to-fine in level-0 units.
216    // Seeding it from a prediction makes the coarsest level start at the
217    // predicted position; everything else is identical to the zero-init path.
218    displacements.clear();
219    match predicted {
220        Some(predicted) => displacements.extend(
221            prev_points
222                .iter()
223                .zip(predicted.iter())
224                .map(|((px, py), (gx, gy))| (gx - px, gy - py)),
225        ),
226        None => displacements.resize(prev_points.len(), (0.0, 0.0)),
227    }
228
229    // One shared gradient buffer sized to the largest (level-0) image; smaller
230    // levels use a prefix slice.
231    let (w0, h0) = prev_pyramid[0].dimensions();
232    grad_x_buf.resize((w0 * h0) as usize, 0);
233    grad_y_buf.resize((w0 * h0) as usize, 0);
234
235    // Initialize results at the input positions; the loop refines them in place.
236    out.clear();
237    out.extend(prev_points.iter().map(|&(x, y)| TrackResult {
238        pos: (x, y),
239        status: TrackStatus::Tracked,
240        error: f32::INFINITY,
241    }));
242
243    // Process levels from top (coarse) to bottom (fine).
244    for level in (0..n_levels).rev() {
245        let scale = 2f32.powi(level as i32);
246        let is_finest = level == 0;
247
248        let prev_img = &prev_pyramid[level];
249        let curr_img = &curr_pyramid[level];
250        let (lw, lh) = prev_img.dimensions();
251        let level_pixels = (lw * lh) as usize;
252
253        compute_gradients_into(
254            prev_img,
255            &mut grad_x_buf[..level_pixels],
256            &mut grad_y_buf[..level_pixels],
257        );
258        let grad_x = &grad_x_buf[..level_pixels];
259        let grad_y = &grad_y_buf[..level_pixels];
260
261        for (idx, (prev_x, prev_y)) in prev_points.iter().enumerate() {
262            // Scale the original point for the current level.
263            let x = *prev_x / scale;
264            let y = *prev_y / scale;
265
266            // Add the current displacement, scaled for this level.
267            let mut dx = displacements[idx].0 / scale;
268            let mut dy = displacements[idx].1 / scale;
269
270            // The window must stay inside the previous image to build the patch.
271            if !in_bounds(prev_img, x, y, radius) {
272                out[idx].status = TrackStatus::OutOfBounds;
273                continue;
274            }
275
276            // Spatial gradient matrix and cached previous/gradient patches.
277            let mut gxx = 0.0f32;
278            let mut gxy = 0.0f32;
279            let mut gyy = 0.0f32;
280
281            for (i, (ox, oy)) in offsets.iter().enumerate() {
282                let sample_x = x + ox;
283                let sample_y = y + oy;
284                let ix = interpolate_i16(grad_x, lw, lh, sample_x, sample_y) / 32.0;
285                let iy = interpolate_i16(grad_y, lw, lh, sample_x, sample_y) / 32.0;
286
287                prev_patch[i] = interpolate(prev_img, sample_x, sample_y);
288                ix_patch[i] = ix;
289                iy_patch[i] = iy;
290                gxx += ix * ix;
291                gxy += ix * iy;
292                gyy += iy * iy;
293            }
294
295            // Reject low-texture windows up front (normalized by window area so
296            // the threshold does not depend on `window_size`).
297            let min_eig = min_eigenvalue(gxx, gxy, gyy) / n_pixels as f32;
298            if min_eig < min_eigen_threshold {
299                out[idx].status = TrackStatus::LowTexture;
300                if is_finest {
301                    out[idx].error =
302                        window_error(curr_img, prev_patch, offsets, x + dx, y + dy, radius);
303                }
304                continue;
305            }
306
307            let Some((inv_h00, inv_h01, inv_h11)) = invert_2x2(gxx, gxy, gyy, det_epsilon) else {
308                out[idx].status = TrackStatus::LowTexture;
309                continue;
310            };
311
312            // Refine the displacement at the current level.
313            let mut converged = false;
314            let mut out_of_bounds = false;
315            let mut diverged = false;
316            for _ in 0..max_iterations {
317                let curr_x = x + dx;
318                let curr_y = y + dy;
319
320                if !in_bounds(curr_img, curr_x, curr_y, radius) {
321                    out_of_bounds = true;
322                    break;
323                }
324
325                let mut bx = 0.0f32;
326                let mut by = 0.0f32;
327
328                for (i, (ox, oy)) in offsets.iter().enumerate() {
329                    let curr = interpolate(curr_img, curr_x + ox, curr_y + oy);
330                    let error = prev_patch[i] - curr;
331                    bx += ix_patch[i] * error;
332                    by += iy_patch[i] * error;
333                }
334
335                let ddx = inv_h00 * bx + inv_h01 * by;
336                let ddy = inv_h01 * bx + inv_h11 * by;
337                dx += ddx;
338                dy += ddy;
339
340                // Guard against runaway steps.
341                if !dx.is_finite()
342                    || !dy.is_finite()
343                    || ddx.abs() > window_size as f32
344                    || ddy.abs() > window_size as f32
345                {
346                    diverged = true;
347                    break;
348                }
349
350                if ddx.abs() < epsilon && ddy.abs() < epsilon {
351                    converged = true;
352                    break;
353                }
354            }
355
356            out[idx].status = if out_of_bounds {
357                TrackStatus::OutOfBounds
358            } else if diverged || !converged {
359                TrackStatus::Diverged
360            } else {
361                TrackStatus::Tracked
362            };
363
364            // Update the total displacement with the current level scale.
365            displacements[idx] = (dx * scale, dy * scale);
366
367            if is_finest {
368                out[idx].error = if out_of_bounds {
369                    f32::INFINITY
370                } else {
371                    window_error(curr_img, prev_patch, offsets, x + dx, y + dy, radius)
372                };
373            }
374        }
375    }
376
377    // Fold accumulated displacements into the reported positions.
378    for (idx, (x, y)) in prev_points.iter().enumerate() {
379        let (dx, dy) = displacements[idx];
380        out[idx].pos = (x + dx, y + dy);
381    }
382}
383
384/// Track points prev->next, then re-track the results next->prev, and flag any
385/// point whose round-trip lands further than `fb_threshold` pixels from where it
386/// started as [`TrackStatus::FbInconsistent`].
387///
388/// This is the standard forward-backward consistency check and is the cheapest
389/// reliable way to reject occlusions and ambiguous matches before they poison a
390/// downstream pose solve. Points that already failed the forward pass keep their
391/// forward status (the round-trip is only evaluated for forward-tracked points).
392///
393/// Both passes reuse the supplied pyramids. The free-function form still
394/// allocates its result vectors internally; for a fully allocation-free
395/// steady-state call, use the equivalent [`TrackerContext`] method.
396///
397/// # Arguments
398/// * `prev_pyramid` / `next_pyramid` - frame pyramids (shared by both passes)
399/// * `prev_points` - points to track (level-0 coordinates)
400/// * `predicted` - optional initial guess for the forward pass, see
401///   [`calc_optical_flow_ex`]
402/// * `window_size`, `max_iterations`, `min_eigen_threshold` - as in
403///   [`calc_optical_flow_ex`]
404/// * `fb_threshold` - maximum allowed round-trip distance in pixels; see
405///   [`DEFAULT_FB_THRESHOLD`]
406#[allow(clippy::too_many_arguments)]
407pub fn calc_optical_flow_fb(
408    prev_pyramid: &[GrayImage],
409    next_pyramid: &[GrayImage],
410    prev_points: &[(f32, f32)],
411    predicted: Option<&[(f32, f32)]>,
412    window_size: usize,
413    max_iterations: usize,
414    min_eigen_threshold: f32,
415    fb_threshold: f32,
416) -> Vec<TrackResult> {
417    let mut scratch = Scratch::default();
418    let mut forward = Vec::new();
419    track_into(
420        prev_pyramid,
421        next_pyramid,
422        prev_points,
423        predicted,
424        window_size,
425        max_iterations,
426        min_eigen_threshold,
427        &mut scratch,
428        &mut forward,
429    );
430
431    let forward_pos: Vec<(f32, f32)> = forward.iter().map(|r| r.pos).collect();
432    let mut backward = Vec::new();
433    // Seed the backward pass at the original points (the round-trip is expected
434    // to return there). This keeps the check robust under large motion, as in
435    // OpenCV's OPTFLOW_USE_INITIAL_FLOW reverse check, without weakening it: a
436    // genuinely wrong forward match still fails to land back within threshold.
437    track_into(
438        next_pyramid,
439        prev_pyramid,
440        &forward_pos,
441        Some(prev_points),
442        window_size,
443        max_iterations,
444        min_eigen_threshold,
445        &mut scratch,
446        &mut backward,
447    );
448
449    mark_fb_inconsistent(&mut forward, &backward, prev_points, fb_threshold);
450    forward
451}
452
453/// Flag forward-tracked points whose backward round-trip exceeds the threshold.
454/// Factored out so the [`TrackerContext`] path can reuse it without allocating.
455fn mark_fb_inconsistent(
456    forward: &mut [TrackResult],
457    backward: &[TrackResult],
458    prev_points: &[(f32, f32)],
459    fb_threshold: f32,
460) {
461    let threshold_sq = fb_threshold * fb_threshold;
462    for (idx, result) in forward.iter_mut().enumerate() {
463        // Only forward-tracked points are eligible; others keep their status.
464        if result.status != TrackStatus::Tracked {
465            continue;
466        }
467
468        let back = &backward[idx];
469        let dx = back.pos.0 - prev_points[idx].0;
470        let dy = back.pos.1 - prev_points[idx].1;
471        if back.status != TrackStatus::Tracked || dx * dx + dy * dy > threshold_sq {
472            result.status = TrackStatus::FbInconsistent;
473        }
474    }
475}
476
477/// Minimum eigenvalue of the symmetric 2x2 matrix `[[a, b], [b, c]]`.
478fn min_eigenvalue(a: f32, b: f32, c: f32) -> f32 {
479    let trace = a + c;
480    let det = a * c - b * b;
481    let disc = (trace * trace - 4.0 * det).max(0.0).sqrt();
482    (trace - disc) / 2.0
483}
484
485/// Mean absolute photometric residual between the cached previous patch and the
486/// next image sampled at `(cx, cy)`. Returns [`f32::INFINITY`] if the window is
487/// out of bounds.
488fn window_error(
489    img: &GrayImage,
490    prev_patch: &[f32],
491    offsets: &[(f32, f32)],
492    cx: f32,
493    cy: f32,
494    radius: usize,
495) -> f32 {
496    if !in_bounds(img, cx, cy, radius) {
497        return f32::INFINITY;
498    }
499
500    let mut sum = 0.0f32;
501    for (i, (ox, oy)) in offsets.iter().enumerate() {
502        let curr = interpolate(img, cx + ox, cy + oy);
503        sum += (prev_patch[i] - curr).abs();
504    }
505    sum / offsets.len() as f32
506}
507
508/// Fills `offsets` with the `(dx, dy)` window sample positions for the given
509/// radius, reusing the existing capacity.
510fn build_window_offsets_into(radius: usize, offsets: &mut Vec<(f32, f32)>) {
511    offsets.clear();
512    offsets.reserve((2 * radius + 1) * (2 * radius + 1));
513
514    for j in -(radius as i32)..=radius as i32 {
515        for i in -(radius as i32)..=radius as i32 {
516            offsets.push((i as f32, j as f32));
517        }
518    }
519}
520
521fn invert_2x2(a00: f32, a01: f32, a11: f32, det_epsilon: f32) -> Option<(f32, f32, f32)> {
522    let det = a00 * a11 - a01 * a01;
523    if det.abs() <= det_epsilon {
524        return None;
525    }
526
527    let inv_det = 1.0 / det;
528    Some((a11 * inv_det, -a01 * inv_det, a00 * inv_det))
529}
530
531/// Checks that the window stays within image bounds
532fn in_bounds(img: &GrayImage, x: f32, y: f32, radius: usize) -> bool {
533    let (w, h) = (img.width() as f32, img.height() as f32);
534    x >= radius as f32 && x < w - radius as f32 && y >= radius as f32 && y < h - radius as f32
535}
536
537/// Bilinear interpolation of the pixel value
538fn interpolate(img: &GrayImage, x: f32, y: f32) -> f32 {
539    let w = img.width() as i32;
540    let h = img.height() as i32;
541    let x0 = x.floor() as i32;
542    let y0 = y.floor() as i32;
543
544    let dx = x - x0 as f32;
545    let dy = y - y0 as f32;
546
547    let data = img.as_raw();
548
549    // Fast path: the full 2x2 footprint is inside the image, so all four reads
550    // are provably in bounds. This is the overwhelmingly common case in the LK
551    // loop (the window is bounds-checked before sampling), and it avoids the
552    // per-sample `Option`/branch machinery of `get_pixel_checked` — the single
553    // hottest cost in tracking, and especially expensive under WASM where
554    // bounds checks are not elided.
555    if x0 >= 0 && y0 >= 0 && x0 + 1 < w && y0 + 1 < h {
556        let stride = w as usize;
557        let base = y0 as usize * stride + x0 as usize;
558        // SAFETY: x0>=0, y0>=0, x0+1<w, y0+1<h proven above, and the buffer is
559        // exactly w*h long, so base, base+1, base+stride, base+stride+1 are all
560        // valid indices.
561        let (p00, p10, p01, p11) = unsafe {
562            (
563                *data.get_unchecked(base) as f32,
564                *data.get_unchecked(base + 1) as f32,
565                *data.get_unchecked(base + stride) as f32,
566                *data.get_unchecked(base + stride + 1) as f32,
567            )
568        };
569        // Same term order as the slow path below for bit-identical results.
570        return p00 * (1.0 - dx) * (1.0 - dy)
571            + p01 * (1.0 - dx) * dy
572            + p10 * dx * (1.0 - dy)
573            + p11 * dx * dy;
574    }
575
576    // Slow path: replicate the zero-padded out-of-bounds behavior at the border.
577    let x1 = x0 + 1;
578    let y1 = y0 + 1;
579    let mut sum = 0.0;
580    for (sx, sy) in &[(x0, y0), (x0, y1), (x1, y0), (x1, y1)] {
581        let px = if *sx >= 0 && *sy >= 0 && *sx < w && *sy < h {
582            data[*sy as usize * w as usize + *sx as usize] as f32
583        } else {
584            0.0
585        };
586
587        let wx = if sx == &x0 { 1.0 - dx } else { dx };
588        let wy = if sy == &y0 { 1.0 - dy } else { dy };
589
590        sum += px * wx * wy;
591    }
592
593    sum
594}
595
596/// Bilinear interpolation over a raw `i16` gradient buffer (`width * height`,
597/// row-major). Out-of-bounds samples read as 0, matching [`interpolate`].
598fn interpolate_i16(data: &[i16], width: u32, height: u32, x: f32, y: f32) -> f32 {
599    let w = width as i32;
600    let h = height as i32;
601    let x0 = x.floor() as i32;
602    let y0 = y.floor() as i32;
603
604    let dx = x - x0 as f32;
605    let dy = y - y0 as f32;
606
607    // Fast path: full 2x2 footprint in bounds (see `interpolate`).
608    if x0 >= 0 && y0 >= 0 && x0 + 1 < w && y0 + 1 < h {
609        let stride = width as usize;
610        let base = y0 as usize * stride + x0 as usize;
611        // SAFETY: footprint proven in bounds and `data` is width*height long.
612        let (p00, p10, p01, p11) = unsafe {
613            (
614                *data.get_unchecked(base) as f32,
615                *data.get_unchecked(base + 1) as f32,
616                *data.get_unchecked(base + stride) as f32,
617                *data.get_unchecked(base + stride + 1) as f32,
618            )
619        };
620        return p00 * (1.0 - dx) * (1.0 - dy)
621            + p01 * (1.0 - dx) * dy
622            + p10 * dx * (1.0 - dy)
623            + p11 * dx * dy;
624    }
625
626    let x1 = x0 + 1;
627    let y1 = y0 + 1;
628    let mut sum = 0.0;
629    for (sx, sy) in &[(x0, y0), (x0, y1), (x1, y0), (x1, y1)] {
630        let px = if *sx >= 0 && *sy >= 0 && *sx < w && *sy < h {
631            data[*sy as usize * width as usize + *sx as usize] as f32
632        } else {
633            0.0
634        };
635
636        let wx = if sx == &x0 { 1.0 - dx } else { dx };
637        let wy = if sy == &y0 { 1.0 - dy } else { dy };
638
639        sum += px * wx * wy;
640    }
641
642    sum
643}
644
645/// Reusable owner of every buffer the tracking hot path touches: both frame
646/// pyramids, the Lucas-Kanade scratch, the result vector and the
647/// forward-backward intermediates.
648///
649/// Create one per tracking thread and call [`prepare`](Self::prepare) then
650/// [`track`](Self::track) / [`track_fb`](Self::track_fb) each frame. After the
651/// first (warm-up) frame, a steady-state step with a fixed image size, level
652/// count, window size and point count performs **no heap allocation** — all
653/// buffers are resized in place. This is the allocation-free path the VIO
654/// front-end runs every frame; the free functions
655/// ([`calc_optical_flow_ex`], [`calc_optical_flow_fb`]) are thin convenience
656/// wrappers that allocate their own scratch.
657#[derive(Default)]
658pub struct TrackerContext {
659    prev_pyramid: Vec<GrayImage>,
660    next_pyramid: Vec<GrayImage>,
661    scratch: Scratch,
662    results: Vec<TrackResult>,
663    forward_pos: Vec<(f32, f32)>,
664    backward: Vec<TrackResult>,
665}
666
667impl TrackerContext {
668    /// Creates an empty context. Buffers grow to fit on the first
669    /// [`prepare`](Self::prepare) / [`track`](Self::track) call and are reused
670    /// thereafter.
671    pub fn new() -> Self {
672        Self::default()
673    }
674
675    /// Builds the previous- and next-frame pyramids into the context's reusable
676    /// buffers. Zero-alloc in steady state (same image size and `levels`).
677    pub fn prepare(&mut self, prev: &GrayImage, next: &GrayImage, levels: usize) {
678        build_pyramid_into(prev, levels, &mut self.prev_pyramid);
679        build_pyramid_into(next, levels, &mut self.next_pyramid);
680    }
681
682    /// The previous-frame pyramid built by the last [`prepare`](Self::prepare).
683    pub fn prev_pyramid(&self) -> &[GrayImage] {
684        &self.prev_pyramid
685    }
686
687    /// The next-frame pyramid built by the last [`prepare`](Self::prepare).
688    pub fn next_pyramid(&self) -> &[GrayImage] {
689        &self.next_pyramid
690    }
691
692    /// Tracks `prev_points` using the prepared pyramids, returning the results
693    /// held inside the context. See [`calc_optical_flow_ex`] for the argument
694    /// semantics. Allocation-free in steady state.
695    pub fn track(
696        &mut self,
697        prev_points: &[(f32, f32)],
698        predicted: Option<&[(f32, f32)]>,
699        window_size: usize,
700        max_iterations: usize,
701        min_eigen_threshold: f32,
702    ) -> &[TrackResult] {
703        track_into(
704            &self.prev_pyramid,
705            &self.next_pyramid,
706            prev_points,
707            predicted,
708            window_size,
709            max_iterations,
710            min_eigen_threshold,
711            &mut self.scratch,
712            &mut self.results,
713        );
714        &self.results
715    }
716
717    /// Forward-backward consistent tracking using the prepared pyramids. See
718    /// [`calc_optical_flow_fb`] for semantics. Reuses the context's scratch and
719    /// intermediate point buffers, so it is allocation-free in steady state.
720    pub fn track_fb(
721        &mut self,
722        prev_points: &[(f32, f32)],
723        predicted: Option<&[(f32, f32)]>,
724        window_size: usize,
725        max_iterations: usize,
726        min_eigen_threshold: f32,
727        fb_threshold: f32,
728    ) -> &[TrackResult] {
729        track_into(
730            &self.prev_pyramid,
731            &self.next_pyramid,
732            prev_points,
733            predicted,
734            window_size,
735            max_iterations,
736            min_eigen_threshold,
737            &mut self.scratch,
738            &mut self.results,
739        );
740
741        self.forward_pos.clear();
742        self.forward_pos.extend(self.results.iter().map(|r| r.pos));
743
744        // Seed the backward pass at the original points (see calc_optical_flow_fb).
745        track_into(
746            &self.next_pyramid,
747            &self.prev_pyramid,
748            &self.forward_pos,
749            Some(prev_points),
750            window_size,
751            max_iterations,
752            min_eigen_threshold,
753            &mut self.scratch,
754            &mut self.backward,
755        );
756
757        mark_fb_inconsistent(&mut self.results, &self.backward, prev_points, fb_threshold);
758        &self.results
759    }
760}
761
762#[cfg(test)]
763mod tests {
764    use super::invert_2x2;
765
766    #[test]
767    fn invert_2x2_returns_inverse_components() {
768        let (inv00, inv01, inv11) = invert_2x2(4.0, 1.0, 3.0, 1e-6).unwrap();
769
770        assert!((inv00 - 3.0 / 11.0).abs() < 1e-6);
771        assert!((inv01 + 1.0 / 11.0).abs() < 1e-6);
772        assert!((inv11 - 4.0 / 11.0).abs() < 1e-6);
773    }
774
775    #[test]
776    fn invert_2x2_rejects_singular_matrix() {
777        assert!(invert_2x2(1.0, 2.0, 4.0, 1e-6).is_none());
778    }
779}