1use image::GrayImage;
2
3use crate::pyramid::build_pyramid_into;
4use crate::utils::fast_gradients::compute_gradients_into;
5
6pub const DEFAULT_MIN_EIGEN_THRESHOLD: f32 = 1e-3;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum TrackStatus {
19 Tracked,
21 OutOfBounds,
23 Diverged,
26 LowTexture,
29 FbInconsistent,
33}
34
35pub const DEFAULT_FB_THRESHOLD: f32 = 0.7;
38
39#[derive(Debug, Clone, Copy, PartialEq)]
47pub struct TrackResult {
48 pub pos: (f32, f32),
50 pub status: TrackStatus,
52 pub error: f32,
59}
60
61#[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
100pub 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#[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#[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 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 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 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 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 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 let x = *prev_x / scale;
264 let y = *prev_y / scale;
265
266 let mut dx = displacements[idx].0 / scale;
268 let mut dy = displacements[idx].1 / scale;
269
270 if !in_bounds(prev_img, x, y, radius) {
272 out[idx].status = TrackStatus::OutOfBounds;
273 continue;
274 }
275
276 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 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 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 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 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 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#[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 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
453fn 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 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
477fn 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
485fn 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
508fn 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
531fn 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
537fn 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 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 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 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 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
596fn 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 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 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#[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 pub fn new() -> Self {
672 Self::default()
673 }
674
675 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 pub fn prev_pyramid(&self) -> &[GrayImage] {
684 &self.prev_pyramid
685 }
686
687 pub fn next_pyramid(&self) -> &[GrayImage] {
689 &self.next_pyramid
690 }
691
692 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 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 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}