1use crate::error::{Result, VisionError};
7use image::GrayImage;
8use scirs2_core::ndarray::{Array1, Array2, Array3};
9use std::time::{Duration, Instant};
10
11#[derive(Debug, Clone)]
30pub struct StereoMatchingParams {
31 pub min_disparity: i32,
33 pub max_disparity: i32,
35 pub block_size: usize,
37 pub cost_function: MatchingCostFunction,
39 pub sub_pixel_refinement: bool,
41 pub lr_consistency_threshold: f32,
43 pub enable_sgm: bool,
45 pub sgm_params: SgmParams,
47}
48
49#[derive(Debug, Clone, Copy)]
51pub enum MatchingCostFunction {
52 SAD,
54 SSD,
56 NCC,
58 Census,
60 MutualInformation,
62 Hybrid,
64}
65
66#[derive(Debug, Clone)]
68pub struct SgmParams {
69 pub p1: f32,
71 pub p2: f32,
73 pub eight_directions: bool,
75 pub uniqueness_ratio: f32,
77 pub speckle_size: usize,
79 pub speckle_range: f32,
81}
82
83#[derive(Debug, Clone)]
85pub struct DepthMapResult {
86 pub disparity_map: Array2<f32>,
88 pub confidence_map: Array2<f32>,
90 pub stats: DepthMapStats,
92}
93
94#[derive(Debug, Clone)]
96pub struct DepthMapStats {
97 pub valid_pixels: usize,
99 pub occluded_pixels: usize,
101 pub avg_matching_cost: f32,
103 pub processing_times: ProcessingTimes,
105}
106
107#[derive(Debug, Clone)]
109pub struct ProcessingTimes {
110 pub cost_computation: Duration,
112 pub cost_aggregation: Duration,
114 pub disparity_optimization: Duration,
116 pub post_processing: Duration,
118 pub total_time: Duration,
120}
121
122impl Default for StereoMatchingParams {
123 fn default() -> Self {
124 Self {
125 min_disparity: 0,
126 max_disparity: 64,
127 block_size: 9,
128 cost_function: MatchingCostFunction::SAD,
129 sub_pixel_refinement: true,
130 lr_consistency_threshold: 1.0,
131 enable_sgm: true,
132 sgm_params: SgmParams::default(),
133 }
134 }
135}
136
137impl Default for SgmParams {
138 fn default() -> Self {
139 Self {
140 p1: 8.0,
141 p2: 32.0,
142 eight_directions: true,
143 uniqueness_ratio: 0.15,
144 speckle_size: 100,
145 speckle_range: 2.0,
146 }
147 }
148}
149
150#[allow(dead_code)]
162pub fn compute_depth_map(
163 left_image: &GrayImage,
164 right_image: &GrayImage,
165 params: &StereoMatchingParams,
166) -> Result<DepthMapResult> {
167 let start_time = Instant::now();
168
169 let (left_width, left_height) = left_image.dimensions();
171 let (right_width, right_height) = right_image.dimensions();
172
173 if left_width != right_width || left_height != right_height {
174 return Err(VisionError::InvalidParameter(
175 "Stereo images must have the same dimensions".to_string(),
176 ));
177 }
178
179 let width = left_width as usize;
180 let _height = left_height as usize;
181
182 let left_array = image_to_array2(left_image);
184 let right_array = image_to_array2(right_image);
185
186 let mut processing_times = ProcessingTimes {
187 cost_computation: Duration::ZERO,
188 cost_aggregation: Duration::ZERO,
189 disparity_optimization: Duration::ZERO,
190 post_processing: Duration::ZERO,
191 total_time: Duration::ZERO,
192 };
193
194 let cost_start = Instant::now();
196 let cost_volume = compute_cost_volume(&left_array, &right_array, params)?;
197 processing_times.cost_computation = cost_start.elapsed();
198
199 let agg_start = Instant::now();
201 let aggregated_costs = if params.enable_sgm {
202 aggregate_costs_sgm(&cost_volume, ¶ms.sgm_params)?
203 } else {
204 cost_volume };
206 processing_times.cost_aggregation = agg_start.elapsed();
207
208 let opt_start = Instant::now();
210 let (mut disparity_map, confidence_map) = compute_disparity_wta(&aggregated_costs, params)?;
211 processing_times.disparity_optimization = opt_start.elapsed();
212
213 let post_start = Instant::now();
215
216 if params.lr_consistency_threshold > 0.0 {
218 let right_disparity = compute_right_disparity(&left_array, &right_array, params)?;
219 disparity_map = apply_lr_consistency_check(
220 &disparity_map,
221 &right_disparity,
222 params.lr_consistency_threshold,
223 );
224 }
225
226 if params.sub_pixel_refinement {
228 disparity_map = apply_subpixel_refinement(&disparity_map, &aggregated_costs)?;
229 }
230
231 disparity_map = fill_holes_and_filter(&disparity_map, ¶ms.sgm_params)?;
233
234 processing_times.post_processing = post_start.elapsed();
235 processing_times.total_time = start_time.elapsed();
236
237 let stats = compute_depth_map_stats(&disparity_map, &confidence_map, processing_times, params);
239
240 Ok(DepthMapResult {
241 disparity_map,
242 confidence_map,
243 stats,
244 })
245}
246
247#[allow(dead_code)]
249fn image_to_array2(image: &GrayImage) -> Array2<f32> {
250 let (width, height) = image.dimensions();
251 Array2::from_shape_fn((height as usize, width as usize), |(y, x)| {
252 image.get_pixel(x as u32, y as u32)[0] as f32 / 255.0
253 })
254}
255
256#[allow(dead_code)]
273fn compute_cost_volume(
274 left_image: &Array2<f32>,
275 right_image: &Array2<f32>,
276 params: &StereoMatchingParams,
277) -> Result<Array3<f32>> {
278 let (height, width) = left_image.dim();
279 let num_disparities = (params.max_disparity - params.min_disparity + 1) as usize;
280 let mut cost_volume = Array3::zeros((height, width, num_disparities));
281
282 let half_block = params.block_size / 2;
283
284 for d in 0..num_disparities {
286 let disparity = params.min_disparity + d as i32;
287
288 for y in half_block..height - half_block {
289 let mut x = half_block;
291 while x < width - half_block - 8 {
292 let batch_size = (width - half_block - x).min(8);
293 let mut costs = Vec::with_capacity(batch_size);
294
295 for i in 0..batch_size {
296 let xi = x + i;
297 let cost = match params.cost_function {
298 MatchingCostFunction::SAD => compute_sad_cost_simd(
299 left_image,
300 right_image,
301 xi,
302 y,
303 disparity,
304 params.block_size,
305 )?,
306 MatchingCostFunction::SSD => compute_ssd_cost_simd(
307 left_image,
308 right_image,
309 xi,
310 y,
311 disparity,
312 params.block_size,
313 )?,
314 MatchingCostFunction::NCC => compute_ncc_cost_simd(
315 left_image,
316 right_image,
317 xi,
318 y,
319 disparity,
320 params.block_size,
321 )?,
322 MatchingCostFunction::Census => compute_census_cost_simd(
323 left_image,
324 right_image,
325 xi,
326 y,
327 disparity,
328 params.block_size,
329 )?,
330 MatchingCostFunction::MutualInformation => compute_mi_cost_simd(
331 left_image,
332 right_image,
333 xi,
334 y,
335 disparity,
336 params.block_size,
337 )?,
338 MatchingCostFunction::Hybrid => compute_hybrid_cost_simd(
339 left_image,
340 right_image,
341 xi,
342 y,
343 disparity,
344 params.block_size,
345 )?,
346 };
347 costs.push(cost);
348 }
349
350 for (i, cost) in costs.iter().enumerate() {
352 if x + i < width - half_block {
353 cost_volume[[y, x + i, d]] = *cost;
354 }
355 }
356
357 x += batch_size;
358 }
359
360 while x < width - half_block {
362 let cost = match params.cost_function {
363 MatchingCostFunction::SAD => compute_sad_cost_simd(
364 left_image,
365 right_image,
366 x,
367 y,
368 disparity,
369 params.block_size,
370 )?,
371 _ => 0.0, };
373 cost_volume[[y, x, d]] = cost;
374 x += 1;
375 }
376 }
377 }
378
379 Ok(cost_volume)
380}
381
382#[allow(dead_code)]
384fn compute_sad_cost_simd(
385 left_image: &Array2<f32>,
386 right_image: &Array2<f32>,
387 x: usize,
388 y: usize,
389 disparity: i32,
390 block_size: usize,
391) -> Result<f32> {
392 use scirs2_core::simd_ops::SimdUnifiedOps;
393
394 let half_block = block_size / 2;
395 let right_x = x as i32 - disparity;
396
397 if right_x < half_block as i32 || right_x >= (right_image.dim().1 - half_block) as i32 {
398 return Ok(f32::INFINITY); }
400
401 let mut total_cost = 0.0f32;
402
403 for dy in -(half_block as i32)..=(half_block as i32) {
405 let ly = (y as i32 + dy) as usize;
406 let ry = ly;
407
408 let left_row: Vec<f32> = (-(half_block as i32)..=(half_block as i32))
410 .map(|dx| left_image[[ly, (x as i32 + dx) as usize]])
411 .collect();
412
413 let right_row: Vec<f32> = (-(half_block as i32)..=(half_block as i32))
414 .map(|dx| right_image[[ry, (right_x + dx) as usize]])
415 .collect();
416
417 let left_array = Array1::from_vec(left_row);
418 let right_array = Array1::from_vec(right_row);
419
420 let diff = f32::simd_sub(&left_array.view(), &right_array.view());
422 let abs_diff = f32::simd_abs(&diff.view());
423 let row_sum = f32::simd_sum(&abs_diff.view());
424
425 total_cost += row_sum;
426 }
427
428 Ok(total_cost)
429}
430
431#[allow(dead_code)]
433fn compute_ssd_cost_simd(
434 left_image: &Array2<f32>,
435 right_image: &Array2<f32>,
436 x: usize,
437 y: usize,
438 disparity: i32,
439 block_size: usize,
440) -> Result<f32> {
441 use scirs2_core::simd_ops::SimdUnifiedOps;
442
443 let half_block = block_size / 2;
444 let right_x = x as i32 - disparity;
445
446 if right_x < half_block as i32 || right_x >= (right_image.dim().1 - half_block) as i32 {
447 return Ok(f32::INFINITY);
448 }
449
450 let mut total_cost = 0.0f32;
451
452 for dy in -(half_block as i32)..=(half_block as i32) {
453 let ly = (y as i32 + dy) as usize;
454 let ry = ly;
455
456 let left_row: Vec<f32> = (-(half_block as i32)..=(half_block as i32))
457 .map(|dx| left_image[[ly, (x as i32 + dx) as usize]])
458 .collect();
459
460 let right_row: Vec<f32> = (-(half_block as i32)..=(half_block as i32))
461 .map(|dx| right_image[[ry, (right_x + dx) as usize]])
462 .collect();
463
464 let left_array = Array1::from_vec(left_row);
465 let right_array = Array1::from_vec(right_row);
466
467 let diff = f32::simd_sub(&left_array.view(), &right_array.view());
469 let sq_diff = f32::simd_mul(&diff.view(), &diff.view());
470 let row_sum = f32::simd_sum(&sq_diff.view());
471
472 total_cost += row_sum;
473 }
474
475 Ok(total_cost)
476}
477
478#[allow(dead_code)]
480fn compute_ncc_cost_simd(
481 left_image: &Array2<f32>,
482 right_image: &Array2<f32>,
483 x: usize,
484 y: usize,
485 disparity: i32,
486 block_size: usize,
487) -> Result<f32> {
488 use scirs2_core::simd_ops::SimdUnifiedOps;
489
490 let half_block = block_size / 2;
491 let right_x = x as i32 - disparity;
492
493 if right_x < half_block as i32 || right_x >= (right_image.dim().1 - half_block) as i32 {
494 return Ok(f32::INFINITY);
495 }
496
497 let mut left_block = Vec::new();
499 let mut right_block = Vec::new();
500
501 for dy in -(half_block as i32)..=(half_block as i32) {
502 for dx in -(half_block as i32)..=(half_block as i32) {
503 let ly = (y as i32 + dy) as usize;
504 let lx = (x as i32 + dx) as usize;
505 let ry = ly;
506 let rx = (right_x + dx) as usize;
507
508 left_block.push(left_image[[ly, lx]]);
509 right_block.push(right_image[[ry, rx]]);
510 }
511 }
512
513 let left_array = Array1::from_vec(left_block);
514 let right_array = Array1::from_vec(right_block);
515
516 let left_mean = f32::simd_sum(&left_array.view()) / left_array.len() as f32;
518 let right_mean = f32::simd_sum(&right_array.view()) / right_array.len() as f32;
519
520 let left_mean_array = Array1::from_elem(left_array.len(), left_mean);
521 let right_mean_array = Array1::from_elem(right_array.len(), right_mean);
522
523 let left_centered = f32::simd_sub(&left_array.view(), &left_mean_array.view());
524 let right_centered = f32::simd_sub(&right_array.view(), &right_mean_array.view());
525
526 let numerator =
527 f32::simd_sum(&f32::simd_mul(&left_centered.view(), &right_centered.view()).view());
528 let left_norm =
529 f32::simd_sum(&f32::simd_mul(&left_centered.view(), &left_centered.view()).view()).sqrt();
530 let right_norm =
531 f32::simd_sum(&f32::simd_mul(&right_centered.view(), &right_centered.view()).view()).sqrt();
532
533 let denominator = left_norm * right_norm;
534
535 if denominator > 1e-6 {
536 let ncc = numerator / denominator;
537 Ok(1.0 - ncc) } else {
539 Ok(f32::INFINITY)
540 }
541}
542
543#[allow(dead_code)]
545fn compute_census_cost_simd(
546 left_image: &Array2<f32>,
547 right_image: &Array2<f32>,
548 x: usize,
549 y: usize,
550 disparity: i32,
551 block_size: usize,
552) -> Result<f32> {
553 let half_block = block_size / 2;
554 let right_x = x as i32 - disparity;
555
556 if right_x < half_block as i32 || right_x >= (right_image.dim().1 - half_block) as i32 {
557 return Ok(f32::INFINITY);
558 }
559
560 let left_census = compute_census_transform(left_image, x, y, block_size);
562 let right_census = compute_census_transform(right_image, right_x as usize, y, block_size);
563
564 let hamming_distance = (left_census ^ right_census).count_ones() as f32;
566
567 Ok(hamming_distance)
568}
569
570#[allow(dead_code)]
572fn compute_census_transform(image: &Array2<f32>, x: usize, y: usize, block_size: usize) -> u32 {
573 let half_block = block_size / 2;
574 let center_value = image[[y, x]];
575 let mut census = 0u32;
576 let mut bit_index = 0;
577
578 for dy in -(half_block as i32)..=(half_block as i32) {
579 for dx in -(half_block as i32)..=(half_block as i32) {
580 if dx == 0 && dy == 0 {
581 continue; }
583
584 let py = (y as i32 + dy) as usize;
585 let px = (x as i32 + dx) as usize;
586
587 if image[[py, px]] < center_value {
588 census |= 1 << bit_index;
589 }
590 bit_index += 1;
591 }
592 }
593
594 census
595}
596
597#[allow(dead_code)]
599fn compute_mi_cost_simd(
600 left_image: &Array2<f32>,
601 right_image: &Array2<f32>,
602 x: usize,
603 y: usize,
604 disparity: i32,
605 block_size: usize,
606) -> Result<f32> {
607 compute_sad_cost_simd(left_image, right_image, x, y, disparity, block_size)
610}
611
612#[allow(dead_code)]
614fn compute_hybrid_cost_simd(
615 left_image: &Array2<f32>,
616 right_image: &Array2<f32>,
617 x: usize,
618 y: usize,
619 disparity: i32,
620 block_size: usize,
621) -> Result<f32> {
622 let sad_cost = compute_sad_cost_simd(left_image, right_image, x, y, disparity, block_size)?;
623 let census_cost =
624 compute_census_cost_simd(left_image, right_image, x, y, disparity, block_size)?;
625
626 Ok(0.7 * sad_cost + 0.3 * census_cost)
628}
629
630#[allow(dead_code)]
646fn aggregate_costs_sgm(cost_volume: &Array3<f32>, sgm_params: &SgmParams) -> Result<Array3<f32>> {
647 let (height, width, num_disparities) = cost_volume.dim();
648 let mut aggregated_costs = Array3::zeros((height, width, num_disparities));
649
650 let directions = if sgm_params.eight_directions {
652 vec![
653 (0, 1), (0, -1), (1, 0), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1), ]
662 } else {
663 vec![(0, 1), (0, -1), (1, 0), (-1, 0)]
664 };
665
666 for &(dy, dx) in &directions {
668 let direction_costs = aggregate_costs_direction(cost_volume, dy, dx, sgm_params)?;
669
670 for y in 0..height {
672 for x in 0..width {
673 for d in 0..num_disparities {
674 aggregated_costs[[y, x, d]] += direction_costs[[y, x, d]];
675 }
676 }
677 }
678 }
679
680 let num_dirs = directions.len() as f32;
682 aggregated_costs.mapv_inplace(|x| x / num_dirs);
683
684 Ok(aggregated_costs)
685}
686
687#[allow(dead_code)]
689fn aggregate_costs_direction(
690 cost_volume: &Array3<f32>,
691 dy: i32,
692 dx: i32,
693 sgm_params: &SgmParams,
694) -> Result<Array3<f32>> {
695 let (height, width, _num_disparities) = cost_volume.dim();
696 let mut direction_costs = cost_volume.clone();
697
698 match dy.cmp(&0) {
700 std::cmp::Ordering::Greater => {
701 for y in 1..height {
703 for x in 0..width {
704 let prev_y = (y as i32 - dy) as usize;
705 let prev_x = if dx != 0 {
706 let px = x as i32 - dx;
707 if px >= 0 && px < width as i32 {
708 px as usize
709 } else {
710 continue;
711 }
712 } else {
713 x
714 };
715
716 if prev_y < height && prev_x < width {
717 aggregate_pixel_costs(
718 &mut direction_costs,
719 y,
720 x,
721 prev_y,
722 prev_x,
723 sgm_params,
724 );
725 }
726 }
727 }
728 }
729 std::cmp::Ordering::Less => {
730 for y in (0..height - 1).rev() {
732 for x in 0..width {
733 let prev_y = (y as i32 - dy) as usize;
734 let prev_x = if dx != 0 {
735 let px = x as i32 - dx;
736 if px >= 0 && px < width as i32 {
737 px as usize
738 } else {
739 continue;
740 }
741 } else {
742 x
743 };
744
745 if prev_y < height && prev_x < width {
746 aggregate_pixel_costs(
747 &mut direction_costs,
748 y,
749 x,
750 prev_y,
751 prev_x,
752 sgm_params,
753 );
754 }
755 }
756 }
757 }
758 std::cmp::Ordering::Equal => {
759 let x_range: Box<dyn Iterator<Item = usize>> = if dx > 0 {
761 Box::new(1..width)
762 } else {
763 Box::new((0..width - 1).rev())
764 };
765
766 for x in x_range {
767 for y in 0..height {
768 let prev_x = (x as i32 - dx) as usize;
769 if prev_x < width {
770 aggregate_pixel_costs(&mut direction_costs, y, x, y, prev_x, sgm_params);
771 }
772 }
773 }
774 }
775 }
776
777 Ok(direction_costs)
778}
779
780#[allow(dead_code)]
782fn aggregate_pixel_costs(
783 direction_costs: &mut Array3<f32>,
784 y: usize,
785 x: usize,
786 prev_y: usize,
787 prev_x: usize,
788 sgm_params: &SgmParams,
789) {
790 let num_disparities = direction_costs.dim().2;
791
792 for d in 0..num_disparities {
793 let raw_cost = direction_costs[[y, x, d]];
794
795 let mut min_aggregated_cost = f32::INFINITY;
797
798 for prev_d in 0..num_disparities {
799 let prev_cost = direction_costs[[prev_y, prev_x, prev_d]];
800
801 let smoothness_penalty = if d == prev_d {
802 0.0 } else if (d as i32 - prev_d as i32).abs() == 1 {
804 sgm_params.p1 } else {
806 sgm_params.p2 };
808
809 let aggregated_cost = prev_cost + smoothness_penalty;
810 if aggregated_cost < min_aggregated_cost {
811 min_aggregated_cost = aggregated_cost;
812 }
813 }
814
815 direction_costs[[y, x, d]] = raw_cost + min_aggregated_cost;
816 }
817}
818
819#[allow(dead_code)]
821fn compute_disparity_wta(
822 cost_volume: &Array3<f32>,
823 params: &StereoMatchingParams,
824) -> Result<(Array2<f32>, Array2<f32>)> {
825 let (height, width, num_disparities) = cost_volume.dim();
826 let mut disparity_map = Array2::zeros((height, width));
827 let mut confidence_map = Array2::zeros((height, width));
828
829 for y in 0..height {
830 for x in 0..width {
831 let mut min_cost = f32::INFINITY;
832 let mut best_disparity = 0;
833 let mut second_min_cost = f32::INFINITY;
834
835 for d in 0..num_disparities {
837 let cost = cost_volume[[y, x, d]];
838 if cost < min_cost {
839 second_min_cost = min_cost;
840 min_cost = cost;
841 best_disparity = d;
842 } else if cost < second_min_cost {
843 second_min_cost = cost;
844 }
845 }
846
847 disparity_map[[y, x]] = (params.min_disparity + best_disparity as i32) as f32;
848
849 let confidence = if second_min_cost > min_cost + 1e-6 {
851 1.0 - min_cost / second_min_cost
852 } else {
853 0.0
854 };
855
856 confidence_map[[y, x]] = confidence.clamp(0.0, 1.0);
857 }
858 }
859
860 Ok((disparity_map, confidence_map))
861}
862
863#[allow(dead_code)]
865fn compute_right_disparity(
866 left_image: &Array2<f32>,
867 right_image: &Array2<f32>,
868 params: &StereoMatchingParams,
869) -> Result<Array2<f32>> {
870 let mut right_params = params.clone();
872 right_params.min_disparity = -params.max_disparity;
873 right_params.max_disparity = -params.min_disparity;
874
875 let cost_volume = compute_cost_volume(right_image, left_image, &right_params)?;
876 let (right_disparity, _) = compute_disparity_wta(&cost_volume, &right_params)?;
877
878 Ok(right_disparity.mapv(|d| -d))
880}
881
882#[allow(dead_code)]
884fn apply_lr_consistency_check(
885 left_disparity: &Array2<f32>,
886 right_disparity: &Array2<f32>,
887 threshold: f32,
888) -> Array2<f32> {
889 let (height, width) = left_disparity.dim();
890 let mut consistent_disparity = left_disparity.clone();
891
892 for y in 0..height {
893 for x in 0..width {
894 let left_d = left_disparity[[y, x]];
895 let right_x = (x as f32 - left_d).round() as i32;
896
897 if right_x >= 0 && right_x < width as i32 {
898 let right_d = right_disparity[[y, right_x as usize]];
899
900 if (left_d - right_d).abs() > threshold {
901 consistent_disparity[[y, x]] = f32::NAN; }
903 } else {
904 consistent_disparity[[y, x]] = f32::NAN;
905 }
906 }
907 }
908
909 consistent_disparity
910}
911
912#[allow(dead_code)]
914fn apply_subpixel_refinement(
915 disparity_map: &Array2<f32>,
916 cost_volume: &Array3<f32>,
917) -> Result<Array2<f32>> {
918 let (height, width) = disparity_map.dim();
919 let mut refined_disparity = disparity_map.clone();
920
921 for y in 0..height {
922 for x in 0..width {
923 let d = disparity_map[[y, x]] as usize;
924
925 if d == 0 || d >= cost_volume.dim().2 - 1 {
927 continue;
928 }
929
930 let c_prev = cost_volume[[y, x, d - 1]];
932 let c_curr = cost_volume[[y, x, d]];
933 let c_next = cost_volume[[y, x, d + 1]];
934
935 let denominator = 2.0 * (c_prev - 2.0 * c_curr + c_next);
936 if denominator.abs() > 1e-6 {
937 let offset = (c_prev - c_next) / denominator;
938 refined_disparity[[y, x]] = d as f32 + offset;
939 }
940 }
941 }
942
943 Ok(refined_disparity)
944}
945
946#[allow(dead_code)]
948fn fill_holes_and_filter(
949 disparity_map: &Array2<f32>,
950 sgm_params: &SgmParams,
951) -> Result<Array2<f32>> {
952 let (height, width) = disparity_map.dim();
953 let mut filtered_disparity = disparity_map.clone();
954
955 for y in 0..height {
957 for x in 0..width {
958 if disparity_map[[y, x]].is_nan() {
959 let mut found = false;
961 for radius in 1..=10 {
962 let mut sum = 0.0;
963 let mut count = 0;
964
965 for dy in -radius..=radius {
966 for dx in -radius..=radius {
967 let ny = y as i32 + dy;
968 let nx = x as i32 + dx;
969
970 if ny >= 0 && ny < height as i32 && nx >= 0 && nx < width as i32 {
971 let val = disparity_map[[ny as usize, nx as usize]];
972 if !val.is_nan() {
973 sum += val;
974 count += 1;
975 }
976 }
977 }
978 }
979
980 if count > 0 {
981 filtered_disparity[[y, x]] = sum / count as f32;
982 found = true;
983 break;
984 }
985 }
986
987 if !found {
988 filtered_disparity[[y, x]] = 0.0;
989 }
990 }
991 }
992 }
993
994 filtered_disparity = apply_median_filter(&filtered_disparity, 3)?;
996
997 filtered_disparity = apply_speckle_filter(&filtered_disparity, sgm_params)?;
999
1000 Ok(filtered_disparity)
1001}
1002
1003#[allow(dead_code)]
1005fn apply_median_filter(disparity_map: &Array2<f32>, window_size: usize) -> Result<Array2<f32>> {
1006 let (height, width) = disparity_map.dim();
1007 let mut filtered = disparity_map.clone();
1008 let half_window = window_size / 2;
1009
1010 for y in half_window..height - half_window {
1011 for x in half_window..width - half_window {
1012 let mut values = Vec::new();
1013
1014 for dy in -(half_window as i32)..=(half_window as i32) {
1015 for dx in -(half_window as i32)..=(half_window as i32) {
1016 let val = disparity_map[[(y as i32 + dy) as usize, (x as i32 + dx) as usize]];
1017 if !val.is_nan() {
1018 values.push(val);
1019 }
1020 }
1021 }
1022
1023 if !values.is_empty() {
1024 values.sort_by(|a, b| a.partial_cmp(b).expect("Operation failed"));
1025 filtered[[y, x]] = values[values.len() / 2];
1026 }
1027 }
1028 }
1029
1030 Ok(filtered)
1031}
1032
1033#[allow(dead_code)]
1035fn apply_speckle_filter(
1036 disparity_map: &Array2<f32>,
1037 sgm_params: &SgmParams,
1038) -> Result<Array2<f32>> {
1039 let (height, width) = disparity_map.dim();
1040 let mut filtered = disparity_map.clone();
1041 let mut visited = Array2::from_elem((height, width), false);
1042
1043 for y in 0..height {
1044 for x in 0..width {
1045 if !visited[[y, x]] && !disparity_map[[y, x]].is_nan() {
1046 let region_size = flood_fill_region_size(
1047 disparity_map,
1048 &mut visited,
1049 x,
1050 y,
1051 disparity_map[[y, x]],
1052 sgm_params.speckle_range,
1053 );
1054
1055 if region_size < sgm_params.speckle_size {
1056 flood_fill_mark_invalid(
1058 &mut filtered,
1059 x,
1060 y,
1061 disparity_map[[y, x]],
1062 sgm_params.speckle_range,
1063 );
1064 }
1065 }
1066 }
1067 }
1068
1069 Ok(filtered)
1070}
1071
1072#[allow(dead_code)]
1074fn flood_fill_region_size(
1075 disparity_map: &Array2<f32>,
1076 visited: &mut Array2<bool>,
1077 start_x: usize,
1078 start_y: usize,
1079 target_disparity: f32,
1080 range: f32,
1081) -> usize {
1082 let (height, width) = disparity_map.dim();
1083 let mut stack = vec![(start_x, start_y)];
1084 let mut region_size = 0;
1085
1086 while let Some((x, y)) = stack.pop() {
1087 if x >= width || y >= height || visited[[y, x]] {
1088 continue;
1089 }
1090
1091 let disparity = disparity_map[[y, x]];
1092 if disparity.is_nan() || (disparity - target_disparity).abs() > range {
1093 continue;
1094 }
1095
1096 visited[[y, x]] = true;
1097 region_size += 1;
1098
1099 if x > 0 {
1101 stack.push((x - 1, y));
1102 }
1103 if x < width - 1 {
1104 stack.push((x + 1, y));
1105 }
1106 if y > 0 {
1107 stack.push((x, y - 1));
1108 }
1109 if y < height - 1 {
1110 stack.push((x, y + 1));
1111 }
1112 }
1113
1114 region_size
1115}
1116
1117#[allow(dead_code)]
1119fn flood_fill_mark_invalid(
1120 disparity_map: &mut Array2<f32>,
1121 start_x: usize,
1122 start_y: usize,
1123 target_disparity: f32,
1124 range: f32,
1125) {
1126 let (height, width) = disparity_map.dim();
1127 let mut stack = vec![(start_x, start_y)];
1128
1129 while let Some((x, y)) = stack.pop() {
1130 if x >= width || y >= height {
1131 continue;
1132 }
1133
1134 let disparity = disparity_map[[y, x]];
1135 if disparity.is_nan() || (disparity - target_disparity).abs() > range {
1136 continue;
1137 }
1138
1139 disparity_map[[y, x]] = f32::NAN;
1140
1141 if x > 0 {
1143 stack.push((x - 1, y));
1144 }
1145 if x < width - 1 {
1146 stack.push((x + 1, y));
1147 }
1148 if y > 0 {
1149 stack.push((x, y - 1));
1150 }
1151 if y < height - 1 {
1152 stack.push((x, y + 1));
1153 }
1154 }
1155}
1156
1157#[allow(dead_code)]
1159fn compute_depth_map_stats(
1160 disparity_map: &Array2<f32>,
1161 confidence_map: &Array2<f32>,
1162 processing_times: ProcessingTimes,
1163 _params: &StereoMatchingParams,
1164) -> DepthMapStats {
1165 let total_pixels = disparity_map.len();
1166 let valid_pixels = disparity_map.iter().filter(|&&d| !d.is_nan()).count();
1167 let occluded_pixels = total_pixels - valid_pixels;
1168
1169 let avg_matching_cost =
1170 confidence_map.iter().filter(|&&c| !c.is_nan()).sum::<f32>() / valid_pixels.max(1) as f32;
1171
1172 DepthMapStats {
1173 valid_pixels,
1174 occluded_pixels,
1175 avg_matching_cost,
1176 processing_times,
1177 }
1178}
1179
1180#[allow(dead_code)]
1192pub fn disparity_to_depth(
1193 disparity_map: &Array2<f32>,
1194 focal_length: f32,
1195 baseline: f32,
1196) -> Array2<f32> {
1197 disparity_map.mapv(|d| {
1198 if d > 0.0 && !d.is_nan() {
1199 (focal_length * baseline) / d
1200 } else {
1201 f32::NAN
1202 }
1203 })
1204}