Skip to main content

spatial_maker/
stereo.rs

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			// scan left to find nearest filled pixel on same row
84			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			// scan right
92			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, // prefer left (background side in DIBR)
102				(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}