1use crate::error::SpatialResult;
2use image::{DynamicImage, ImageBuffer, Rgb};
3use ndarray::Array2;
4
5pub fn generate_stereo_pair(
6 image: &DynamicImage,
7 depth: &Array2<f32>,
8 max_disparity: u32,
9) -> SpatialResult<(DynamicImage, DynamicImage)> {
10 let img_rgb = image.to_rgb8();
11 let width = img_rgb.width() as usize;
12 let height = img_rgb.height() as usize;
13
14 let mut right_rgb: ImageBuffer<Rgb<u8>, Vec<u8>> = ImageBuffer::new(width as u32, height as u32);
15 let mut depth_buffer = vec![f32::NEG_INFINITY; width * height];
16 let mut filled = vec![false; width * height];
17
18 for y in 0..height {
19 for x in 0..width {
20 let depth_val = get_depth_at(depth, x, y, width, height);
21 let disparity = (depth_val * max_disparity as f32).round() as i32;
22 let x_right = x as i32 - disparity;
23
24 if x_right >= 0 && x_right < width as i32 {
25 let idx = y * width + x_right as usize;
26 if depth_val > depth_buffer[idx] {
27 depth_buffer[idx] = depth_val;
28 filled[idx] = true;
29 if let Some(pixel) = img_rgb.get_pixel_checked(x as u32, y as u32) {
30 right_rgb.put_pixel(x_right as u32, y as u32, *pixel);
31 }
32 }
33 }
34 }
35 }
36
37 fill_disocclusions(&mut right_rgb, &filled, width, height);
38
39 let left_image = image.clone();
40 let right_image = DynamicImage::ImageRgb8(right_rgb);
41
42 Ok((left_image, right_image))
43}
44
45fn get_depth_at(
46 depth: &Array2<f32>,
47 x: usize,
48 y: usize,
49 img_width: usize,
50 img_height: usize,
51) -> f32 {
52 let (depth_height, depth_width) = depth.dim();
53
54 if depth_height == img_height && depth_width == img_width {
55 depth[[y, x]]
56 } else {
57 let scaled_x = (x as f32 * depth_width as f32 / img_width as f32)
58 .min(depth_width as f32 - 1.0) as usize;
59 let scaled_y = (y as f32 * depth_height as f32 / img_height as f32)
60 .min(depth_height as f32 - 1.0) as usize;
61
62 if scaled_y < depth_height && scaled_x < depth_width {
63 depth[[scaled_y, scaled_x]]
64 } else {
65 0.5
66 }
67 }
68}
69
70fn fill_disocclusions(
71 image: &mut ImageBuffer<Rgb<u8>, Vec<u8>>,
72 filled: &[bool],
73 width: usize,
74 height: usize,
75) {
76 let original = image.clone();
77
78 for y in 0..height {
79 for x in 0..width {
80 if filled[y * width + x] {
81 continue;
82 }
83 let mut left_pixel = None;
85 for lx in (0..x).rev() {
86 if filled[y * width + lx] {
87 left_pixel = Some(*original.get_pixel(lx as u32, y as u32));
88 break;
89 }
90 }
91 let mut right_pixel = None;
93 for rx in (x + 1)..width {
94 if filled[y * width + rx] {
95 right_pixel = Some(*original.get_pixel(rx as u32, y as u32));
96 break;
97 }
98 }
99
100 let fill = match (left_pixel, right_pixel) {
101 (Some(l), Some(_)) => l, (Some(l), None) => l,
103 (None, Some(r)) => r,
104 (None, None) => continue,
105 };
106 image.put_pixel(x as u32, y as u32, fill);
107 }
108 }
109}