Address NDT review feedback

Author: Xiang-Zeng

CHANGELOG.md

@@ -51,7 +51,7 @@
 -   Fix macOS arm64 builds, add CI runner for macOS arm64 (PR #6695)
 -   Fix KDTreeFlann possibly using a dangling pointer instead of internal storage and simplified its members (PR #6734)
 -   Fix RANSAC early stop if no inliers in a specific iteration (PR #6789)
--   Add 3D Normal Distributions Transform registration with C++ and Python APIs.
+-   Add 3D Normal Distributions Transform registration with C++ and Python APIs (PR #7517).
 -   Fix segmentation fault (infinite recursion) of DetectPlanarPatches if multiple points have same coordinates (PR #6794)
 -   `TriangleMesh`'s `+=` operator appends UVs regardless of the presence of existing features (PR #6728)
 -   Fix build with fmt v10.2.0 (#6783)

cpp/open3d/pipelines/registration/NormalDistributionsTransform.cpp

@@ -62,6 +62,23 @@ struct VoxelGaussian {
 
 using VoxelMap = std::unordered_map<VoxelKey, VoxelGaussian, VoxelKeyHash>;
 
+struct VoxelAccumulator {
+    int count = 0;
+    Eigen::Vector3d mean = Eigen::Vector3d::Zero();
+    Eigen::Matrix3d covariance_accumulator = Eigen::Matrix3d::Zero();
+
+    void AddPoint(const Eigen::Vector3d &point) {
+        ++count;
+        const Eigen::Vector3d delta = point - mean;
+        mean += delta / static_cast<double>(count);
+        const Eigen::Vector3d delta_after_update = point - mean;
+        covariance_accumulator += delta * delta_after_update.transpose();
+    }
+};
+
+using VoxelAccumulatorMap =
+        std::unordered_map<VoxelKey, VoxelAccumulator, VoxelKeyHash>;
+
 struct NDTLinearSystem {
     Eigen::Matrix6d JTJ = Eigen::Matrix6d::Zero();
     Eigen::Vector6d JTr = Eigen::Vector6d::Zero();
@@ -152,39 +169,24 @@ void ValidateNDTOption(const NormalDistributionsTransformOption &option) {
 VoxelMap BuildVoxelGaussians(const geometry::PointCloud &target,
                              const NormalDistributionsTransformOption &option) {
     const double inv_voxel_size = 1.0 / option.voxel_size_;
-    std::unordered_map<VoxelKey, std::vector<int>, VoxelKeyHash> voxel_indices;
-    for (int i = 0; i < static_cast<int>(target.points_.size()); ++i) {
-        voxel_indices[GetVoxelKey(target.points_[i], inv_voxel_size)].push_back(
-                i);
+    VoxelAccumulatorMap voxel_accumulators;
+    for (const Eigen::Vector3d &point : target.points_) {
+        voxel_accumulators[GetVoxelKey(point, inv_voxel_size)].AddPoint(point);
     }
 
     VoxelMap voxel_map;
-    for (const auto &item : voxel_indices) {
-        const auto &indices = item.second;
-        if (static_cast<int>(indices.size()) < option.min_points_per_voxel_) {
+    for (const auto &item : voxel_accumulators) {
+        const VoxelAccumulator &accumulator = item.second;
+        if (accumulator.count < option.min_points_per_voxel_) {
             continue;
         }
 
         VoxelGaussian gaussian;
-        gaussian.count = static_cast<int>(indices.size());
-        for (const int idx : indices) {
-            gaussian.mean += target.points_[idx];
-        }
-        gaussian.mean /= static_cast<double>(indices.size());
-
-        Eigen::Matrix3d covariance = Eigen::Matrix3d::Zero();
-        double representative_distance2 = std::numeric_limits<double>::max();
-        for (const int idx : indices) {
-            const Eigen::Vector3d centered =
-                    target.points_[idx] - gaussian.mean;
-            covariance += centered * centered.transpose();
-            const double distance2 = centered.squaredNorm();
-            if (distance2 < representative_distance2) {
-                representative_distance2 = distance2;
-                gaussian.representative_index = idx;
-            }
-        }
-        covariance /= static_cast<double>(indices.size() - 1);
+        gaussian.count = accumulator.count;
+        gaussian.mean = accumulator.mean;
+        const Eigen::Matrix3d covariance =
+                accumulator.covariance_accumulator /
+                static_cast<double>(accumulator.count - 1);
 
         Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> solver(covariance);
         if (solver.info() != Eigen::Success) {
@@ -205,6 +207,22 @@ VoxelMap BuildVoxelGaussians(const geometry::PointCloud &target,
                                solver.eigenvectors().transpose();
         voxel_map.emplace(item.first, gaussian);
     }
+
+    for (int i = 0; i < static_cast<int>(target.points_.size()); ++i) {
+        const Eigen::Vector3d &point = target.points_[i];
+        auto voxel_itr = voxel_map.find(GetVoxelKey(point, inv_voxel_size));
+        if (voxel_itr == voxel_map.end()) {
+            continue;
+        }
+        VoxelGaussian &gaussian = voxel_itr->second;
+        const double distance2 = (point - gaussian.mean).squaredNorm();
+        if (gaussian.representative_index < 0 ||
+            distance2 < (target.points_[gaussian.representative_index] -
+                         gaussian.mean)
+                                .squaredNorm()) {
+            gaussian.representative_index = i;
+        }
+    }
     return voxel_map;
 }
 

cpp/open3d/pipelines/registration/NormalDistributionsTransform.h

@@ -79,7 +79,9 @@ class NormalDistributionsTransformOption {
 ///
 /// This implementation builds a voxelized Gaussian model from the target point
 /// cloud and optimizes the rigid source-to-target transformation with
-/// Gauss-Newton iterations.
+/// Gauss-Newton iterations. The returned correspondence set uses the target
+/// point closest to each accepted voxel mean as its representative, and the
+/// inlier RMSE is computed from those representative point correspondences.
 ///
 /// \param source The source point cloud.
 /// \param target The target point cloud.

cpp/pybind/pipelines/registration/registration.cpp

@@ -418,18 +418,8 @@ Sets :math:`c = 1` if ``with_scaling`` is ``False``.
     py::detail::bind_copy_functions<NormalDistributionsTransformOption>(
             ndt_option);
     ndt_option
-            .def(py::init([](double voxel_size, int min_points_per_voxel,
-                             double covariance_regularization,
-                             double transformation_epsilon,
-                             double relative_objective, int max_iteration,
-                             double outlier_threshold,
-                             int neighbor_search_type) {
-                     return new NormalDistributionsTransformOption(
-                             voxel_size, min_points_per_voxel,
-                             covariance_regularization, transformation_epsilon,
-                             relative_objective, max_iteration,
-                             outlier_threshold, neighbor_search_type);
-                 }),
+            .def(py::init<double, int, double, double, double, int, double,
+                          int>(),
                  "voxel_size"_a = 1.0, "min_points_per_voxel"_a = 6,
                  "covariance_regularization"_a = 1e-3,
                  "transformation_epsilon"_a = 1e-6,
@@ -773,8 +763,6 @@ must hold true for all edges.)");
                      "the "
                      "source point's correspondence is itself."},
                     {"option", "Registration option"},
-                    {"ndt_option",
-                     "Normal Distributions Transform registration option."},
                     {"ransac_n",
                      "Fit ransac with ``ransac_n`` correspondences"},
                     {"source_feature", "Source point cloud feature."},
@@ -834,8 +822,11 @@ must hold true for all edges.)");
                        "source"_a, "target"_a,
                        "option"_a = NormalDistributionsTransformOption(),
                        "init"_a = Eigen::Matrix4d::Identity());
+    auto ndt_argument_docstrings = map_shared_argument_docstrings;
+    ndt_argument_docstrings["option"] =
+            "Normal Distributions Transform registration option.";
     docstring::FunctionDocInject(m_registration, "registration_ndt",
-                                 map_shared_argument_docstrings);
+                                 ndt_argument_docstrings);
 
     m_registration.def(
             "registration_ransac_based_on_correspondence",

docs/tutorial/pipelines/ndt_registration.rst

@@ -15,7 +15,10 @@ Open3D exposes NDT through
 collected in ``NormalDistributionsTransformOption``, including both voxel
 Gaussian model parameters and convergence criteria. Optimization stops when
 the pose update is small or the relative change in mean Mahalanobis objective
-falls below the configured threshold:
+falls below the configured threshold. In the returned ``RegistrationResult``,
+each accepted voxel is represented by the target point closest to its mean,
+and ``inlier_rmse`` is the Euclidean RMSE over those representative point
+correspondences:
 
 .. literalinclude:: ../../../examples/python/pipelines/ndt_registration.py
    :language: python