Move NN search back to the VoxelHashMap (#388)

* Now we can search neighbors from the voxel hash * Integrate Nacho's Comment Co-authored-by: Ignacio Vizzo <ignacio@dexory.com> --------- Co-authored-by: Ignacio Vizzo <ignacio@dexory.com>
PRBonn · Aug 8, 2024 · 624e46e · 624e46e
1 parent f27ced4
commit 624e46e
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 42 deletions.
diff --git a/cpp/kiss_icp/core/Registration.cpp b/cpp/kiss_icp/core/Registration.cpp
@@ -55,41 +55,6 @@ void TransformPoints(const Sophus::SE3d &T, std::vector<Eigen::Vector3d> &points
                    [&](const auto &point) { return T * point; });
 }
 
-using Voxel = kiss_icp::Voxel;
-std::vector<Voxel> GetAdjacentVoxels(const Voxel &voxel, int adjacent_voxels = 1) {
-    std::vector<Voxel> voxel_neighborhood;
-    for (int i = voxel.x() - adjacent_voxels; i < voxel.x() + adjacent_voxels + 1; ++i) {
-        for (int j = voxel.y() - adjacent_voxels; j < voxel.y() + adjacent_voxels + 1; ++j) {
-            for (int k = voxel.z() - adjacent_voxels; k < voxel.z() + adjacent_voxels + 1; ++k) {
-                voxel_neighborhood.emplace_back(i, j, k);
-            }
-        }
-    }
-    return voxel_neighborhood;
-}
-
-std::tuple<Eigen::Vector3d, double> GetClosestNeighbor(const Eigen::Vector3d &point,
-                                                       const kiss_icp::VoxelHashMap &voxel_map) {
-    // Convert the point to voxel coordinates
-    const auto &voxel = kiss_icp::PointToVoxel(point, voxel_map.voxel_size_);
-    // Get nearby voxels on the map
-    const auto &query_voxels = GetAdjacentVoxels(voxel);
-    // Extract the points contained within the neighborhood voxels
-    const auto &neighbors = voxel_map.GetPoints(query_voxels);
-
-    // Find the nearest neighbor
-    Eigen::Vector3d closest_neighbor = Eigen::Vector3d::Zero();
-    double closest_distance = std::numeric_limits<double>::max();
-    std::for_each(neighbors.cbegin(), neighbors.cend(), [&](const auto &neighbor) {
-        double distance = (neighbor - point).norm();
-        if (distance < closest_distance) {
-            closest_neighbor = neighbor;
-            closest_distance = distance;
-        }
-    });
-    return std::make_tuple(closest_neighbor, closest_distance);
-}
-
 Correspondences DataAssociation(const std::vector<Eigen::Vector3d> &points,
                                 const kiss_icp::VoxelHashMap &voxel_map,
                                 const double max_correspondance_distance) {
@@ -105,7 +70,7 @@ Correspondences DataAssociation(const std::vector<Eigen::Vector3d> &points,
         [&](const tbb::blocked_range<points_iterator> &r, Correspondences res) -> Correspondences {
             res.reserve(r.size());
             std::for_each(r.begin(), r.end(), [&](const auto &point) {
-                const auto &[closest_neighbor, distance] = GetClosestNeighbor(point, voxel_map);
+                const auto &[closest_neighbor, distance] = voxel_map.GetClosestNeighbor(point);
                 if (distance < max_correspondance_distance) {
                     res.emplace_back(point, closest_neighbor);
                 }

diff --git a/cpp/kiss_icp/core/VoxelHashMap.cpp b/cpp/kiss_icp/core/VoxelHashMap.cpp
@@ -29,14 +29,27 @@
 
 #include "VoxelUtils.hpp"
 
-namespace kiss_icp {
+namespace {
+using kiss_icp::Voxel;
 
-std::vector<Eigen::Vector3d> VoxelHashMap::GetPoints(const std::vector<Voxel> &query_voxels) const {
+std::vector<Voxel> GetAdjacentVoxels(const Voxel &voxel, int adjacent_voxels = 1) {
+    std::vector<Voxel> voxel_neighborhood;
+    for (int i = voxel.x() - adjacent_voxels; i < voxel.x() + adjacent_voxels + 1; ++i) {
+        for (int j = voxel.y() - adjacent_voxels; j < voxel.y() + adjacent_voxels + 1; ++j) {
+            for (int k = voxel.z() - adjacent_voxels; k < voxel.z() + adjacent_voxels + 1; ++k) {
+                voxel_neighborhood.emplace_back(i, j, k);
+            }
+        }
+    }
+    return voxel_neighborhood;
+}
+std::vector<Eigen::Vector3d> GetPoints(const std::vector<Voxel> &query_voxels,
+                                       const kiss_icp::VoxelHashMap &voxel_map) {
     std::vector<Eigen::Vector3d> points;
-    points.reserve(query_voxels.size() * static_cast<size_t>(max_points_per_voxel_));
+    points.reserve(query_voxels.size() * static_cast<size_t>(voxel_map.max_points_per_voxel_));
     std::for_each(query_voxels.cbegin(), query_voxels.cend(), [&](const auto &query) {
-        auto search = map_.find(query);
-        if (search != map_.end()) {
+        auto search = voxel_map.map_.find(query);
+        if (search != voxel_map.map_.end()) {
             const auto &voxel_points = search.value();
             points.insert(points.end(), voxel_points.cbegin(), voxel_points.cend());
         }
@@ -45,6 +58,32 @@ std::vector<Eigen::Vector3d> VoxelHashMap::GetPoints(const std::vector<Voxel> &q
     return points;
 }
 
+}  // namespace
+
+namespace kiss_icp {
+
+std::tuple<Eigen::Vector3d, double> VoxelHashMap::GetClosestNeighbor(
+    const Eigen::Vector3d &point) const {
+    // Convert the point to voxel coordinates
+    const auto &voxel = PointToVoxel(point, voxel_size_);
+    // Get nearby voxels on the map
+    const auto &query_voxels = GetAdjacentVoxels(voxel);
+    // Extract the points contained within the neighborhood voxels
+    const auto &neighbors = GetPoints(query_voxels, *this);
+
+    // Find the nearest neighbor
+    Eigen::Vector3d closest_neighbor = Eigen::Vector3d::Zero();
+    double closest_distance = std::numeric_limits<double>::max();
+    std::for_each(neighbors.cbegin(), neighbors.cend(), [&](const auto &neighbor) {
+        double distance = (neighbor - point).norm();
+        if (distance < closest_distance) {
+            closest_neighbor = neighbor;
+            closest_distance = distance;
+        }
+    });
+    return std::make_tuple(closest_neighbor, closest_distance);
+}
+
 std::vector<Eigen::Vector3d> VoxelHashMap::Pointcloud() const {
     std::vector<Eigen::Vector3d> points;
     points.reserve(map_.size() * static_cast<size_t>(max_points_per_voxel_));

diff --git a/cpp/kiss_icp/core/VoxelHashMap.hpp b/cpp/kiss_icp/core/VoxelHashMap.hpp
@@ -48,7 +48,7 @@ struct VoxelHashMap {
     void AddPoints(const std::vector<Eigen::Vector3d> &points);
     void RemovePointsFarFromLocation(const Eigen::Vector3d &origin);
     std::vector<Eigen::Vector3d> Pointcloud() const;
-    std::vector<Eigen::Vector3d> GetPoints(const std::vector<Voxel> &query_voxels) const;
+    std::tuple<Eigen::Vector3d, double> GetClosestNeighbor(const Eigen::Vector3d &point) const;
 
     double voxel_size_;
     double max_distance_;