yscv_imgproc/ops/
stereo.rs

1use yscv_tensor::Tensor;
2
3use super::super::ImgProcError;
4use super::super::shape::hwc_shape;
5
6/// Configuration for stereo block matching.
7#[derive(Debug, Clone)]
8pub struct StereoConfig {
9    /// Number of disparity levels to search (must be > 0).
10    pub num_disparities: usize,
11    /// Side length of the matching block (must be odd and > 0).
12    pub block_size: usize,
13    /// Minimum disparity offset (can be 0).
14    pub min_disparity: usize,
15}
16
17impl Default for StereoConfig {
18    fn default() -> Self {
19        Self {
20            num_disparities: 64,
21            block_size: 9,
22            min_disparity: 0,
23        }
24    }
25}
26
27/// Compute a disparity map from a rectified stereo pair using block matching (SAD).
28///
29/// Both `left` and `right` must be single-channel images with shape `[H, W, 1]`.
30/// Returns an `[H, W, 1]` disparity map of type `f32`.  Pixels where no valid
31/// disparity could be found are set to `0.0`.
32pub fn stereo_block_matching(
33    left: &Tensor,
34    right: &Tensor,
35    config: &StereoConfig,
36) -> Result<Tensor, ImgProcError> {
37    let (h, w, c) = hwc_shape(left)?;
38    if c != 1 {
39        return Err(ImgProcError::InvalidChannelCount {
40            expected: 1,
41            got: c,
42        });
43    }
44    let (rh, rw, rc) = hwc_shape(right)?;
45    if rh != h || rw != w || rc != 1 {
46        return Err(ImgProcError::ShapeMismatch {
47            expected: vec![h, w, 1],
48            got: vec![rh, rw, rc],
49        });
50    }
51    if config.block_size == 0 || config.block_size.is_multiple_of(2) {
52        return Err(ImgProcError::InvalidBlockSize {
53            block_size: config.block_size,
54        });
55    }
56    if config.num_disparities == 0 {
57        return Err(ImgProcError::InvalidSize {
58            height: 0,
59            width: config.num_disparities,
60        });
61    }
62
63    let half = (config.block_size / 2) as isize;
64    let left_data = left.data();
65    let right_data = right.data();
66    let mut out = vec![0.0f32; h * w];
67
68    for y in 0..h {
69        for x in 0..w {
70            let mut best_sad = f32::MAX;
71            let mut best_d: usize = 0;
72
73            for d in config.min_disparity..config.min_disparity + config.num_disparities {
74                // If the block in the right image would go out of bounds, skip.
75                if (x as isize - d as isize) < half {
76                    continue;
77                }
78
79                let mut sad = 0.0f32;
80                for ky in -half..=half {
81                    let sy = y as isize + ky;
82                    if sy < 0 || sy >= h as isize {
83                        continue;
84                    }
85                    for kx in -half..=half {
86                        let lx = x as isize + kx;
87                        let rx = lx - d as isize;
88                        if lx < 0 || lx >= w as isize || rx < 0 || rx >= w as isize {
89                            continue;
90                        }
91                        let li = sy as usize * w + lx as usize;
92                        let ri = sy as usize * w + rx as usize;
93                        sad += (left_data[li] - right_data[ri]).abs();
94                    }
95                }
96
97                if sad < best_sad {
98                    best_sad = sad;
99                    best_d = d;
100                }
101            }
102
103            out[y * w + x] = best_d as f32;
104        }
105    }
106
107    Tensor::from_vec(vec![h, w, 1], out).map_err(Into::into)
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    fn make_gray(h: usize, w: usize, data: Vec<f32>) -> Tensor {
115        Tensor::from_vec(vec![h, w, 1], data).unwrap()
116    }
117
118    #[test]
119    fn test_zero_disparity_identical_images() {
120        let data = vec![1.0; 8 * 8];
121        let img = make_gray(8, 8, data);
122        let config = StereoConfig {
123            num_disparities: 4,
124            block_size: 3,
125            min_disparity: 0,
126        };
127        let disp = stereo_block_matching(&img, &img, &config).unwrap();
128        assert_eq!(disp.shape(), &[8, 8, 1]);
129        // Identical images → disparity should be 0 everywhere.
130        for &v in disp.data() {
131            assert!((v - 0.0).abs() < f32::EPSILON, "expected 0, got {v}");
132        }
133    }
134
135    #[test]
136    fn test_known_shift() {
137        // Create a 1-row image with a clear feature shifted by 2 pixels.
138        let h = 1;
139        let w = 16;
140        let left_data: Vec<f32> = (0..w).map(|i| if i == 8 { 100.0 } else { 0.0 }).collect();
141        let right_data: Vec<f32> = (0..w).map(|i| if i == 6 { 100.0 } else { 0.0 }).collect();
142        let left = make_gray(h, w, left_data);
143        let right = make_gray(h, w, right_data);
144        let config = StereoConfig {
145            num_disparities: 8,
146            block_size: 1,
147            min_disparity: 0,
148        };
149        let disp = stereo_block_matching(&left, &right, &config).unwrap();
150        // At x=8 the best match in right is at x=6 → disparity = 2.
151        assert!((disp.data()[8] - 2.0).abs() < f32::EPSILON);
152    }
153
154    #[test]
155    fn test_shape_mismatch_error() {
156        let a = make_gray(4, 4, vec![0.0; 16]);
157        let b = make_gray(4, 5, vec![0.0; 20]);
158        let config = StereoConfig::default();
159        assert!(stereo_block_matching(&a, &b, &config).is_err());
160    }
161
162    #[test]
163    fn test_invalid_block_size() {
164        let img = make_gray(4, 4, vec![0.0; 16]);
165        let config = StereoConfig {
166            num_disparities: 4,
167            block_size: 4, // even – invalid
168            min_disparity: 0,
169        };
170        assert!(stereo_block_matching(&img, &img, &config).is_err());
171    }
172}
yscv_imgproc/ops/stereo.rs

yscv_imgproc/ops/
stereo.rs
