ruvector_robotics/perception/
clustering.rs1use crate::bridge::{Point3D, PointCloud};
13use std::collections::HashMap;
14
15pub fn cluster_point_cloud(cloud: &PointCloud, cell_size: f64) -> Vec<Vec<Point3D>> {
21 if cloud.points.is_empty() || cell_size <= 0.0 {
22 return Vec::new();
23 }
24
25 let mut cell_map: HashMap<(i64, i64, i64), Vec<usize>> = HashMap::new();
27 for (idx, p) in cloud.points.iter().enumerate() {
28 let key = cell_key(p, cell_size);
29 cell_map.entry(key).or_default().push(idx);
30 }
31
32 let cells: Vec<(i64, i64, i64)> = cell_map.keys().copied().collect();
34 let cell_count = cells.len();
35 let cell_idx: HashMap<(i64, i64, i64), usize> = cells
36 .iter()
37 .enumerate()
38 .map(|(i, &k)| (k, i))
39 .collect();
40
41 let mut parent: Vec<usize> = (0..cell_count).collect();
42 let mut rank: Vec<u8> = vec![0; cell_count];
43
44 for &(cx, cy, cz) in &cells {
45 let a = cell_idx[&(cx, cy, cz)];
46 for dx in -1..=1_i64 {
47 for dy in -1..=1_i64 {
48 for dz in -1..=1_i64 {
49 let neighbor = (cx + dx, cy + dy, cz + dz);
50 if let Some(&b) = cell_idx.get(&neighbor) {
51 uf_union(&mut parent, &mut rank, a, b);
52 }
53 }
54 }
55 }
56 }
57
58 let mut groups: HashMap<usize, Vec<Point3D>> = HashMap::new();
60 for (key, point_indices) in &cell_map {
61 let ci = cell_idx[key];
62 let root = uf_find(&mut parent, ci);
63 let entry = groups.entry(root).or_default();
64 for &pi in point_indices {
65 entry.push(cloud.points[pi]);
66 }
67 }
68
69 groups.into_values().collect()
70}
71
72#[inline]
74fn cell_key(p: &Point3D, cell_size: f64) -> (i64, i64, i64) {
75 (
76 (p.x as f64 / cell_size).floor() as i64,
77 (p.y as f64 / cell_size).floor() as i64,
78 (p.z as f64 / cell_size).floor() as i64,
79 )
80}
81
82#[inline]
84fn uf_find(parent: &mut [usize], mut i: usize) -> usize {
85 while parent[i] != i {
86 parent[i] = parent[parent[i]];
87 i = parent[i];
88 }
89 i
90}
91
92#[inline]
94fn uf_union(parent: &mut [usize], rank: &mut [u8], a: usize, b: usize) {
95 let ra = uf_find(parent, a);
96 let rb = uf_find(parent, b);
97 if ra != rb {
98 match rank[ra].cmp(&rank[rb]) {
99 std::cmp::Ordering::Less => parent[ra] = rb,
100 std::cmp::Ordering::Greater => parent[rb] = ra,
101 std::cmp::Ordering::Equal => {
102 parent[rb] = ra;
103 rank[ra] += 1;
104 }
105 }
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112
113 fn make_cloud(pts: &[[f32; 3]]) -> PointCloud {
114 let points: Vec<Point3D> = pts.iter().map(|a| Point3D::new(a[0], a[1], a[2])).collect();
115 PointCloud::new(points, 0)
116 }
117
118 #[test]
119 fn test_empty_cloud() {
120 let cloud = PointCloud::default();
121 let clusters = cluster_point_cloud(&cloud, 1.0);
122 assert!(clusters.is_empty());
123 }
124
125 #[test]
126 fn test_single_cluster() {
127 let cloud = make_cloud(&[
128 [1.0, 1.0, 0.0],
129 [1.1, 1.0, 0.0],
130 [1.0, 1.1, 0.0],
131 ]);
132 let clusters = cluster_point_cloud(&cloud, 0.5);
133 assert_eq!(clusters.len(), 1);
134 assert_eq!(clusters[0].len(), 3);
135 }
136
137 #[test]
138 fn test_two_clusters() {
139 let cloud = make_cloud(&[
140 [0.0, 0.0, 0.0],
141 [0.1, 0.0, 0.0],
142 [10.0, 10.0, 0.0],
143 [10.1, 10.0, 0.0],
144 ]);
145 let clusters = cluster_point_cloud(&cloud, 0.5);
146 assert_eq!(clusters.len(), 2);
147 }
148
149 #[test]
150 fn test_negative_coordinates() {
151 let cloud = make_cloud(&[
152 [-1.0, -1.0, 0.0],
153 [-0.9, -1.0, 0.0],
154 [1.0, 1.0, 0.0],
155 ]);
156 let clusters = cluster_point_cloud(&cloud, 0.5);
157 assert_eq!(clusters.len(), 2);
158 }
159}