spatialrust_vision/
optical_flow.rs

1//! Sparse pyramidal Lucas–Kanade point tracking.
2
3use spatialrust_image::{Image, ImageView};
4use spatialrust_math::Vec2;
5
6use crate::{PixelComponent, VisionError, VisionResult};
7
8/// Configuration for sparse Lucas–Kanade tracking.
9#[derive(Clone, Copy, Debug, PartialEq)]
10pub struct LucasKanadeOptions {
11    /// Half-window radius in pixels at each pyramid level.
12    pub window_radius: usize,
13    /// Number of pyramid levels including the full-resolution level.
14    pub pyramid_levels: usize,
15    /// Maximum Gauss–Newton iterations per pyramid level.
16    pub max_iterations: usize,
17    /// Convergence threshold on the update length in pixels.
18    pub epsilon: f64,
19    /// Minimum structure-tensor determinant accepted as trackable.
20    pub min_eigenvalue: f64,
21}
22
23impl Default for LucasKanadeOptions {
24    fn default() -> Self {
25        Self {
26            window_radius: 5,
27            pyramid_levels: 3,
28            max_iterations: 30,
29            epsilon: 1e-3,
30            min_eigenvalue: 1e-4,
31        }
32    }
33}
34
35impl LucasKanadeOptions {
36    /// Validates window, pyramid, and numerical thresholds.
37    pub fn validate(self) -> VisionResult<Self> {
38        if self.window_radius == 0 || self.window_radius > isize::MAX as usize {
39            return Err(VisionError::InvalidParameter(
40                "Lucas–Kanade window_radius must be positive and fit isize".into(),
41            ));
42        }
43        if self.pyramid_levels == 0 || self.pyramid_levels > u32::BITS as usize {
44            return Err(VisionError::InvalidParameter(
45                "Lucas–Kanade pyramid_levels must be in 1..=32".into(),
46            ));
47        }
48        if self.max_iterations == 0 {
49            return Err(VisionError::InvalidParameter(
50                "Lucas–Kanade max_iterations must be positive".into(),
51            ));
52        }
53        if !self.epsilon.is_finite() || self.epsilon <= 0.0 {
54            return Err(VisionError::InvalidParameter(
55                "Lucas–Kanade epsilon must be finite and positive".into(),
56            ));
57        }
58        if !self.min_eigenvalue.is_finite() || self.min_eigenvalue <= 0.0 {
59            return Err(VisionError::InvalidParameter(
60                "Lucas–Kanade min_eigenvalue must be finite and positive".into(),
61            ));
62        }
63        Ok(self)
64    }
65}
66
67/// Result of sparse point tracking between two frames.
68#[derive(Clone, Debug, PartialEq)]
69pub struct TrackedPoints {
70    next_points: Vec<Vec2<f64>>,
71    status: Vec<bool>,
72}
73
74impl TrackedPoints {
75    /// Returns one next-frame coordinate per input point.
76    pub fn next_points(&self) -> &[Vec2<f64>] {
77        &self.next_points
78    }
79
80    /// Returns whether each point was successfully tracked.
81    pub fn status(&self) -> &[bool] {
82        &self.status
83    }
84}
85
86/// Tracks sparse points from `previous` to `next` with pyramidal Lucas–Kanade.
87pub fn track_points_lucas_kanade<T: PixelComponent>(
88    previous: ImageView<'_, T, 1>,
89    next: ImageView<'_, T, 1>,
90    points: &[Vec2<f64>],
91    options: LucasKanadeOptions,
92) -> VisionResult<TrackedPoints> {
93    let pair = LucasKanadePair::try_new(previous, next, options)?;
94    Ok(pair.track(points, false))
95}
96
97/// Owns the immutable CPU pyramids for one pair of frames.
98pub(crate) struct LucasKanadePair {
99    previous_pyramid: Vec<Image<f32, 1>>,
100    next_pyramid: Vec<Image<f32, 1>>,
101    options: LucasKanadeOptions,
102}
103
104impl LucasKanadePair {
105    pub(crate) fn try_new<T: PixelComponent>(
106        previous: ImageView<'_, T, 1>,
107        next: ImageView<'_, T, 1>,
108        options: LucasKanadeOptions,
109    ) -> VisionResult<Self> {
110        let options = options.validate()?;
111        if previous.width() != next.width() || previous.height() != next.height() {
112            return Err(VisionError::ShapeMismatch(
113                "Lucas–Kanade frames must share width and height".into(),
114            ));
115        }
116        if previous.width() < 3 || previous.height() < 3 {
117            return Err(VisionError::InvalidDimensions(
118                "Lucas–Kanade requires at least 3x3 images".into(),
119            ));
120        }
121        let previous_pyramid = build_pyramid(previous, options.pyramid_levels)?;
122        let next_pyramid = build_pyramid(next, options.pyramid_levels)?;
123        Ok(Self { previous_pyramid, next_pyramid, options })
124    }
125
126    pub(crate) fn track(&self, points: &[Vec2<f64>], reverse: bool) -> TrackedPoints {
127        let options = self.options;
128        let (previous_pyramid, next_pyramid) = if reverse {
129            (&self.next_pyramid, &self.previous_pyramid)
130        } else {
131            (&self.previous_pyramid, &self.next_pyramid)
132        };
133        let mut next_points = points.to_vec();
134        let mut status = vec![true; points.len()];
135        let scale = 1.0 / f64::from(1u32 << (options.pyramid_levels.saturating_sub(1) as u32));
136        for point in &mut next_points {
137            point.x *= scale;
138            point.y *= scale;
139        }
140        let mut previous_guess = next_points.clone();
141        for level in (0..options.pyramid_levels).rev() {
142            let level_scale = 1.0 / f64::from(1u32 << (level as u32));
143            for (index, previous_point) in points.iter().enumerate() {
144                if !status[index] {
145                    continue;
146                }
147                let mut guess = previous_guess[index];
148                let target =
149                    Vec2 { x: previous_point.x * level_scale, y: previous_point.y * level_scale };
150                match track_one_level(
151                    previous_pyramid[level].view(),
152                    next_pyramid[level].view(),
153                    target,
154                    guess,
155                    options,
156                ) {
157                    Ok(updated) => {
158                        guess = updated;
159                        status[index] = true;
160                    }
161                    Err(_) => status[index] = false,
162                }
163                previous_guess[index] = guess;
164                if level > 0 {
165                    previous_guess[index].x *= 2.0;
166                    previous_guess[index].y *= 2.0;
167                }
168            }
169        }
170        TrackedPoints { next_points: previous_guess, status }
171    }
172}
173
174fn track_one_level(
175    previous: ImageView<'_, f32, 1>,
176    next: ImageView<'_, f32, 1>,
177    previous_point: Vec2<f64>,
178    mut next_point: Vec2<f64>,
179    options: LucasKanadeOptions,
180) -> VisionResult<Vec2<f64>> {
181    if options.window_radius >= previous.width().min(previous.height()) {
182        return Err(VisionError::InvalidParameter("Lucas–Kanade patch exceeds the image".into()));
183    }
184    let radius = options.window_radius as f64;
185    let mut patch_cache: LkTemplateCache =
186        if options.window_radius <= 5 && options.max_iterations > 1 {
187            Some([None; 121])
188        } else {
189            None
190        };
191    for _ in 0..options.max_iterations {
192        let mut gxx = 0.0;
193        let mut gxy = 0.0;
194        let mut gyy = 0.0;
195        let mut bx = 0.0;
196        let mut by = 0.0;
197        let mut samples = 0usize;
198        let start = -(options.window_radius as isize);
199        let end = options.window_radius as isize;
200        for dy in start..=end {
201            for dx in start..=end {
202                let px = previous_point.x + dx as f64;
203                let py = previous_point.y + dy as f64;
204                let qx = next_point.x + dx as f64;
205                let qy = next_point.y + dy as f64;
206                if !(in_bounds(previous, px, py, radius) && in_bounds(next, qx, qy, radius)) {
207                    continue;
208                }
209                let (ix, iy, it) = if let Some(cache) = patch_cache.as_mut() {
210                    // Only radii <= 5 enter this branch, so the row-major slot
211                    // fits 11*11 entries. Both frame bounds were checked above.
212                    let slot = (dy - start) as usize * (2 * options.window_radius + 1)
213                        + (dx - start) as usize;
214                    if let Some(saved) = cache[slot] {
215                        (saved.ix, saved.iy, sample(next, qx, qy)? - saved.value)
216                    } else {
217                        let ix = 0.5
218                            * (sample(previous, px + 1.0, py)? - sample(previous, px - 1.0, py)?);
219                        let iy = 0.5
220                            * (sample(previous, px, py + 1.0)? - sample(previous, px, py - 1.0)?);
221                        let current = sample(next, qx, qy)?;
222                        let value = sample(previous, px, py)?;
223                        cache[slot] = Some(LkTemplateSample { ix, iy, value });
224                        (ix, iy, current - value)
225                    }
226                } else {
227                    let ix =
228                        0.5 * (sample(previous, px + 1.0, py)? - sample(previous, px - 1.0, py)?);
229                    let iy =
230                        0.5 * (sample(previous, px, py + 1.0)? - sample(previous, px, py - 1.0)?);
231                    let it = sample(next, qx, qy)? - sample(previous, px, py)?;
232                    (ix, iy, it)
233                };
234                gxx += ix * ix;
235                gxy += ix * iy;
236                gyy += iy * iy;
237                bx += ix * it;
238                by += iy * it;
239                samples += 1;
240            }
241        }
242        if samples < 4 {
243            return Err(VisionError::InvalidParameter("Lucas–Kanade patch left the image".into()));
244        }
245        let det = gxx * gyy - gxy * gxy;
246        if det.abs() < options.min_eigenvalue {
247            return Err(VisionError::InvalidParameter(
248                "Lucas–Kanade structure tensor is singular".into(),
249            ));
250        }
251        let dx = (-gyy * bx + gxy * by) / det;
252        let dy = (gxy * bx - gxx * by) / det;
253        next_point.x += dx;
254        next_point.y += dy;
255        if dx.hypot(dy) < options.epsilon {
256            break;
257        }
258    }
259    if !in_bounds(next, next_point.x, next_point.y, 0.0) {
260        return Err(VisionError::InvalidParameter(
261            "Lucas–Kanade tracked point left the image".into(),
262        ));
263    }
264    Ok(next_point)
265}
266
267#[derive(Clone, Copy)]
268struct LkTemplateSample {
269    ix: f64,
270    iy: f64,
271    value: f64,
272}
273
274// One point/level owns this bounded array; it never survives a tracking call.
275type LkTemplateCache = Option<[Option<LkTemplateSample>; 121]>;
276const _: () = assert!(std::mem::size_of::<LkTemplateCache>() <= 4096);
277
278fn build_pyramid<T: PixelComponent>(
279    image: ImageView<'_, T, 1>,
280    levels: usize,
281) -> VisionResult<Vec<Image<f32, 1>>> {
282    let mut pyramid = Vec::with_capacity(levels);
283    pyramid.push(to_f32_image(image)?);
284    for _ in 1..levels {
285        let previous = pyramid.last().expect("pyramid starts non-empty");
286        pyramid.push(downsample(previous.view())?);
287    }
288    Ok(pyramid)
289}
290
291fn to_f32_image<T: PixelComponent>(image: ImageView<'_, T, 1>) -> VisionResult<Image<f32, 1>> {
292    let mut data = Vec::with_capacity(image.width() * image.height());
293    for y in 0..image.height() {
294        for x in 0..image.width() {
295            data.push(image.get(x, y).expect("in-bounds")[0].to_f64() as f32);
296        }
297    }
298    Ok(Image::try_new(image.width(), image.height(), data)?)
299}
300
301fn downsample(image: ImageView<'_, f32, 1>) -> VisionResult<Image<f32, 1>> {
302    let width = (image.width() / 2).max(1);
303    let height = (image.height() / 2).max(1);
304    let mut data = Vec::with_capacity(width * height);
305    for y in 0..height {
306        for x in 0..width {
307            let sx = x * 2;
308            let sy = y * 2;
309            let mut sum = 0.0_f64;
310            let mut count = 0.0_f64;
311            for dy in 0..2 {
312                for dx in 0..2 {
313                    if let Some(pixel) = image.get(sx + dx, sy + dy) {
314                        sum += f64::from(pixel[0]);
315                        count += 1.0;
316                    }
317                }
318            }
319            data.push((sum / count.max(1.0)) as f32);
320        }
321    }
322    Ok(Image::try_new(width, height, data)?)
323}
324
325fn sample(image: ImageView<'_, f32, 1>, x: f64, y: f64) -> VisionResult<f64> {
326    if !x.is_finite() || !y.is_finite() {
327        return Err(VisionError::InvalidParameter(
328            "Lucas–Kanade sample coordinates must be finite".into(),
329        ));
330    }
331    let x0 = x.floor() as isize;
332    let y0 = y.floor() as isize;
333    let x1 = x0 + 1;
334    let y1 = y0 + 1;
335    if x0 < 0 || y0 < 0 || x1 >= image.width() as isize || y1 >= image.height() as isize {
336        return Err(VisionError::InvalidParameter("Lucas–Kanade sample is out of bounds".into()));
337    }
338    let ax = x - x0 as f64;
339    let ay = y - y0 as f64;
340    let i00 = f64::from(image.get(x0 as usize, y0 as usize).expect("in-bounds")[0]);
341    let i10 = f64::from(image.get(x1 as usize, y0 as usize).expect("in-bounds")[0]);
342    let i01 = f64::from(image.get(x0 as usize, y1 as usize).expect("in-bounds")[0]);
343    let i11 = f64::from(image.get(x1 as usize, y1 as usize).expect("in-bounds")[0]);
344    Ok((1.0 - ax) * (1.0 - ay) * i00
345        + ax * (1.0 - ay) * i10
346        + (1.0 - ax) * ay * i01
347        + ax * ay * i11)
348}
349
350fn in_bounds(image: ImageView<'_, f32, 1>, x: f64, y: f64, margin: f64) -> bool {
351    x >= margin
352        && y >= margin
353        && x < image.width() as f64 - 1.0 - margin
354        && y < image.height() as f64 - 1.0 - margin
355}
356
357#[cfg(test)]
358#[path = "optical_flow_template_tests.rs"]
359mod template_tests;
360
361#[cfg(test)]
362mod tests {
363    use super::{track_points_lucas_kanade, LucasKanadeOptions};
364    use spatialrust_image::Image;
365    use spatialrust_math::Vec2;
366
367    #[test]
368    fn rejects_unrepresentable_scales_and_windows_before_pyramid_allocation() {
369        for levels in [0, 33, usize::MAX] {
370            let options = LucasKanadeOptions { pyramid_levels: levels, ..Default::default() };
371            assert!(options.validate().is_err());
372        }
373        for radius in [0, isize::MAX as usize + 1, usize::MAX] {
374            let options = LucasKanadeOptions { window_radius: radius, ..Default::default() };
375            assert!(options.validate().is_err());
376        }
377        assert!(LucasKanadeOptions { pyramid_levels: 32, ..Default::default() }.validate().is_ok());
378    }
379
380    #[test]
381    fn oversized_representable_patch_fails_without_iterating_the_window() {
382        let image = Image::<u8, 1>::try_new(8, 8, vec![0; 64]).unwrap();
383        let result = track_points_lucas_kanade(
384            image.view(),
385            image.view(),
386            &[Vec2 { x: 4.0, y: 4.0 }],
387            LucasKanadeOptions {
388                window_radius: isize::MAX as usize,
389                pyramid_levels: 1,
390                ..Default::default()
391            },
392        )
393        .unwrap();
394        assert_eq!(result.status(), &[false]);
395    }
396
397    #[test]
398    fn tracks_integer_translation() {
399        let width = 64;
400        let height = 48;
401        let mut previous = vec![0u8; width * height];
402        for y in 0..height {
403            for x in 0..width {
404                previous[y * width + x] = ((x * 13 + y * 7) % 200 + 20) as u8;
405            }
406        }
407        let shift = 3isize;
408        let mut next = vec![0u8; width * height];
409        for y in 0..height {
410            for x in 0..width {
411                let sx = x as isize - shift;
412                let sy = y as isize;
413                if (0..width as isize).contains(&sx) && (0..height as isize).contains(&sy) {
414                    next[y * width + x] = previous[sy as usize * width + sx as usize];
415                }
416            }
417        }
418        let previous = Image::<u8, 1>::try_new(width, height, previous).unwrap();
419        let next = Image::<u8, 1>::try_new(width, height, next).unwrap();
420        let points = [Vec2 { x: 24.0, y: 20.0 }];
421        let tracked = track_points_lucas_kanade(
422            previous.view(),
423            next.view(),
424            &points,
425            LucasKanadeOptions {
426                window_radius: 4,
427                pyramid_levels: 1,
428                max_iterations: 50,
429                epsilon: 1e-4,
430                min_eigenvalue: 1e-8,
431            },
432        )
433        .unwrap();
434        assert!(tracked.status().iter().all(|&ok| ok), "status={:?}", tracked.status());
435        let actual = tracked.next_points()[0];
436        assert!(
437            (actual.x - (points[0].x + shift as f64)).abs() < 1.5,
438            "x actual={} expected={}",
439            actual.x,
440            points[0].x + shift as f64
441        );
442        assert!(
443            (actual.y - points[0].y).abs() < 1.5,
444            "y actual={} expected={}",
445            actual.y,
446            points[0].y
447        );
448    }
449}
spatialrust_vision/optical_flow.rs

spatialrust_vision/
optical_flow.rs
