Skip to main content

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 {
39            return Err(VisionError::InvalidParameter(
40                "Lucas–Kanade window_radius must be positive".into(),
41            ));
42        }
43        if self.pyramid_levels == 0 {
44            return Err(VisionError::InvalidParameter(
45                "Lucas–Kanade pyramid_levels must be positive".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 options = options.validate()?;
94    if previous.width() != next.width() || previous.height() != next.height() {
95        return Err(VisionError::ShapeMismatch(
96            "Lucas–Kanade frames must share width and height".into(),
97        ));
98    }
99    if previous.width() < 3 || previous.height() < 3 {
100        return Err(VisionError::InvalidDimensions(
101            "Lucas–Kanade requires at least 3x3 images".into(),
102        ));
103    }
104    let previous_pyramid = build_pyramid(previous, options.pyramid_levels)?;
105    let next_pyramid = build_pyramid(next, options.pyramid_levels)?;
106    let mut next_points = points.to_vec();
107    let mut status = vec![true; points.len()];
108    let scale = 1.0 / f64::from(1u32 << (options.pyramid_levels.saturating_sub(1) as u32));
109    for point in &mut next_points {
110        point.x *= scale;
111        point.y *= scale;
112    }
113    let mut previous_guess = next_points.clone();
114    for level in (0..options.pyramid_levels).rev() {
115        let level_scale = 1.0 / f64::from(1u32 << (level as u32));
116        for (index, previous_point) in points.iter().enumerate() {
117            if !status[index] {
118                continue;
119            }
120            let mut guess = previous_guess[index];
121            let target =
122                Vec2 { x: previous_point.x * level_scale, y: previous_point.y * level_scale };
123            match track_one_level(
124                previous_pyramid[level].view(),
125                next_pyramid[level].view(),
126                target,
127                guess,
128                options,
129            ) {
130                Ok(updated) => {
131                    guess = updated;
132                    status[index] = true;
133                }
134                Err(_) => status[index] = false,
135            }
136            previous_guess[index] = guess;
137            if level > 0 {
138                previous_guess[index].x *= 2.0;
139                previous_guess[index].y *= 2.0;
140            }
141        }
142    }
143    Ok(TrackedPoints { next_points: previous_guess, status })
144}
145
146fn track_one_level(
147    previous: ImageView<'_, f32, 1>,
148    next: ImageView<'_, f32, 1>,
149    previous_point: Vec2<f64>,
150    mut next_point: Vec2<f64>,
151    options: LucasKanadeOptions,
152) -> VisionResult<Vec2<f64>> {
153    let radius = options.window_radius as f64;
154    for _ in 0..options.max_iterations {
155        let mut gxx = 0.0;
156        let mut gxy = 0.0;
157        let mut gyy = 0.0;
158        let mut bx = 0.0;
159        let mut by = 0.0;
160        let mut samples = 0usize;
161        let start = -(options.window_radius as isize);
162        let end = options.window_radius as isize;
163        for dy in start..=end {
164            for dx in start..=end {
165                let px = previous_point.x + dx as f64;
166                let py = previous_point.y + dy as f64;
167                let qx = next_point.x + dx as f64;
168                let qy = next_point.y + dy as f64;
169                if !(in_bounds(previous, px, py, radius) && in_bounds(next, qx, qy, radius)) {
170                    continue;
171                }
172                let ix = 0.5 * (sample(previous, px + 1.0, py)? - sample(previous, px - 1.0, py)?);
173                let iy = 0.5 * (sample(previous, px, py + 1.0)? - sample(previous, px, py - 1.0)?);
174                let it = sample(next, qx, qy)? - sample(previous, px, py)?;
175                gxx += ix * ix;
176                gxy += ix * iy;
177                gyy += iy * iy;
178                bx += ix * it;
179                by += iy * it;
180                samples += 1;
181            }
182        }
183        if samples < 4 {
184            return Err(VisionError::InvalidParameter("Lucas–Kanade patch left the image".into()));
185        }
186        let det = gxx * gyy - gxy * gxy;
187        if det.abs() < options.min_eigenvalue {
188            return Err(VisionError::InvalidParameter(
189                "Lucas–Kanade structure tensor is singular".into(),
190            ));
191        }
192        let dx = (-gyy * bx + gxy * by) / det;
193        let dy = (gxy * bx - gxx * by) / det;
194        next_point.x += dx;
195        next_point.y += dy;
196        if dx.hypot(dy) < options.epsilon {
197            break;
198        }
199    }
200    if !in_bounds(next, next_point.x, next_point.y, 0.0) {
201        return Err(VisionError::InvalidParameter(
202            "Lucas–Kanade tracked point left the image".into(),
203        ));
204    }
205    Ok(next_point)
206}
207
208fn build_pyramid<T: PixelComponent>(
209    image: ImageView<'_, T, 1>,
210    levels: usize,
211) -> VisionResult<Vec<Image<f32, 1>>> {
212    let mut pyramid = Vec::with_capacity(levels);
213    pyramid.push(to_f32_image(image)?);
214    for _ in 1..levels {
215        let previous = pyramid.last().expect("pyramid starts non-empty");
216        pyramid.push(downsample(previous.view())?);
217    }
218    Ok(pyramid)
219}
220
221fn to_f32_image<T: PixelComponent>(image: ImageView<'_, T, 1>) -> VisionResult<Image<f32, 1>> {
222    let mut data = Vec::with_capacity(image.width() * image.height());
223    for y in 0..image.height() {
224        for x in 0..image.width() {
225            data.push(image.get(x, y).expect("in-bounds")[0].to_f64() as f32);
226        }
227    }
228    Ok(Image::try_new(image.width(), image.height(), data)?)
229}
230
231fn downsample(image: ImageView<'_, f32, 1>) -> VisionResult<Image<f32, 1>> {
232    let width = (image.width() / 2).max(1);
233    let height = (image.height() / 2).max(1);
234    let mut data = Vec::with_capacity(width * height);
235    for y in 0..height {
236        for x in 0..width {
237            let sx = x * 2;
238            let sy = y * 2;
239            let mut sum = 0.0_f64;
240            let mut count = 0.0_f64;
241            for dy in 0..2 {
242                for dx in 0..2 {
243                    if let Some(pixel) = image.get(sx + dx, sy + dy) {
244                        sum += f64::from(pixel[0]);
245                        count += 1.0;
246                    }
247                }
248            }
249            data.push((sum / count.max(1.0)) as f32);
250        }
251    }
252    Ok(Image::try_new(width, height, data)?)
253}
254
255fn sample(image: ImageView<'_, f32, 1>, x: f64, y: f64) -> VisionResult<f64> {
256    if !x.is_finite() || !y.is_finite() {
257        return Err(VisionError::InvalidParameter(
258            "Lucas–Kanade sample coordinates must be finite".into(),
259        ));
260    }
261    let x0 = x.floor() as isize;
262    let y0 = y.floor() as isize;
263    let x1 = x0 + 1;
264    let y1 = y0 + 1;
265    if x0 < 0 || y0 < 0 || x1 >= image.width() as isize || y1 >= image.height() as isize {
266        return Err(VisionError::InvalidParameter("Lucas–Kanade sample is out of bounds".into()));
267    }
268    let ax = x - x0 as f64;
269    let ay = y - y0 as f64;
270    let i00 = f64::from(image.get(x0 as usize, y0 as usize).expect("in-bounds")[0]);
271    let i10 = f64::from(image.get(x1 as usize, y0 as usize).expect("in-bounds")[0]);
272    let i01 = f64::from(image.get(x0 as usize, y1 as usize).expect("in-bounds")[0]);
273    let i11 = f64::from(image.get(x1 as usize, y1 as usize).expect("in-bounds")[0]);
274    Ok((1.0 - ax) * (1.0 - ay) * i00
275        + ax * (1.0 - ay) * i10
276        + (1.0 - ax) * ay * i01
277        + ax * ay * i11)
278}
279
280fn in_bounds(image: ImageView<'_, f32, 1>, x: f64, y: f64, margin: f64) -> bool {
281    x >= margin
282        && y >= margin
283        && x < image.width() as f64 - 1.0 - margin
284        && y < image.height() as f64 - 1.0 - margin
285}
286
287#[cfg(test)]
288mod tests {
289    use super::{track_points_lucas_kanade, LucasKanadeOptions};
290    use spatialrust_image::Image;
291    use spatialrust_math::Vec2;
292
293    #[test]
294    fn tracks_integer_translation() {
295        let width = 64;
296        let height = 48;
297        let mut previous = vec![0u8; width * height];
298        for y in 0..height {
299            for x in 0..width {
300                previous[y * width + x] = ((x * 13 + y * 7) % 200 + 20) as u8;
301            }
302        }
303        let shift = 3isize;
304        let mut next = vec![0u8; width * height];
305        for y in 0..height {
306            for x in 0..width {
307                let sx = x as isize - shift;
308                let sy = y as isize;
309                if (0..width as isize).contains(&sx) && (0..height as isize).contains(&sy) {
310                    next[y * width + x] = previous[sy as usize * width + sx as usize];
311                }
312            }
313        }
314        let previous = Image::<u8, 1>::try_new(width, height, previous).unwrap();
315        let next = Image::<u8, 1>::try_new(width, height, next).unwrap();
316        let points = [Vec2 { x: 24.0, y: 20.0 }];
317        let tracked = track_points_lucas_kanade(
318            previous.view(),
319            next.view(),
320            &points,
321            LucasKanadeOptions {
322                window_radius: 4,
323                pyramid_levels: 1,
324                max_iterations: 50,
325                epsilon: 1e-4,
326                min_eigenvalue: 1e-8,
327            },
328        )
329        .unwrap();
330        assert!(tracked.status().iter().all(|&ok| ok), "status={:?}", tracked.status());
331        let actual = tracked.next_points()[0];
332        assert!(
333            (actual.x - (points[0].x + shift as f64)).abs() < 1.5,
334            "x actual={} expected={}",
335            actual.x,
336            points[0].x + shift as f64
337        );
338        assert!(
339            (actual.y - points[0].y).abs() < 1.5,
340            "y actual={} expected={}",
341            actual.y,
342            points[0].y
343        );
344    }
345}