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
16 let bg = Rgb([64u8, 64u8, 64u8]);
17 for pixel in right_rgb.pixels_mut() {
18 *pixel = bg;
19 }
20
21 for y in 0..height {
22 for x in 0..width {
23 let depth_val = get_depth_at(depth, x, y, width, height);
24 let disparity = (depth_val * max_disparity as f32).round() as i32;
25 let x_right = x as i32 - disparity;
26
27 if x_right >= 0 && x_right < width as i32 {
28 if let Some(pixel) = img_rgb.get_pixel_checked(x as u32, y as u32) {
29 right_rgb.put_pixel(x_right as u32, y as u32, *pixel);
30 }
31 }
32 }
33 }
34
35 fill_disocclusions(&mut right_rgb);
36
37 let left_image = image.clone();
38 let right_image = DynamicImage::ImageRgb8(right_rgb);
39
40 Ok((left_image, right_image))
41}
42
43fn get_depth_at(
44 depth: &Array2<f32>,
45 x: usize,
46 y: usize,
47 img_width: usize,
48 img_height: usize,
49) -> f32 {
50 let (depth_height, depth_width) = depth.dim();
51
52 if depth_height == img_height && depth_width == img_width {
53 depth[[y, x]]
54 } else {
55 let scaled_x = (x as f32 * depth_width as f32 / img_width as f32)
56 .min(depth_width as f32 - 1.0) as usize;
57 let scaled_y = (y as f32 * depth_height as f32 / img_height as f32)
58 .min(depth_height as f32 - 1.0) as usize;
59
60 if scaled_y < depth_height && scaled_x < depth_width {
61 depth[[scaled_y, scaled_x]]
62 } else {
63 0.5
64 }
65 }
66}
67
68fn fill_disocclusions(image: &mut ImageBuffer<Rgb<u8>, Vec<u8>>) {
69 let width = image.width() as usize;
70 let height = image.height() as usize;
71 let bg = Rgb([64u8, 64u8, 64u8]);
72
73 let original = image.clone();
74
75 for y in 0..height {
76 for x in 0..width {
77 let pixel = original.get_pixel(x as u32, y as u32);
78 if pixel[0] == bg[0] && pixel[1] == bg[1] && pixel[2] == bg[2] {
79 if let Some(nearest) = find_nearest_valid(&original, x, y, bg) {
80 image.put_pixel(x as u32, y as u32, nearest);
81 }
82 }
83 }
84 }
85}
86
87fn find_nearest_valid(
88 image: &ImageBuffer<Rgb<u8>, Vec<u8>>,
89 cx: usize,
90 cy: usize,
91 bg: Rgb<u8>,
92) -> Option<Rgb<u8>> {
93 let width = image.width() as usize;
94 let height = image.height() as usize;
95
96 for radius in 1..=20 {
97 for dy in -(radius as i32)..=(radius as i32) {
98 for dx in -(radius as i32)..=(radius as i32) {
99 if dx.abs() != radius as i32 && dy.abs() != radius as i32 {
100 continue;
101 }
102 let nx = (cx as i32 + dx) as usize;
103 let ny = (cy as i32 + dy) as usize;
104 if nx < width && ny < height {
105 let pixel = image.get_pixel(nx as u32, ny as u32);
106 if pixel[0] != bg[0] || pixel[1] != bg[1] || pixel[2] != bg[2] {
107 return Some(*pixel);
108 }
109 }
110 }
111 }
112 }
113
114 None
115}