1use spatialrust_image::{Image, ImageView};
4use spatialrust_math::{Mat3, Vec3};
5
6use crate::{CameraMatrix3, PixelComponent, RelativePose, VisionError, VisionResult};
7
8pub const INVALID_DISPARITY: f32 = -1.0;
10
11#[derive(Clone, Copy, Debug, PartialEq)]
13pub struct StereoRig {
14 left: CameraMatrix3,
15 right: CameraMatrix3,
16 pose: RelativePose,
17}
18
19impl StereoRig {
20 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 pub const fn left(self) -> CameraMatrix3 {
32 self.left
33 }
34
35 pub const fn right(self) -> CameraMatrix3 {
37 self.right
38 }
39
40 pub const fn pose(self) -> RelativePose {
42 self.pose
43 }
44
45 #[must_use]
47 pub fn baseline(self) -> f64 {
48 self.pose.translation().length()
49 }
50}
51
52#[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 pub const fn left_map_x(&self) -> &Image<f32, 1> {
67 &self.left_map_x
68 }
69
70 pub const fn left_map_y(&self) -> &Image<f32, 1> {
72 &self.left_map_y
73 }
74
75 pub const fn right_map_x(&self) -> &Image<f32, 1> {
77 &self.right_map_x
78 }
79
80 pub const fn right_map_y(&self) -> &Image<f32, 1> {
82 &self.right_map_y
83 }
84
85 pub const fn rectified_left(&self) -> CameraMatrix3 {
87 self.rectified_left
88 }
89
90 pub const fn rectified_right(&self) -> CameraMatrix3 {
92 self.rectified_right
93 }
94
95 pub const fn baseline(&self) -> f64 {
97 self.baseline
98 }
99}
100
101#[derive(Clone, Copy, Debug, PartialEq)]
103pub struct StereoBmOptions {
104 pub window_size: usize,
106 pub min_disparity: i32,
108 pub num_disparities: i32,
110 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 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
142pub 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
187pub 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
253pub 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
280pub 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 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}