Nacho/strip nn search from voxel hash map

cpp/kiss_icp/core/Registration.cpp

@@ -31,44 +31,45 @@
 #include <sophus/so3.hpp>
 #include <tuple>
 
+#include "VoxelHashMap.hpp"
+
 namespace Eigen {
 using Matrix6d = Eigen::Matrix<double, 6, 6>;
 using Matrix3_6d = Eigen::Matrix<double, 3, 6>;
 using Vector6d = Eigen::Matrix<double, 6, 1>;
 }  // namespace Eigen
 
 namespace {
+// aliases
+using Vector3dVectorTuple = std::tuple<std::vector<Eigen::Vector3d>, std::vector<Eigen::Vector3d>>;
 
 inline double square(double x) { return x * x; }
 
-struct ResultTuple {
-    ResultTuple() {
-        JTJ.setZero();
-        JTr.setZero();
-    }
-
-    ResultTuple operator+(const ResultTuple &other) {
-        this->JTJ += other.JTJ;
-        this->JTr += other.JTr;
-        return *this;
-    }
-
-    Eigen::Matrix6d JTJ;
-    Eigen::Vector6d JTr;
-};
-
 void TransformPoints(const Sophus::SE3d &T, std::vector<Eigen::Vector3d> &points) {
     std::transform(points.cbegin(), points.cend(), points.begin(),
                    [&](const auto &point) { return T * point; });
 }
 
-constexpr int MAX_NUM_ITERATIONS_ = 500;
-constexpr double ESTIMATION_THRESHOLD_ = 0.0001;
-
 std::tuple<Eigen::Matrix6d, Eigen::Vector6d> BuildLinearSystem(
     const std::vector<Eigen::Vector3d> &source,
     const std::vector<Eigen::Vector3d> &target,
     double kernel) {
+    struct LinearSystemRes {
+        LinearSystemRes() {
+            JTJ.setZero();
+            JTr.setZero();
+        }
+
+        LinearSystemRes operator+(const LinearSystemRes &other) {
+            this->JTJ += other.JTJ;
+            this->JTr += other.JTr;
+            return *this;
+        }
+
+        Eigen::Matrix6d JTJ;
+        Eigen::Vector6d JTr;
+    };
+
     auto compute_jacobian_and_residual = [&](auto i) {
         const Eigen::Vector3d residual = source[i] - target[i];
         Eigen::Matrix3_6d J_r;
@@ -81,9 +82,9 @@ std::tuple<Eigen::Matrix6d, Eigen::Vector6d> BuildLinearSystem(
         // Range
         tbb::blocked_range<size_t>{0, source.size()},
         // Identity
-        ResultTuple(),
+        LinearSystemRes(),
         // 1st Lambda: Parallel computation
-        [&](const tbb::blocked_range<size_t> &r, ResultTuple J) -> ResultTuple {
+        [&](const tbb::blocked_range<size_t> &r, LinearSystemRes J) -> LinearSystemRes {
             auto Weight = [&](double residual2) {
                 return square(kernel) / square(kernel + residual2);
             };
@@ -97,19 +98,66 @@ std::tuple<Eigen::Matrix6d, Eigen::Vector6d> BuildLinearSystem(
             return J;
         },
         // 2nd Lambda: Parallel reduction of the private Jacboians
-        [&](ResultTuple a, const ResultTuple &b) -> ResultTuple { return a + b; });
+        [&](LinearSystemRes a, const LinearSystemRes &b) -> LinearSystemRes { return a + b; });
 
-    return std::make_tuple(JTJ, JTr);
+    return {JTJ, JTr};
+}
+
+Vector3dVectorTuple GetCorrespondences(const std::vector<Eigen::Vector3d> &points,
+                                       const kiss_icp::VoxelHashMap &voxel_map,
+                                       double max_correspondance_distance) {
+    struct CorrespondenceSet {
+        CorrespondenceSet(std::size_t n) {
+            source.reserve(n);
+            target.reserve(n);
+        }
+        std::vector<Eigen::Vector3d> source;
+        std::vector<Eigen::Vector3d> target;
+    };
+
+    using points_iterator = std::vector<Eigen::Vector3d>::const_iterator;
+    const auto [source, target] = tbb::parallel_reduce(
+        // Range
+        tbb::blocked_range<points_iterator>{points.cbegin(), points.cend()},
+        // Identity
+        CorrespondenceSet(points.size()),
+        // 1st lambda: Parallel computation
+        [&](const tbb::blocked_range<points_iterator> &r,
+            CorrespondenceSet res) -> CorrespondenceSet {
+            auto &[src, tgt] = res;
+            src.reserve(r.size());
+            tgt.reserve(r.size());
+            for (const auto &point : r) {
+                Eigen::Vector3d closest_neighbor = voxel_map.GetClosestNeighbor(point);
+                if ((closest_neighbor - point).norm() < max_correspondance_distance) {
+                    src.emplace_back(point);
+                    tgt.emplace_back(closest_neighbor);
+                }
+            }
+            return res;
+        },
+        // 2nd lambda: Parallel reduction
+        [](CorrespondenceSet a, const CorrespondenceSet &b) -> CorrespondenceSet {
+            auto &[src, tgt] = a;
+            const auto &[srcp, tgtp] = b;
+            src.insert(src.end(),  //
+                       std::make_move_iterator(srcp.begin()), std::make_move_iterator(srcp.end()));
+            tgt.insert(tgt.end(),  //
+                       std::make_move_iterator(tgtp.begin()), std::make_move_iterator(tgtp.end()));
+            return a;
+        });
+
+    return {source, target};
 }
 }  // namespace
 
 namespace kiss_icp {
 
-Sophus::SE3d RegisterFrame(const std::vector<Eigen::Vector3d> &frame,
-                           const VoxelHashMap &voxel_map,
-                           const Sophus::SE3d &initial_guess,
-                           double max_correspondence_distance,
-                           double kernel) {
+Sophus::SE3d Registration::AlignCloudToMap(const std::vector<Eigen::Vector3d> &frame,
+                                           const VoxelHashMap &voxel_map,
+                                           const Sophus::SE3d &initial_guess,
+                                           double max_distance,
+                                           double kernel) {
     if (voxel_map.Empty()) return initial_guess;
 
     // Equation (9)
@@ -118,9 +166,9 @@ Sophus::SE3d RegisterFrame(const std::vector<Eigen::Vector3d> &frame,
 
     // ICP-loop
     Sophus::SE3d T_icp = Sophus::SE3d();
-    for (int j = 0; j < MAX_NUM_ITERATIONS_; ++j) {
+    for (int j = 0; j < max_num_iterations_; ++j) {
         // Equation (10)
-        const auto &[src, tgt] = voxel_map.GetCorrespondences(source, max_correspondence_distance);
+        const auto &[src, tgt] = GetCorrespondences(source, voxel_map, max_distance);
         // Equation (11)
         const auto &[JTJ, JTr] = BuildLinearSystem(src, tgt, kernel);
         const Eigen::Vector6d dx = JTJ.ldlt().solve(-JTr);
@@ -130,7 +178,7 @@ Sophus::SE3d RegisterFrame(const std::vector<Eigen::Vector3d> &frame,
         // Update iterations
         T_icp = estimation * T_icp;
         // Termination criteria
-        if (dx.norm() < ESTIMATION_THRESHOLD_) break;
+        if (dx.norm() < estimation_threshold_) break;
     }
     // Spit the final transformation
     return T_icp * initial_guess;

cpp/kiss_icp/core/Registration.hpp

@@ -30,9 +30,19 @@
 
 namespace kiss_icp {
 
-Sophus::SE3d RegisterFrame(const std::vector<Eigen::Vector3d> &frame,
-                           const VoxelHashMap &voxel_map,
-                           const Sophus::SE3d &initial_guess,
-                           double max_correspondence_distance,
-                           double kernel);
+struct Registration {
+    explicit Registration(int max_num_iteration, double estimation_threshold)
+        : max_num_iterations_(max_num_iteration), estimation_threshold_(estimation_threshold) {}
+
+    // Register a point cloud to the given internal map representation. The config input parameter
+    // contains all the necessary parametrization for the ICP loop
+    Sophus::SE3d AlignCloudToMap(const std::vector<Eigen::Vector3d> &frame,
+                                 const VoxelHashMap &voxel_map,
+                                 const Sophus::SE3d &initial_guess,
+                                 double max_distance,
+                                 double kernel);
+
+    int max_num_iterations_;
+    double estimation_threshold_;
+};
 }  // namespace kiss_icp

cpp/kiss_icp/core/VoxelHashMap.cpp

@@ -22,107 +22,55 @@
 // SOFTWARE.
 #include "VoxelHashMap.hpp"
 
-#include <tbb/blocked_range.h>
-#include <tbb/parallel_reduce.h>
-
 #include <Eigen/Core>
 #include <algorithm>
 #include <limits>
 #include <tuple>
 #include <utility>
 #include <vector>
 
-// This parameters are not intended to be changed, therefore we do not expose it
-namespace {
-struct ResultTuple {
-    ResultTuple(std::size_t n) {
-        source.reserve(n);
-        target.reserve(n);
-    }
-    std::vector<Eigen::Vector3d> source;
-    std::vector<Eigen::Vector3d> target;
-};
-}  // namespace
-
 namespace kiss_icp {
 
-VoxelHashMap::Vector3dVectorTuple VoxelHashMap::GetCorrespondences(
-    const Vector3dVector &points, double max_correspondance_distance) const {
-    // Lambda Function to obtain the KNN of one point, maybe refactor
-    auto GetClosestNeighboor = [&](const Eigen::Vector3d &point) {
-        auto kx = static_cast<int>(point[0] / voxel_size_);
-        auto ky = static_cast<int>(point[1] / voxel_size_);
-        auto kz = static_cast<int>(point[2] / voxel_size_);
-        std::vector<Voxel> voxels;
-        voxels.reserve(27);
-        for (int i = kx - 1; i < kx + 1 + 1; ++i) {
-            for (int j = ky - 1; j < ky + 1 + 1; ++j) {
-                for (int k = kz - 1; k < kz + 1 + 1; ++k) {
-                    voxels.emplace_back(i, j, k);
-                }
+// Function to obtain the KNN of one point
+Eigen::Vector3d VoxelHashMap::GetClosestNeighbor(const Eigen::Vector3d &point) const {
+    auto kx = static_cast<int>(point[0] / voxel_size_);
+    auto ky = static_cast<int>(point[1] / voxel_size_);
+    auto kz = static_cast<int>(point[2] / voxel_size_);
+    std::vector<kiss_icp::VoxelHashMap::Voxel> voxels;
+    voxels.reserve(27);
+    for (int i = kx - 1; i < kx + 1 + 1; ++i) {
+        for (int j = ky - 1; j < ky + 1 + 1; ++j) {
+            for (int k = kz - 1; k < kz + 1 + 1; ++k) {
+                voxels.emplace_back(i, j, k);
             }
         }
+    }
 
-        using Vector3dVector = std::vector<Eigen::Vector3d>;
-        Vector3dVector neighboors;
-        neighboors.reserve(27 * max_points_per_voxel_);
-        std::for_each(voxels.cbegin(), voxels.cend(), [&](const auto &voxel) {
-            auto search = map_.find(voxel);
-            if (search != map_.end()) {
-                const auto &points = search->second.points;
-                if (!points.empty()) {
-                    for (const auto &point : points) {
-                        neighboors.emplace_back(point);
-                    }
+    std::vector<Eigen::Vector3d> neighbors;
+    neighbors.reserve(static_cast<size_t>(27 * max_points_per_voxel_));
+    std::for_each(voxels.cbegin(), voxels.cend(), [&](const auto &voxel) {
+        auto search = map_.find(voxel);
+        if (search != map_.end()) {
+            const auto &points = search->second.points;
+            if (!points.empty()) {
+                for (const auto &point : points) {
+                    neighbors.emplace_back(point);
                 }
             }
-        });
-
-        Eigen::Vector3d closest_neighbor;
-        double closest_distance2 = std::numeric_limits<double>::max();
-        std::for_each(neighboors.cbegin(), neighboors.cend(), [&](const auto &neighbor) {
-            double distance = (neighbor - point).squaredNorm();
-            if (distance < closest_distance2) {
-                closest_neighbor = neighbor;
-                closest_distance2 = distance;
-            }
-        });
+        }
+    });
 
-        return closest_neighbor;
-    };
-    using points_iterator = std::vector<Eigen::Vector3d>::const_iterator;
-    const auto [source, target] = tbb::parallel_reduce(
-        // Range
-        tbb::blocked_range<points_iterator>{points.cbegin(), points.cend()},
-        // Identity
-        ResultTuple(points.size()),
-        // 1st lambda: Parallel computation
-        [max_correspondance_distance, &GetClosestNeighboor](
-            const tbb::blocked_range<points_iterator> &r, ResultTuple res) -> ResultTuple {
-            auto &[src, tgt] = res;
-            src.reserve(r.size());
-            tgt.reserve(r.size());
-            for (const auto &point : r) {
-                Eigen::Vector3d closest_neighboors = GetClosestNeighboor(point);
-                if ((closest_neighboors - point).norm() < max_correspondance_distance) {
-                    src.emplace_back(point);
-                    tgt.emplace_back(closest_neighboors);
-                }
-            }
-            return res;
-        },
-        // 2nd lambda: Parallel reduction
-        [](ResultTuple a, const ResultTuple &b) -> ResultTuple {
-            auto &[src, tgt] = a;
-            const auto &[srcp, tgtp] = b;
-            src.insert(src.end(),  //
-                       std::make_move_iterator(srcp.begin()), std::make_move_iterator(srcp.end()));
-            tgt.insert(tgt.end(),  //
-                       std::make_move_iterator(tgtp.begin()), std::make_move_iterator(tgtp.end()));
-            return a;
-        });
+    Eigen::Vector3d closest_neighbor;
+    double closest_distance2 = std::numeric_limits<double>::max();
+    std::for_each(neighbors.cbegin(), neighbors.cend(), [&](const auto &neighbor) {
+        double distance = (neighbor - point).squaredNorm();
+        if (distance < closest_distance2) {
+            closest_neighbor = neighbor;
+            closest_distance2 = distance;
+        }
+    });
 
-    return std::make_tuple(source, target);
+    return closest_neighbor;
 }
 
 std::vector<Eigen::Vector3d> VoxelHashMap::Pointcloud() const {
@@ -137,13 +85,14 @@ std::vector<Eigen::Vector3d> VoxelHashMap::Pointcloud() const {
     return points;
 }
 
-void VoxelHashMap::Update(const Vector3dVector &points, const Eigen::Vector3d &origin) {
+void VoxelHashMap::Update(const std::vector<Eigen::Vector3d> &points,
+                          const Eigen::Vector3d &origin) {
     AddPoints(points);
     RemovePointsFarFromLocation(origin);
 }
 
-void VoxelHashMap::Update(const Vector3dVector &points, const Sophus::SE3d &pose) {
-    Vector3dVector points_transformed(points.size());
+void VoxelHashMap::Update(const std::vector<Eigen::Vector3d> &points, const Sophus::SE3d &pose) {
+    std::vector<Eigen::Vector3d> points_transformed(points.size());
     std::transform(points.cbegin(), points.cend(), points_transformed.begin(),
                    [&](const auto &point) { return pose * point; });
     const Eigen::Vector3d &origin = pose.translation();

cpp/kiss_icp/core/VoxelHashMap.hpp

@@ -34,8 +34,6 @@
 
 namespace kiss_icp {
 struct VoxelHashMap {
-    using Vector3dVector = std::vector<Eigen::Vector3d>;
-    using Vector3dVectorTuple = std::tuple<Vector3dVector, Vector3dVector>;
     using Voxel = Eigen::Vector3i;
     struct VoxelBlock {
         // buffer of points with a max limit of n_points
@@ -57,15 +55,14 @@ struct VoxelHashMap {
           max_distance_(max_distance),
           max_points_per_voxel_(max_points_per_voxel) {}
 
-    Vector3dVectorTuple GetCorrespondences(const Vector3dVector &points,
-                                           double max_correspondance_distance) const;
     inline void Clear() { map_.clear(); }
     inline bool Empty() const { return map_.empty(); }
     void Update(const std::vector<Eigen::Vector3d> &points, const Eigen::Vector3d &origin);
     void Update(const std::vector<Eigen::Vector3d> &points, const Sophus::SE3d &pose);
     void AddPoints(const std::vector<Eigen::Vector3d> &points);
     void RemovePointsFarFromLocation(const Eigen::Vector3d &origin);
     std::vector<Eigen::Vector3d> Pointcloud() const;
+    Eigen::Vector3d GetClosestNeighbor(const Eigen::Vector3d &point) const;
 
     double voxel_size_;
     double max_distance_;

cpp/kiss_icp/pipeline/KissICP.cpp

@@ -68,11 +68,11 @@ KissICP::Vector3dVectorTuple KissICP::RegisterFrame(const std::vector<Eigen::Vec
     const auto initial_guess = last_pose * prediction;
 
     // Run icp
-    const Sophus::SE3d new_pose = kiss_icp::RegisterFrame(source,         //
-                                                          local_map_,     //
-                                                          initial_guess,  //
-                                                          3.0 * sigma,    //
-                                                          sigma / 3.0);
+    const Sophus::SE3d new_pose = registration_.AlignCloudToMap(source,         //
+                                                                local_map_,     //
+                                                                initial_guess,  //
+                                                                3.0 * sigma,    //
+                                                                sigma / 3.0);
     const auto model_deviation = initial_guess.inverse() * new_pose;
     adaptive_threshold_.UpdateModelDeviation(model_deviation);
     local_map_.Update(frame_downsample, new_pose);