Skip to main content

spatialrust_vision/
stereo.rs

1//! Stereo rig contracts, rectification maps, and block-matching disparity.
2
3use spatialrust_image::{Image, ImageView};
4use spatialrust_math::{Mat3, Vec3};
5
6use crate::{CameraMatrix3, PixelComponent, RelativePose, VisionError, VisionResult};
7
8/// Invalid disparity sentinel written by [`stereo_block_match`].
9pub const INVALID_DISPARITY: f32 = -1.0;
10
11/// Calibrated two-camera stereo geometry.
12#[derive(Clone, Copy, Debug, PartialEq)]
13pub struct StereoRig {
14    left: CameraMatrix3,
15    right: CameraMatrix3,
16    pose: RelativePose,
17}
18
19impl StereoRig {
20    /// Creates a rig from left/right intrinsics and the right-camera pose in the
21    /// left-camera frame (`X_right = R X_left + t`).
22    pub fn try_new(
23        left: CameraMatrix3,
24        right: CameraMatrix3,
25        pose: RelativePose,
26    ) -> VisionResult<Self> {
27        Ok(Self { left, right, pose })
28    }
29
30    /// Returns the left intrinsic matrix.
31    pub const fn left(self) -> CameraMatrix3 {
32        self.left
33    }
34
35    /// Returns the right intrinsic matrix.
36    pub const fn right(self) -> CameraMatrix3 {
37        self.right
38    }
39
40    /// Returns the right camera pose expressed in the left camera frame.
41    pub const fn pose(self) -> RelativePose {
42        self.pose
43    }
44
45    /// Returns the absolute baseline length `|t|`.
46    #[must_use]
47    pub fn baseline(self) -> f64 {
48        self.pose.translation().length()
49    }
50}
51
52/// Remap grids produced by stereo rectification for explicit `warp::remap`.
53#[derive(Clone, Debug, PartialEq)]
54pub struct StereoRectifyMaps {
55    left_map_x: Image<f32, 1>,
56    left_map_y: Image<f32, 1>,
57    right_map_x: Image<f32, 1>,
58    right_map_y: Image<f32, 1>,
59    rectified_left: CameraMatrix3,
60    rectified_right: CameraMatrix3,
61    baseline: f64,
62}
63
64impl StereoRectifyMaps {
65    /// Returns the left absolute-x remap image.
66    pub const fn left_map_x(&self) -> &Image<f32, 1> {
67        &self.left_map_x
68    }
69
70    /// Returns the left absolute-y remap image.
71    pub const fn left_map_y(&self) -> &Image<f32, 1> {
72        &self.left_map_y
73    }
74
75    /// Returns the right absolute-x remap image.
76    pub const fn right_map_x(&self) -> &Image<f32, 1> {
77        &self.right_map_x
78    }
79
80    /// Returns the right absolute-y remap image.
81    pub const fn right_map_y(&self) -> &Image<f32, 1> {
82        &self.right_map_y
83    }
84
85    /// Returns the shared rectified left intrinsics.
86    pub const fn rectified_left(&self) -> CameraMatrix3 {
87        self.rectified_left
88    }
89
90    /// Returns the shared rectified right intrinsics.
91    pub const fn rectified_right(&self) -> CameraMatrix3 {
92        self.rectified_right
93    }
94
95    /// Returns the positive rectified baseline along +X.
96    pub const fn baseline(&self) -> f64 {
97        self.baseline
98    }
99}
100
101/// Block-matching stereo options.
102#[derive(Clone, Copy, Debug, PartialEq)]
103pub struct StereoBmOptions {
104    /// Odd SAD window size in pixels.
105    pub window_size: usize,
106    /// Inclusive minimum positive disparity in pixels.
107    pub min_disparity: i32,
108    /// Number of disparities searched (must be positive and even-friendly).
109    pub num_disparities: i32,
110    /// Uniqueness ratio in percent; candidates failing it are invalid.
111    pub uniqueness_ratio: f32,
112}
113
114impl Default for StereoBmOptions {
115    fn default() -> Self {
116        Self { window_size: 15, min_disparity: 0, num_disparities: 64, uniqueness_ratio: 15.0 }
117    }
118}
119
120impl StereoBmOptions {
121    /// Validates window and disparity search settings.
122    pub fn validate(self) -> VisionResult<Self> {
123        if self.window_size < 3 || self.window_size % 2 == 0 {
124            return Err(VisionError::InvalidParameter(
125                "stereo BM window_size must be odd and at least 3".into(),
126            ));
127        }
128        if self.num_disparities <= 0 {
129            return Err(VisionError::InvalidParameter(
130                "stereo BM num_disparities must be positive".into(),
131            ));
132        }
133        if !self.uniqueness_ratio.is_finite() || self.uniqueness_ratio < 0.0 {
134            return Err(VisionError::InvalidParameter(
135                "stereo BM uniqueness_ratio must be finite and non-negative".into(),
136            ));
137        }
138        Ok(self)
139    }
140}
141
142/// Builds left/right absolute remap grids that make epipolar lines horizontal.
143///
144/// Callers feed these maps into `warp::remap`. Identical fronto-parallel
145/// cameras with a pure +X baseline produce identity remaps.
146pub fn stereo_rectify(
147    rig: StereoRig,
148    width: usize,
149    height: usize,
150) -> VisionResult<StereoRectifyMaps> {
151    if width == 0 || height == 0 {
152        return Err(VisionError::InvalidDimensions(
153            "stereo rectify requires positive width and height".into(),
154        ));
155    }
156    let translation = rig.pose().translation();
157    let baseline = translation.length();
158    if baseline <= f64::EPSILON {
159        return Err(VisionError::InvalidParameter("stereo baseline must be non-zero".into()));
160    }
161    let e1 = translation.normalize();
162    let helper = if e1.x.abs() < 0.9 { Vec3::new(1.0, 0.0, 0.0) } else { Vec3::new(0.0, 1.0, 0.0) };
163    let e2 = e1.cross(helper).normalize();
164    let e3 = e1.cross(e2).normalize();
165    let r_rect = Mat3::from_rows([e1.x, e1.y, e1.z], [e2.x, e2.y, e2.z], [e3.x, e3.y, e3.z]);
166    let left_rotation = r_rect;
167    let right_rotation = r_rect.mul_mat3(rig.pose().rotation());
168    let fx = 0.5 * (rig.left().matrix().m[0][0] + rig.right().matrix().m[0][0]);
169    let fy = 0.5 * (rig.left().matrix().m[1][1] + rig.right().matrix().m[1][1]);
170    let cx = (width as f64 - 1.0) * 0.5;
171    let cy = (height as f64 - 1.0) * 0.5;
172    let new_k = Mat3::from_rows([fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]);
173    let new_camera = CameraMatrix3::try_from_pinhole(fx, fy, cx, cy)?;
174    let left_maps = build_rectify_maps(rig.left(), left_rotation, new_k, width, height)?;
175    let right_maps = build_rectify_maps(rig.right(), right_rotation, new_k, width, height)?;
176    Ok(StereoRectifyMaps {
177        left_map_x: left_maps.0,
178        left_map_y: left_maps.1,
179        right_map_x: right_maps.0,
180        right_map_y: right_maps.1,
181        rectified_left: new_camera,
182        rectified_right: new_camera,
183        baseline,
184    })
185}
186
187/// Dense SAD block matching on already-rectified grayscale stereo images.
188///
189/// Invalid disparities are set to [`INVALID_DISPARITY`]. Search looks for
190/// matches of the left pixel in the right image at `x - d`.
191pub fn stereo_block_match<T: PixelComponent>(
192    left: ImageView<'_, T, 1>,
193    right: ImageView<'_, T, 1>,
194    options: StereoBmOptions,
195) -> VisionResult<Image<f32, 1>> {
196    let options = options.validate()?;
197    if left.width() != right.width() || left.height() != right.height() {
198        return Err(VisionError::ShapeMismatch(
199            "stereo BM frames must share width and height".into(),
200        ));
201    }
202    let width = left.width();
203    let height = left.height();
204    let radius = options.window_size / 2;
205    let mut data = vec![INVALID_DISPARITY; width * height];
206    for y in radius..(height.saturating_sub(radius)) {
207        for x in radius..(width.saturating_sub(radius)) {
208            let mut best_cost = f64::INFINITY;
209            let mut second_cost = f64::INFINITY;
210            let mut best_d = 0i32;
211            let d0 = options.min_disparity;
212            let d1 = options.min_disparity + options.num_disparities;
213            for disparity in d0..d1 {
214                let xr = x as i32 - disparity;
215                if xr < radius as i32 || xr >= (width - radius) as i32 {
216                    continue;
217                }
218                let mut cost = 0.0;
219                for dy in -(radius as isize)..=(radius as isize) {
220                    for dx in -(radius as isize)..=(radius as isize) {
221                        let ly = (y as isize + dy) as usize;
222                        let lx = (x as isize + dx) as usize;
223                        let ry = ly;
224                        let rx = (xr as isize + dx) as usize;
225                        let left_value = left.get(lx, ly).expect("in-bounds")[0].to_f64();
226                        let right_value = right.get(rx, ry).expect("in-bounds")[0].to_f64();
227                        cost += (left_value - right_value).abs();
228                    }
229                }
230                if cost < best_cost {
231                    second_cost = best_cost;
232                    best_cost = cost;
233                    best_d = disparity;
234                } else if cost < second_cost {
235                    second_cost = cost;
236                }
237            }
238            let unique = if !best_cost.is_finite() {
239                false
240            } else if !second_cost.is_finite() || second_cost <= best_cost {
241                true
242            } else {
243                second_cost >= best_cost * (1.0 + f64::from(options.uniqueness_ratio) / 100.0)
244            };
245            if unique && best_d > options.min_disparity {
246                data[y * width + x] = best_d as f32;
247            }
248        }
249    }
250    Ok(Image::try_new(width, height, data)?)
251}
252
253/// Converts horizontal disparity to metric depth with `Z = f * B / d`.
254pub fn disparity_to_depth(
255    disparity: ImageView<'_, f32, 1>,
256    focal_length: f64,
257    baseline: f64,
258) -> VisionResult<Image<f32, 1>> {
259    if !focal_length.is_finite() || focal_length <= 0.0 {
260        return Err(VisionError::InvalidParameter(
261            "disparity_to_depth focal_length must be finite and positive".into(),
262        ));
263    }
264    if !baseline.is_finite() || baseline <= 0.0 {
265        return Err(VisionError::InvalidParameter(
266            "disparity_to_depth baseline must be finite and positive".into(),
267        ));
268    }
269    let mut data = vec![0.0_f32; disparity.width() * disparity.height()];
270    for y in 0..disparity.height() {
271        for x in 0..disparity.width() {
272            let d = f64::from(disparity.get(x, y).expect("in-bounds")[0]);
273            data[y * disparity.width() + x] =
274                if d > 0.0 && d.is_finite() { (focal_length * baseline / d) as f32 } else { 0.0 };
275        }
276    }
277    Ok(Image::try_new(disparity.width(), disparity.height(), data)?)
278}
279
280/// Reprojects disparity into left-camera XYZ using rectified intrinsics.
281pub fn disparity_to_xyz(
282    disparity: ImageView<'_, f32, 1>,
283    camera: CameraMatrix3,
284    baseline: f64,
285) -> VisionResult<Image<f32, 3>> {
286    if !baseline.is_finite() || baseline <= 0.0 {
287        return Err(VisionError::InvalidParameter(
288            "disparity_to_xyz baseline must be finite and positive".into(),
289        ));
290    }
291    let fx = camera.matrix().m[0][0];
292    let fy = camera.matrix().m[1][1];
293    let cx = camera.matrix().m[0][2];
294    let cy = camera.matrix().m[1][2];
295    let mut data = vec![0.0_f32; disparity.width() * disparity.height() * 3];
296    for y in 0..disparity.height() {
297        for x in 0..disparity.width() {
298            let d = f64::from(disparity.get(x, y).expect("in-bounds")[0]);
299            let index = (y * disparity.width() + x) * 3;
300            if d > 0.0 && d.is_finite() {
301                let z = fx * baseline / d;
302                let xx = (x as f64 - cx) * z / fx;
303                let yy = (y as f64 - cy) * z / fy;
304                data[index] = xx as f32;
305                data[index + 1] = yy as f32;
306                data[index + 2] = z as f32;
307            }
308        }
309    }
310    Ok(Image::try_new(disparity.width(), disparity.height(), data)?)
311}
312
313fn build_rectify_maps(
314    camera: CameraMatrix3,
315    rotation: Mat3<f64>,
316    new_k: Mat3<f64>,
317    width: usize,
318    height: usize,
319) -> VisionResult<(Image<f32, 1>, Image<f32, 1>)> {
320    let mut map_x = vec![0.0_f32; width * height];
321    let mut map_y = vec![0.0_f32; width * height];
322    let new_inverse = invert_intrinsic(new_k)?;
323    let map_matrix = camera.matrix().mul_mat3(rotation.transpose()).mul_mat3(new_inverse);
324    for y in 0..height {
325        for x in 0..width {
326            let destination = Vec3::new(x as f64, y as f64, 1.0);
327            let source = map_matrix.mul_vec3(destination);
328            if source.z.abs() <= 1e-12 {
329                map_x[y * width + x] = -1.0;
330                map_y[y * width + x] = -1.0;
331            } else {
332                map_x[y * width + x] = (source.x / source.z) as f32;
333                map_y[y * width + x] = (source.y / source.z) as f32;
334            }
335        }
336    }
337    Ok((Image::try_new(width, height, map_x)?, Image::try_new(width, height, map_y)?))
338}
339
340fn invert_intrinsic(matrix: Mat3<f64>) -> VisionResult<Mat3<f64>> {
341    let fx = matrix.m[0][0];
342    let fy = matrix.m[1][1];
343    let cx = matrix.m[0][2];
344    let cy = matrix.m[1][2];
345    if fx.abs() <= f64::EPSILON || fy.abs() <= f64::EPSILON {
346        return Err(VisionError::InvalidParameter(
347            "rectified intrinsics must have non-zero focal lengths".into(),
348        ));
349    }
350    Ok(Mat3::from_rows([1.0 / fx, 0.0, -cx / fx], [0.0, 1.0 / fy, -cy / fy], [0.0, 0.0, 1.0]))
351}
352
353#[cfg(test)]
354mod tests {
355    use super::{
356        disparity_to_depth, stereo_block_match, stereo_rectify, StereoBmOptions, StereoRig,
357        INVALID_DISPARITY,
358    };
359    use crate::{CameraMatrix3, RelativePose};
360    use spatialrust_camera::CameraIntrinsics;
361    use spatialrust_image::Image;
362    use spatialrust_math::{Mat3, Vec3};
363
364    fn camera() -> CameraMatrix3 {
365        let intrinsics = CameraIntrinsics::try_new(400.0, 400.0, 80.0, 60.0, 160, 120).unwrap();
366        CameraMatrix3::from_intrinsics(intrinsics)
367    }
368
369    #[test]
370    fn fronto_parallel_stereo_recovers_plane_depth() {
371        let camera = camera();
372        let baseline = 0.1;
373        let pose =
374            RelativePose::try_new(Mat3::<f64>::identity(), Vec3::new(baseline, 0.0, 0.0)).unwrap();
375        let rig = StereoRig::try_new(camera, camera, pose).unwrap();
376        let maps = stereo_rectify(rig, 160, 120).unwrap();
377        assert!((maps.baseline() - baseline).abs() < 1e-12);
378
379        // Synthetic textured fronto-parallel plane at Z=2 with disparity = f*B/Z = 20.
380        let depth = 2.0;
381        let disparity = (400.0 * baseline / depth) as i32;
382        let width = 160;
383        let height = 120;
384        let mut left = vec![0u8; width * height];
385        let mut right = vec![0u8; width * height];
386        for y in 0..height {
387            for x in 0..width {
388                let value = (((x * 17 + y * 29) % 200) + 20) as u8;
389                left[y * width + x] = value;
390                let xr = x as i32 - disparity;
391                if (0..width as i32).contains(&xr) {
392                    right[y * width + xr as usize] = value;
393                }
394            }
395        }
396        let left = Image::<u8, 1>::try_new(width, height, left).unwrap();
397        let right = Image::<u8, 1>::try_new(width, height, right).unwrap();
398        let disparity_map = stereo_block_match(
399            left.view(),
400            right.view(),
401            StereoBmOptions {
402                window_size: 11,
403                min_disparity: 1,
404                num_disparities: 64,
405                uniqueness_ratio: 5.0,
406            },
407        )
408        .unwrap();
409        let center = disparity_map.get(80, 60).unwrap()[0];
410        assert!(
411            (center - disparity as f32).abs() <= 1.0,
412            "center disparity {center}, expected {disparity}"
413        );
414        assert_ne!(center, INVALID_DISPARITY);
415        let depth_map = disparity_to_depth(disparity_map.view(), 400.0, baseline).unwrap();
416        let recovered = depth_map.get(80, 60).unwrap()[0];
417        assert!((f64::from(recovered) - depth).abs() < 0.15);
418    }
419}