Feaure matching API (#5964)

* expose correspondences from features api for external feature usages and debugging

* Improve FPFH feature matching time

cpp/open3d/pipelines/registration/Feature.cpp

@@ -139,6 +139,68 @@ std::shared_ptr<Feature> ComputeFPFHFeature(
  return feature;
 }
 
+CorrespondenceSet CorrespondencesFromFeatures(const Feature &source_features,
+ const Feature &target_features,
+ bool mutual_filter,
+ float mutual_consistent_ratio) {
+ const int num_searches = mutual_filter ? 2 : 1;
+
+ // Access by reference, since Eigen Matrix could be copied
+ std::array<std::reference_wrapper<const Feature>, 2> features{
+ std::reference_wrapper<const Feature>(source_features),
+ std::reference_wrapper<const Feature>(target_features)};
+ std::array<int, 2> num_pts{int(source_features.data_.cols()),
+ int(target_features.data_.cols())};
+ std::vector<CorrespondenceSet> corres(num_searches);
+
+ const int kMaxThreads = utility::EstimateMaxThreads();
+
+ const int kOuterThreads = std::min(kMaxThreads, num_searches);
+ const int kInnerThreads = std::max(kMaxThreads / num_searches, 1);
+#pragma omp parallel for num_threads(kOuterThreads)
+ for (int k = 0; k < num_searches; ++k) {
+ geometry::KDTreeFlann kdtree(features[1 - k]);
+
+ int num_pts_k = num_pts[k];
+ corres[k] = CorrespondenceSet(num_pts_k);
+#pragma omp parallel for num_threads(kInnerThreads)
+ for (int i = 0; i < num_pts_k; i++) {
+ std::vector<int> corres_tmp(1);
+ std::vector<double> dist_tmp(1);
+
+ kdtree.SearchKNN(Eigen::VectorXd(features[k].get().data_.col(i)), 1,
+ corres_tmp, dist_tmp);
+ int j = corres_tmp[0];
+ corres[k][i] = Eigen::Vector2i(i, j);
+ }
+ }
+
+ // corres[0]: corres_ij, corres[1]: corres_ji
+ if (!mutual_filter) return corres[0];
+
+ // should not use parallel for due to emplace back
+ CorrespondenceSet corres_mutual;
+ int num_src_pts = num_pts[0];
+ for (int i = 0; i < num_src_pts; ++i) {
+ int j = corres[0][i](1);
+ if (corres[1][j](1) == i) {
+ corres_mutual.emplace_back(i, j);
+ }
+ }
+
+ // Empirically mutual correspondence set should not be too small
+ if (int(corres_mutual.size()) >=
+ int(mutual_consistent_ratio * num_src_pts)) {
+ utility::LogDebug("{:d} correspondences remain after mutual filter",
+ corres_mutual.size());
+ return corres_mutual;
+ }
+ utility::LogWarning(
+ "Too few correspondences ({:d}) after mutual filter, fall back to "
+ "original correspondences.",
+ corres_mutual.size());
+ return corres[0];
+}
 } // namespace registration
 } // namespace pipelines
 } // namespace open3d

cpp/open3d/pipelines/registration/Feature.h

@@ -22,6 +22,8 @@ class PointCloud;
 namespace pipelines {
 namespace registration {
 
+typedef std::vector<Eigen::Vector2i> CorrespondenceSet;
+
 /// \class Feature
 ///
 /// \brief Class to store featrues for registration.
@@ -54,6 +56,28 @@ std::shared_ptr<Feature> ComputeFPFHFeature(
  const geometry::KDTreeSearchParam &search_param =
  geometry::KDTreeSearchParamKNN());
 
+/// \brief Function to find correspondences via 1-nearest neighbor feature
+/// matching. Target is used to construct a nearest neighbor search
+/// object, in order to query source.
+/// \param source_features (D, N) feature
+/// \param target_features (D, M) feature
+/// \param mutual_filter Boolean flag, only return correspondences (i, j) s.t.
+/// source_features[i] and target_features[j] are mutually the nearest neighbor.
+/// \param mutual_consistency_ratio Float threshold to decide whether the number
+/// of correspondences is sufficient. Only used when mutual_filter is set to
+/// True.
+/// \return A CorrespondenceSet. When mutual_filter is disabled: the first
+/// column is arange(0, N) of source, and the second column is the corresponding
+/// index of target. When mutual_filter is enabled, return the filtering subset
+/// of the aforementioned correspondence set where source[i] and target[j] are
+/// mutually the nearest neighbor. If the subset size is smaller than
+/// mutual_consistency_ratio * N, return the unfiltered set.
+CorrespondenceSet CorrespondencesFromFeatures(
+ const Feature &source_features,
+ const Feature &target_features,
+ bool mutual_filter = false,
+ float mutual_consistency_ratio = 0.1);
+
 } // namespace registration
 } // namespace pipelines
 } // namespace open3d

cpp/open3d/pipelines/registration/Registration.cpp

@@ -280,8 +280,8 @@ RegistrationResult RegistrationRANSACBasedOnCorrespondence(
 RegistrationResult RegistrationRANSACBasedOnFeatureMatching(
  const geometry::PointCloud &source,
  const geometry::PointCloud &target,
- const Feature &source_feature,
- const Feature &target_feature,
+ const Feature &source_features,
+ const Feature &target_features,
  bool mutual_filter,
  double max_correspondence_distance,
  const TransformationEstimation
@@ -295,62 +295,11 @@ RegistrationResult RegistrationRANSACBasedOnFeatureMatching(
  return RegistrationResult();
  }
 
- int num_src_pts = int(source.points_.size());
- int num_tgt_pts = int(target.points_.size());
-
- geometry::KDTreeFlann kdtree_target(target_feature);
- pipelines::registration::CorrespondenceSet corres_ij(num_src_pts);
-
-#pragma omp parallel for num_threads(utility::EstimateMaxThreads())
- for (int i = 0; i < num_src_pts; i++) {
- std::vector<int> corres_tmp(1);
- std::vector<double> dist_tmp(1);
-
- kdtree_target.SearchKNN(Eigen::VectorXd(source_feature.data_.col(i)), 1,
- corres_tmp, dist_tmp);
- int j = corres_tmp[0];
- corres_ij[i] = Eigen::Vector2i(i, j);
- }
-
- // Do reverse check if mutual_filter is enabled
- if (mutual_filter) {
- geometry::KDTreeFlann kdtree_source(source_feature);
- pipelines::registration::CorrespondenceSet corres_ji(num_tgt_pts);
-
-#pragma omp parallel for num_threads(utility::EstimateMaxThreads())
- for (int j = 0; j < num_tgt_pts; ++j) {
- std::vector<int> corres_tmp(1);
- std::vector<double> dist_tmp(1);
- kdtree_source.SearchKNN(
- Eigen::VectorXd(target_feature.data_.col(j)), 1, corres_tmp,
- dist_tmp);
- int i = corres_tmp[0];
- corres_ji[j] = Eigen::Vector2i(i, j);
- }
-
- pipelines::registration::CorrespondenceSet corres_mutual;
- for (int i = 0; i < num_src_pts; ++i) {
- int j = corres_ij[i](1);
- if (corres_ji[j](0) == i) {
- corres_mutual.emplace_back(i, j);
- }
- }
-
- // Empirically mutual correspondence set should not be too small
- if (int(corres_mutual.size()) >= ransac_n * 3) {
- utility::LogDebug("{:d} correspondences remain after mutual filter",
- corres_mutual.size());
- return RegistrationRANSACBasedOnCorrespondence(
- source, target, corres_mutual, max_correspondence_distance,
- estimation, ransac_n, checkers, criteria);
- }
- utility::LogDebug(
- "Too few correspondences after mutual filter, fall back to "
- "original correspondences.");
- }
+ CorrespondenceSet corres = CorrespondencesFromFeatures(
+ source_features, target_features, mutual_filter);
 
  return RegistrationRANSACBasedOnCorrespondence(
- source, target, corres_ij, max_correspondence_distance, estimation,
+ source, target, corres, max_correspondence_distance, estimation,
  ransac_n, checkers, criteria);
 }
 

cpp/open3d/pipelines/registration/Registration.h

@@ -185,8 +185,8 @@ RegistrationResult RegistrationRANSACBasedOnCorrespondence(
 ///
 /// \param source The source point cloud.
 /// \param target The target point cloud.
-/// \param source_feature Source point cloud feature.
-/// \param target_feature Target point cloud feature.
+/// \param source_features Source point cloud feature.
+/// \param target_features Target point cloud feature.
 /// \param mutual_filter Enables mutual filter such that the correspondence of
 /// the source point's correspondence is itself.
 /// \param max_correspondence_distance Maximum correspondence points-pair
@@ -197,8 +197,8 @@ RegistrationResult RegistrationRANSACBasedOnCorrespondence(
 RegistrationResult RegistrationRANSACBasedOnFeatureMatching(
  const geometry::PointCloud &source,
  const geometry::PointCloud &target,
- const Feature &source_feature,
- const Feature &target_feature,
+ const Feature &source_features,
+ const Feature &target_features,
  bool mutual_filter,
  double max_correspondence_distance,
  const TransformationEstimation &estimation =

cpp/open3d/t/pipelines/registration/Feature.cpp

@@ -10,6 +10,7 @@
 #include "open3d/core/nns/NearestNeighborSearch.h"
 #include "open3d/t/geometry/PointCloud.h"
 #include "open3d/t/pipelines/kernel/Feature.h"
+#include "open3d/utility/Parallel.h"
 
 namespace open3d {
 namespace t {
@@ -90,6 +91,63 @@ core::Tensor ComputeFPFHFeature(const geometry::PointCloud &input,
  return fpfh;
 }
 
+core::Tensor CorrespondencesFromFeatures(const core::Tensor &source_features,
+ const core::Tensor &target_features,
+ bool mutual_filter,
+ float mutual_consistent_ratio) {
+ const int num_searches = mutual_filter ? 2 : 1;
+
+ std::array<core::Tensor, 2> features{source_features, target_features};
+ std::vector<core::Tensor> corres(num_searches);
+
+ const int kMaxThreads = utility::EstimateMaxThreads();
+ const int kOuterThreads = std::min(kMaxThreads, num_searches);
+
+ // corres[0]: corres_ij, corres[1]: corres_ji
+#pragma omp parallel for num_threads(kOuterThreads)
+ for (int i = 0; i < num_searches; ++i) {
+ core::nns::NearestNeighborSearch nns(features[1 - i],
+ core::Dtype::Int64);
+ nns.KnnIndex();
+ auto result = nns.KnnSearch(features[i], 1);
+
+ corres[i] = result.first.View({-1});
+ }
+
+ auto corres_ij = corres[0];
+ core::Tensor arange_source =
+ core::Tensor::Arange(0, source_features.GetLength(), 1,
+ corres_ij.GetDtype(), corres_ij.GetDevice());
+
+ // Change view for the appending axis
+ core::Tensor result_ij =
+ arange_source.View({-1, 1}).Append(corres_ij.View({-1, 1}), 1);
+
+ if (!mutual_filter) {
+ return result_ij;
+ }
+
+ auto corres_ji = corres[1];
+ // Mutually consistent
+ core::Tensor corres_ii = corres_ji.IndexGet({corres_ij});
+ core::Tensor identical = corres_ii.Eq(arange_source);
+ core::Tensor result_mutual = corres_ij.IndexGet({identical});
+ if (result_mutual.GetLength() >
+ mutual_consistent_ratio * arange_source.GetLength()) {
+ utility::LogDebug("{:d} correspondences remain after mutual filter",
+ result_mutual.GetLength());
+ return arange_source.IndexGet({identical})
+ .View({-1, 1})
+ .Append(result_mutual.View({-1, 1}), 1);
+ }
+ // fall back to full correspondences
+ utility::LogWarning(
+ "Too few correspondences ({:d}) after mutual filter, fall back to "
+ "original correspondences.",
+ result_mutual.GetLength());
+ return result_ij;
+}
+
 } // namespace registration
 } // namespace pipelines
 } // namespace t

cpp/open3d/t/pipelines/registration/Feature.h

@@ -37,6 +37,26 @@ core::Tensor ComputeFPFHFeature(
  const utility::optional<int> max_nn = 100,
  const utility::optional<double> radius = utility::nullopt);
 
+/// \brief Function to find correspondences via 1-nearest neighbor feature
+/// matching. Target is used to construct a nearest neighbor search
+/// object, in order to query source.
+/// \param source_feats (N, D) tensor
+/// \param target_feats (M, D) tensor
+/// \param mutual_filter Boolean flag, only return correspondences (i, j) s.t.
+/// source_features[i] and target_features[j] are mutually the nearest neighbor.
+/// \param mutual_consistency_ratio Float threshold to decide whether the number
+/// of correspondences is sufficient. Only used when mutual_filter is set to
+/// True.
+/// \return (K, 2, Int64) tensor. When mutual_filter is disabled: the first
+/// column is arange(0, N) of source, and the second column is the corresponding
+/// index of target. When mutual_filter is enabled, return the filtering subset
+/// of the aforementioned correspondence set where source[i] and target[j] are
+/// mutually the nearest neighbor. If the subset size is smaller than
+/// mutual_consistency_ratio * N, return the unfiltered set.
+core::Tensor CorrespondencesFromFeatures(const core::Tensor &source_features,
+ const core::Tensor &target_features,
+ bool mutual_filter = false,
+ float mutual_consistency_ratio = 0.1);
 } // namespace registration
 } // namespace pipelines
 } // namespace t

cpp/pybind/pipelines/registration/feature.cpp

@@ -53,6 +53,23 @@ void pybind_feature_methods(py::module &m) {
  m, "compute_fpfh_feature",
  {{"input", "The Input point cloud."},
  {"search_param", "KDTree KNN search parameter."}});
+
+ m.def("correspondences_from_features", &CorrespondencesFromFeatures,
+ "Function to find nearest neighbor correspondences from features",
+ "source_features"_a, "target_features"_a, "mutual_filter"_a = false,
+ "mutual_consistency_ratio"_a = 0.1f);
+ docstring::FunctionDocInject(
+ m, "correspondences_from_features",
+ {{"source_features", "The source features stored in (dim, N)."},
+ {"target_features", "The target features stored in (dim, M)."},
+ {"mutual_filter",
+ "filter correspondences and return the collection of (i, j) s.t. "
+ "source_features[i] and target_features[j] are mutually the "
+ "nearest neighbor."},
+ {"mutual_consistency_ratio",
+ "Threshold to decide whether the number of filtered "
+ "correspondences is sufficient. Only used when mutual_filter is "
+ "enabled."}});
 }
 
 } // namespace registration

cpp/pybind/t/pipelines/registration/feature.cpp

@@ -35,6 +35,26 @@ parameter is provided, and Hybrid search (Recommended) if both are provided.)",
  {"radius",
  "[optional] Neighbor search radius parameter. [Recommended ~5x "
  "voxel size]"}});
+
+ m.def("correspondences_from_features", &CorrespondencesFromFeatures,
+ py::call_guard<py::gil_scoped_release>(),
+ R"(Function to query nearest neighbors of source_features in target_features.)",
+ "source_features"_a, "target_features"_a, "mutual_filter"_a = false,
+ "mutual_consistency_ratio"_a = 0.1f);
+ docstring::FunctionDocInject(
+ m, "correspondences_from_features",
+ {{"source_features", "The source features in shape (N, dim)."},
+ {"target_features", "The target features in shape (M, dim)."},
+ {"mutual_filter",
+ "filter correspondences and return the collection of (i, j) "
+ "s.t. "
+ "source_features[i] and target_features[j] are mutually the "
+ "nearest neighbor."},
+ {"mutual_consistency_ratio",
+ "Threshold to decide whether the number of filtered "
+ "correspondences is sufficient. Only used when "
+ "mutual_filter is "
+ "enabled."}});
 }
 
 } // namespace registration