Skip to main content

spatialrust_vision/
odometry.rs

1//! Robust feature tracking and calibrated visual odometry building blocks.
2
3use spatialrust_image::ImageView;
4use spatialrust_math::{Vec2, Vec3};
5
6use crate::{
7    estimate_essential_ransac, recover_relative_pose, solve_pnp_ransac, track_points_lucas_kanade,
8    AbsolutePose, CameraMatrix3, Keypoint2, LucasKanadeOptions, ObjectImageCorrespondence,
9    PointCorrespondence2, RelativePose, RobustEstimationOptions, VisionError, VisionResult,
10};
11
12/// Controls deterministic spatial distribution of candidate keypoints.
13#[derive(Clone, Copy, Debug, PartialEq, Eq)]
14pub struct GridSelectionOptions {
15    /// Width of one grid cell in pixels.
16    pub cell_width: usize,
17    /// Height of one grid cell in pixels.
18    pub cell_height: usize,
19    /// Maximum retained candidates per cell.
20    pub max_per_cell: usize,
21}
22
23impl Default for GridSelectionOptions {
24    fn default() -> Self {
25        Self { cell_width: 32, cell_height: 32, max_per_cell: 4 }
26    }
27}
28
29/// Retains the strongest keypoints in each image cell.
30///
31/// Output is ordered by cell scan order, then decreasing response, with the
32/// input index as the deterministic tie breaker.
33pub fn select_keypoints_grid(
34    keypoints: &[Keypoint2],
35    width: usize,
36    height: usize,
37    options: GridSelectionOptions,
38) -> VisionResult<Vec<Keypoint2>> {
39    if width == 0
40        || height == 0
41        || options.cell_width == 0
42        || options.cell_height == 0
43        || options.max_per_cell == 0
44    {
45        return Err(VisionError::InvalidParameter(
46            "grid selection dimensions and limits must be positive".into(),
47        ));
48    }
49    let columns = width.div_ceil(options.cell_width);
50    let rows = height.div_ceil(options.cell_height);
51    let mut cells = vec![Vec::<(usize, Keypoint2)>::new(); columns * rows];
52    for (index, &keypoint) in keypoints.iter().enumerate() {
53        if keypoint.x() < 0.0
54            || keypoint.y() < 0.0
55            || keypoint.x() >= width as f32
56            || keypoint.y() >= height as f32
57        {
58            continue;
59        }
60        let column = keypoint.x() as usize / options.cell_width;
61        let row = keypoint.y() as usize / options.cell_height;
62        cells[row * columns + column].push((index, keypoint));
63    }
64    let mut selected = Vec::new();
65    for cell in &mut cells {
66        cell.sort_by(|(left_index, left), (right_index, right)| {
67            right.response().total_cmp(&left.response()).then(left_index.cmp(right_index))
68        });
69        selected.extend(cell.iter().take(options.max_per_cell).map(|(_, point)| *point));
70    }
71    Ok(selected)
72}
73
74/// Forward/backward consistency settings layered over pyramidal LK.
75#[derive(Clone, Copy, Debug, PartialEq)]
76pub struct RobustTrackOptions {
77    /// Underlying LK settings in both directions.
78    pub lucas_kanade: LucasKanadeOptions,
79    /// Maximum accepted round-trip distance in pixels.
80    pub max_forward_backward_error: f64,
81}
82
83impl Default for RobustTrackOptions {
84    fn default() -> Self {
85        Self { lucas_kanade: LucasKanadeOptions::default(), max_forward_backward_error: 1.0 }
86    }
87}
88
89/// One target coordinate, validity decision, and round-trip error per source point.
90#[derive(Clone, Debug, PartialEq)]
91pub struct RobustTracks {
92    next_points: Vec<Vec2<f64>>,
93    status: Vec<bool>,
94    forward_backward_errors: Vec<f64>,
95}
96
97impl RobustTracks {
98    /// Returns one next-frame coordinate per source point.
99    pub fn next_points(&self) -> &[Vec2<f64>] {
100        &self.next_points
101    }
102    /// Returns the combined forward, backward, and threshold decision.
103    pub fn status(&self) -> &[bool] {
104        &self.status
105    }
106    /// Returns round-trip pixel error, or infinity when either direction failed.
107    pub fn forward_backward_errors(&self) -> &[f64] {
108        &self.forward_backward_errors
109    }
110    /// Returns the accepted track count.
111    pub fn accepted_count(&self) -> usize {
112        self.status.iter().filter(|&&value| value).count()
113    }
114}
115
116/// Tracks points in both directions and rejects inconsistent round trips.
117pub fn track_points_forward_backward<T: crate::PixelComponent>(
118    previous: ImageView<'_, T, 1>,
119    next: ImageView<'_, T, 1>,
120    points: &[Vec2<f64>],
121    options: RobustTrackOptions,
122) -> VisionResult<RobustTracks> {
123    if !options.max_forward_backward_error.is_finite() || options.max_forward_backward_error < 0.0 {
124        return Err(VisionError::InvalidParameter(
125            "forward/backward threshold must be finite and non-negative".into(),
126        ));
127    }
128    // Custom components may observe conversion calls; keep their original flow.
129    let scalar = std::any::TypeId::of::<T>();
130    let built_in = [
131        std::any::TypeId::of::<u8>(),
132        std::any::TypeId::of::<u16>(),
133        std::any::TypeId::of::<f32>(),
134        std::any::TypeId::of::<f64>(),
135    ]
136    .contains(&scalar);
137    let (forward, backward) = if built_in {
138        let pair =
139            crate::optical_flow::LucasKanadePair::try_new(previous, next, options.lucas_kanade)?;
140        let forward = pair.track(points, false);
141        let backward = pair.track(forward.next_points(), true);
142        (forward, backward)
143    } else {
144        let forward = track_points_lucas_kanade(previous, next, points, options.lucas_kanade)?;
145        let backward =
146            track_points_lucas_kanade(next, previous, forward.next_points(), options.lucas_kanade)?;
147        (forward, backward)
148    };
149    let mut status = Vec::with_capacity(points.len());
150    let mut errors = Vec::with_capacity(points.len());
151    for (index, point) in points.iter().enumerate() {
152        let valid = forward.status()[index] && backward.status()[index];
153        let error = if valid {
154            let dx = backward.next_points()[index].x - point.x;
155            let dy = backward.next_points()[index].y - point.y;
156            dx.hypot(dy)
157        } else {
158            f64::INFINITY
159        };
160        status.push(valid && error <= options.max_forward_backward_error);
161        errors.push(error);
162    }
163    Ok(RobustTracks {
164        next_points: forward.next_points().to_vec(),
165        status,
166        forward_backward_errors: errors,
167    })
168}
169
170/// Scale-ambiguous monocular odometry result.
171#[derive(Clone, Debug, PartialEq)]
172pub struct MonocularOdometryEstimate {
173    /// Source-to-target pose; translation has unit, arbitrary scale.
174    pub pose: RelativePose,
175    /// Essential-matrix RANSAC inliers.
176    pub inliers: Vec<bool>,
177    /// Correspondences triangulated in front of both cameras.
178    pub positive_depth_count: usize,
179}
180
181/// Estimates calibrated monocular motion with essential RANSAC and cheirality.
182pub fn estimate_monocular_odometry(
183    correspondences: &[PointCorrespondence2],
184    camera: CameraMatrix3,
185    options: RobustEstimationOptions,
186) -> VisionResult<MonocularOdometryEstimate> {
187    let estimate = estimate_essential_ransac(correspondences, camera, camera, options)?;
188    let inlier_pairs = correspondences
189        .iter()
190        .zip(estimate.inliers())
191        .filter_map(|(&pair, &inlier)| inlier.then_some(pair))
192        .collect::<Vec<_>>();
193    if inlier_pairs.len() < 8 {
194        return Err(VisionError::InvalidParameter(
195            "monocular odometry requires at least eight essential inliers".into(),
196        ));
197    }
198    let recovered = recover_relative_pose(*estimate.model(), &inlier_pairs, camera, camera)?;
199    Ok(MonocularOdometryEstimate {
200        pose: recovered.pose(),
201        inliers: estimate.inliers().to_vec(),
202        positive_depth_count: recovered.positive_depth_count(),
203    })
204}
205
206/// Depth filtering and PnP RANSAC controls for metric RGB-D odometry.
207#[derive(Clone, Copy, Debug, PartialEq)]
208pub struct RgbdOdometryOptions {
209    /// Multiplier converting stored depth to metres.
210    pub depth_scale: f64,
211    /// Inclusive minimum accepted metric depth.
212    pub min_depth: f64,
213    /// Inclusive maximum accepted metric depth.
214    pub max_depth: f64,
215    /// PnP robust-estimation settings (threshold in target pixels).
216    pub robust: RobustEstimationOptions,
217}
218
219impl Default for RgbdOdometryOptions {
220    fn default() -> Self {
221        Self {
222            depth_scale: 1.0,
223            min_depth: 0.1,
224            max_depth: 100.0,
225            robust: RobustEstimationOptions { threshold: 1.0, ..Default::default() },
226        }
227    }
228}
229
230/// Metric previous-camera to current-camera odometry result.
231#[derive(Clone, Debug, PartialEq)]
232pub struct RgbdOdometryEstimate {
233    /// Full metric source-to-target pose.
234    pub pose: AbsolutePose,
235    /// RANSAC decisions over correspondences that had valid source depth.
236    pub inliers: Vec<bool>,
237    /// Number of input rows discarded for missing/out-of-range source depth.
238    pub rejected_depth_count: usize,
239}
240
241/// Estimates metric motion from source depth and source-to-target pixel tracks.
242pub fn estimate_rgbd_odometry<T: crate::PixelComponent>(
243    previous_depth: ImageView<'_, T, 1>,
244    correspondences: &[PointCorrespondence2],
245    camera: CameraMatrix3,
246    options: RgbdOdometryOptions,
247) -> VisionResult<RgbdOdometryEstimate> {
248    if !options.depth_scale.is_finite()
249        || options.depth_scale <= 0.0
250        || !options.min_depth.is_finite()
251        || options.min_depth <= 0.0
252        || options.max_depth.is_nan()
253        || options.max_depth < options.min_depth
254    {
255        return Err(VisionError::InvalidParameter("invalid RGB-D odometry depth range".into()));
256    }
257    let mut pairs = Vec::new();
258    let mut rejected = 0;
259    for &pair in correspondences {
260        let pixel = pair.source();
261        let x = pixel.x.round() as isize;
262        let y = pixel.y.round() as isize;
263        let Some(sample) =
264            (x >= 0 && y >= 0).then(|| previous_depth.get(x as usize, y as usize)).flatten()
265        else {
266            rejected += 1;
267            continue;
268        };
269        let depth = sample[0].to_f64() * options.depth_scale;
270        if !depth.is_finite() || depth < options.min_depth || depth > options.max_depth {
271            rejected += 1;
272            continue;
273        }
274        let ray = camera.normalize_pixel(pixel);
275        pairs.push(ObjectImageCorrespondence::try_new(
276            Vec3::new(ray.x * depth, ray.y * depth, depth),
277            pair.target(),
278        )?);
279    }
280    if pairs.len() < 6 {
281        return Err(VisionError::InvalidParameter(
282            "RGB-D odometry requires at least six tracks with valid depth".into(),
283        ));
284    }
285    let estimate = solve_pnp_ransac(&pairs, camera, options.robust)?;
286    Ok(RgbdOdometryEstimate {
287        pose: *estimate.model(),
288        inliers: estimate.inliers().to_vec(),
289        rejected_depth_count: rejected,
290    })
291}
292
293#[cfg(test)]
294#[path = "odometry_pair_tests.rs"]
295mod pair_tests;
296
297#[cfg(test)]
298mod tests {
299    use super::{
300        estimate_rgbd_odometry, select_keypoints_grid, GridSelectionOptions, RgbdOdometryOptions,
301    };
302    use crate::{CameraMatrix3, Keypoint2, PointCorrespondence2, RobustEstimationOptions};
303    use spatialrust_camera::CameraIntrinsics;
304    use spatialrust_image::Image;
305    use spatialrust_math::Vec2;
306
307    #[test]
308    fn grid_selection_keeps_strongest_per_cell() {
309        let points = vec![
310            Keypoint2::try_new(2.0, 2.0, 1.0).unwrap(),
311            Keypoint2::try_new(3.0, 3.0, 4.0).unwrap(),
312            Keypoint2::try_new(12.0, 2.0, 2.0).unwrap(),
313        ];
314        let selected = select_keypoints_grid(
315            &points,
316            20,
317            10,
318            GridSelectionOptions { cell_width: 10, cell_height: 10, max_per_cell: 1 },
319        )
320        .unwrap();
321        assert_eq!(selected.len(), 2);
322        assert_eq!(selected[0].response(), 4.0);
323        assert_eq!(selected[1].response(), 2.0);
324    }
325
326    #[test]
327    fn rgbd_odometry_recovers_metric_translation() {
328        let camera = CameraMatrix3::from_intrinsics(
329            CameraIntrinsics::try_new(100.0, 100.0, 50.0, 40.0, 100, 80).unwrap(),
330        );
331        let mut depth = Image::<f32, 1>::from_pixel(100, 80, [f32::NAN]).unwrap();
332        let mut pairs = Vec::new();
333        for (index, (x, y)) in [
334            (20, 20),
335            (35, 20),
336            (50, 20),
337            (65, 20),
338            (80, 20),
339            (20, 35),
340            (35, 35),
341            (50, 35),
342            (65, 35),
343            (80, 35),
344            (25, 55),
345            (45, 55),
346            (65, 55),
347            (80, 55),
348        ]
349        .into_iter()
350        .enumerate()
351        {
352            let z = 1.5 + index as f64 * 0.07;
353            depth.get_mut(x, y).unwrap()[0] = z as f32;
354            let source = Vec2 { x: x as f64, y: y as f64 };
355            let target = Vec2 { x: source.x + 100.0 * 0.1 / z, y: source.y };
356            pairs.push(PointCorrespondence2::try_new(source, target).unwrap());
357        }
358        let estimate = estimate_rgbd_odometry(
359            depth.view(),
360            &pairs,
361            camera,
362            RgbdOdometryOptions {
363                robust: RobustEstimationOptions {
364                    threshold: 0.05,
365                    max_iterations: 500,
366                    ..Default::default()
367                },
368                ..Default::default()
369            },
370        )
371        .unwrap();
372        assert!((estimate.pose.translation().x - 0.1).abs() < 1e-4);
373        assert!(estimate.pose.translation().y.abs() < 1e-4);
374        assert!(estimate.pose.translation().z.abs() < 1e-4);
375        assert_eq!(estimate.inliers.iter().filter(|&&value| value).count(), pairs.len());
376    }
377}