iris/segmentation/
components.rs

1use crate::core::types::Rect;
2use crate::error::Result;
3use crate::image::Image;
4use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
5
6/// Component statistics output.
7#[derive(Clone, Debug, PartialEq)]
8pub struct ComponentStats {
9    pub label: usize,
10    pub bbox: Rect<usize>,
11    pub area: usize,
12    pub centroid: (f64, f64),
13}
14
15impl<B: Backend> Image<B> {
16    /// Identifies connected components in a binary image and calculates statistics for each component.
17    /// Returns labeled image tensor of shape [H, W] and stats list.
18    pub fn connected_components_with_stats(
19        &self,
20    ) -> Result<(Tensor<B, 2, Int>, Vec<ComponentStats>)> {
21        let gray = self.grayscale()?;
22        let dims = gray.tensor.dims();
23        let h = dims[1];
24        let w = dims[2];
25
26        let device = gray.tensor.device();
27        let tensor_data = gray.tensor.clone().into_data();
28        let flat_vals: Vec<f32> = tensor_data.iter::<f32>().collect();
29
30        let mut labels = vec![0usize; h * w];
31        let mut stats = Vec::new();
32        let mut current_label = 0;
33
34        // Depth-First Search Labeling
35        for y in 0..h {
36            for x in 0..w {
37                let idx = y * w + x;
38                if flat_vals[idx] > 0.5 && labels[idx] == 0 {
39                    current_label += 1;
40
41                    // Connected component statistics tracker
42                    let mut area = 0;
43                    let mut min_x = x;
44                    let mut max_x = x;
45                    let mut min_y = y;
46                    let mut max_y = y;
47                    let mut sum_x = 0;
48                    let mut sum_y = 0;
49
50                    // BFS queue
51                    let mut queue = std::collections::VecDeque::new();
52                    queue.push_back((x, y));
53                    labels[idx] = current_label;
54
55                    while let Some((cx, cy)) = queue.pop_front() {
56                        area += 1;
57                        min_x = min_x.min(cx);
58                        max_x = max_x.max(cx);
59                        min_y = min_y.min(cy);
60                        max_y = max_y.max(cy);
61                        sum_x += cx;
62                        sum_y += cy;
63
64                        // 4-connectivity neighbors
65                        let neighbors = [
66                            (cx as isize + 1, cy as isize),
67                            (cx as isize - 1, cy as isize),
68                            (cx as isize, cy as isize + 1),
69                            (cx as isize, cy as isize - 1),
70                        ];
71
72                        for &(nx, ny) in &neighbors {
73                            if nx >= 0 && nx < w as isize && ny >= 0 && ny < h as isize {
74                                let nidx = (ny as usize) * w + (nx as usize);
75                                if flat_vals[nidx] > 0.5 && labels[nidx] == 0 {
76                                    labels[nidx] = current_label;
77                                    queue.push_back((nx as usize, ny as usize));
78                                }
79                            }
80                        }
81                    }
82
83                    stats.push(ComponentStats {
84                        label: current_label,
85                        bbox: Rect::new(min_x, min_y, max_x - min_x + 1, max_y - min_y + 1),
86                        area,
87                        centroid: (sum_x as f64 / area as f64, sum_y as f64 / area as f64),
88                    });
89                }
90            }
91        }
92
93        // Reconstruct labeled tensor as integer tensor
94        let labels_i32: Vec<i32> = labels.iter().map(|&l| l as i32).collect();
95        let labels_data = TensorData::new(labels_i32, [h, w]);
96        let labels_tensor = Tensor::<B, 2, Int>::from_data(labels_data, &device);
97
98        Ok((labels_tensor, stats))
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use crate::test_helpers::{TestBackend, test_device};
106
107    #[test]
108    fn test_connected_components() {
109        let device = test_device();
110        // 5x5 image with two separated components
111        let mut flat_data = vec![0.0f32; 5 * 5];
112        flat_data[0] = 1.0;
113        flat_data[1] = 1.0;
114        flat_data[23] = 1.0;
115        flat_data[24] = 1.0;
116
117        let tensor =
118            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [1, 5, 5]), &device);
119        let img = Image::new(tensor);
120
121        let (labels, stats) = img.connected_components_with_stats().unwrap();
122        assert_eq!(labels.dims(), [5, 5]);
123        assert_eq!(stats.len(), 2);
124        assert_eq!(stats[0].area, 2);
125        assert_eq!(stats[1].area, 2);
126    }
127}
iris/segmentation/components.rs

iris/segmentation/
components.rs
