simple_icp/
voxel_hash_map.rs

1use nalgebra as na;
2use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
3use std::collections::HashMap;
4
5use crate::{
6    point3d::{self},
7    voxel_util::{self, na_vec_to_voxel},
8};
9
10pub type Voxel = na::Vector3<i32>;
11type VoxelPoints = Vec<point3d::Point3d>;
12
13pub struct VoxelHashMap {
14    voxel_size: f32,
15    max_distance: f64,
16    max_points_per_voxel: usize,
17    map: HashMap<na::Vector3<i32>, VoxelPoints>,
18    pub last_batch_points: VoxelPoints,
19}
20
21fn get_adjacent_voxels(voxel: &Voxel, adjacent_voxels: i32) -> Vec<Voxel> {
22    let mut voxel_neighborhood = Vec::<Voxel>::new();
23    for x in voxel.x - adjacent_voxels..voxel.x + adjacent_voxels + 1 {
24        for y in voxel.y - adjacent_voxels..voxel.y + adjacent_voxels + 1 {
25            for z in voxel.z - adjacent_voxels..voxel.z + adjacent_voxels + 1 {
26                voxel_neighborhood.push(Voxel::new(x, y, z));
27            }
28        }
29    }
30    voxel_neighborhood
31}
32
33impl VoxelHashMap {
34    pub fn default_values() -> VoxelHashMap {
35        VoxelHashMap {
36            voxel_size: 1.0,
37            max_distance: 100.0,
38            max_points_per_voxel: 20,
39            map: HashMap::new(),
40            last_batch_points: Vec::new(),
41        }
42    }
43
44    #[inline]
45    pub fn is_empty(&self) -> bool {
46        self.map.is_empty()
47    }
48
49    pub fn map_len(&self) -> usize {
50        self.map.iter().fold(0, |acc, (_, v)| acc + v.len())
51    }
52
53    fn update(&mut self, points: &VoxelPoints, current_origin: &na::Vector3<f64>) {
54        self.add_points(points);
55        self.remove_points_too_far(current_origin);
56    }
57
58    pub fn get_na_points(&self) -> Vec<na::Vector3<f64>> {
59        self.map
60            .values()
61            .map(|v| {
62                v.iter()
63                    .map(|p| p.to_na_vec_f64())
64                    .collect::<Vec<na::Vector3<f64>>>()
65            })
66            .reduce(|mut acc, mut e| {
67                acc.append(&mut e);
68                acc
69            })
70            .unwrap()
71    }
72
73    pub fn update_with_pose(
74        &mut self,
75        points: &[point3d::Point3d],
76        t_origin_current: &na::Isometry3<f64>,
77    ) {
78        let transformed_points: VoxelPoints = points
79            .par_iter()
80            .map(|pt| {
81                let transformed_na = t_origin_current * pt.to_na_point3_f64();
82                point3d::Point3d {
83                    x: transformed_na.x as f32,
84                    y: transformed_na.y as f32,
85                    z: transformed_na.z as f32,
86                    intensity: pt.intensity,
87                }
88            })
89            .collect();
90        self.update(&transformed_points, &t_origin_current.translation.vector);
91    }
92
93    fn add_points(&mut self, points: &VoxelPoints) {
94        let mut last_batch = Vec::new();
95        let map_resolution =
96            (self.voxel_size * self.voxel_size / self.max_points_per_voxel as f32).sqrt() as f64;
97        points.iter().for_each(|pt| {
98            let voxel = na_vec_to_voxel(&pt.to_na_vec_f64(), self.voxel_size as f64);
99            if let Some(voxel_points) = self.map.get_mut(&voxel) {
100                if voxel_points.len() >= self.max_points_per_voxel
101                    || voxel_points.iter().any(|vpt| {
102                        (vpt.to_na_vec_f64() - pt.to_na_vec_f64()).norm() < map_resolution
103                    })
104                {
105                    // return;
106                } else {
107                    voxel_points.push(*pt);
108                    last_batch.push(*pt);
109                }
110            } else {
111                self.map.insert(voxel, vec![*pt]);
112                last_batch.push(*pt);
113            }
114        });
115        self.last_batch_points = last_batch;
116    }
117    fn remove_points_too_far(&mut self, current_origin: &na::Vector3<f64>) {
118        let max_distance2 = self.max_distance * self.max_distance;
119        let keys_too_far: Vec<Voxel> = self
120            .map
121            .par_iter()
122            .filter_map(|(k, vps)| {
123                if (vps[0].to_na_vec_f64() - current_origin).norm_squared() >= max_distance2 {
124                    Some(k.to_owned())
125                } else {
126                    None
127                }
128            })
129            .collect();
130        keys_too_far.iter().for_each(|k| {
131            self.map.remove(k);
132        });
133    }
134
135    pub fn get_closest_neighbor(
136        &self,
137        point: &point3d::Point3d,
138    ) -> Option<(point3d::Point3d, f64)> {
139        let voxel = voxel_util::point_to_voxel(point, self.voxel_size);
140        let query_voxels = get_adjacent_voxels(&voxel, 1);
141        let point_na = point.to_na_vec_f64();
142        let neighbors: Vec<(point3d::Point3d, f64)> = query_voxels
143            .iter()
144            .filter_map(|query_voxel| {
145                if let Some(voxel_points) = self.map.get(query_voxel) {
146                    let neighbor = voxel_points
147                        .iter()
148                        .reduce(|acc, pt| {
149                            if (acc.to_na_vec_f64() - point_na).norm()
150                                < (pt.to_na_vec_f64() - point_na).norm()
151                            {
152                                acc
153                            } else {
154                                pt
155                            }
156                        })
157                        .unwrap();
158                    let distance = (neighbor.to_na_vec_f64() - point_na).norm();
159                    Some((*neighbor, distance))
160                } else {
161                    None
162                }
163            })
164            .collect();
165        neighbors
166            .iter()
167            .reduce(
168                |acc, neighbor| {
169                    if acc.1 < neighbor.1 {
170                        acc
171                    } else {
172                        neighbor
173                    }
174                },
175            )
176            .copied()
177    }
178}
simple_icp/voxel_hash_map.rs

simple_icp/
voxel_hash_map.rs
