Add optional indices arg for fast computation of a small subset of FPFH features

CHANGELOG.md

@@ -55,6 +55,7 @@
 -   Fix render to depth image on Apple Retina displays (PR #7001)
 -   Fix infinite loop in segment_plane if num_points < ransac_n (PR #7032)
 -   Add select_by_index method to Feature class (PR #7039)
+-   Add optional indices arg for fast computation of a small subset of FPFH features (PR #7118).
 
 
 ## 0.13

cpp/open3d/pipelines/registration/Feature.cpp

@@ -21,23 +21,38 @@ namespace registration {
 std::shared_ptr<Feature> Feature::SelectByIndex(
         const std::vector<size_t> &indices, bool invert /* = false */) const {
     auto output = std::make_shared<Feature>();
-    output->Resize(data_.rows(), indices.size());
 
     std::vector<bool> mask = std::vector<bool>(data_.cols(), invert);
+    size_t n_output_features = 0;
     for (size_t i : indices) {
-        mask[i] = !invert;
+        if (i < mask.size()) {
+            if (mask[i] == invert) {
+                mask[i] = !invert;
+                n_output_features++;
+            }
+        } else {
+            utility::LogWarning(
+                    "[SelectByIndex] contains index {} that is "
+                    "not within the bounds",
+                    (int)i);
+        }
+    }
+    if (invert) {
+        n_output_features = data_.cols() - n_output_features;
     }
 
-    size_t current_col_feature = 0;
-    for (size_t i = 0; i < static_cast<size_t>(data_.cols()); i++) {
+    output->Resize(data_.rows(), n_output_features);
+
+    for (size_t i = 0, current_col_feature = 0;
+         i < static_cast<size_t>(data_.cols()); i++) {
         if (mask[i]) {
-            output->data_.col(current_col_feature) = data_.col(i);
-            current_col_feature++;
+            output->data_.col(current_col_feature++) = data_.col(i);
         }
     }
 
     utility::LogDebug(
-            "Feature group down sampled from {:d} features to {:d} features.",
+            "[SelectByIndex] Feature group down sampled from {:d} features to "
+            "{:d} features.",
             (int)data_.cols(), (int)output->data_.cols());
 
     return output;
@@ -80,14 +95,23 @@ static Eigen::Vector4d ComputePairFeatures(const Eigen::Vector3d &p1,
 static std::shared_ptr<Feature> ComputeSPFHFeature(
         const geometry::PointCloud &input,
         const geometry::KDTreeFlann &kdtree,
-        const geometry::KDTreeSearchParam &search_param) {
+        const geometry::KDTreeSearchParam &search_param,
+        const utility::optional<std::vector<size_t>> &indices =
+                utility::nullopt) {
+    const bool filter_spfh = indices.has_value();
+    const auto spfh_indices = indices.value_or(std::vector<size_t>());
+
+    const size_t n_spfh =
+            filter_spfh ? spfh_indices.size() : input.points_.size();
     auto feature = std::make_shared<Feature>();
-    feature->Resize(33, (int)input.points_.size());
+    feature->Resize(33, (int)n_spfh);
+
 #pragma omp parallel for schedule(static) \
         num_threads(utility::EstimateMaxThreads())
-    for (int i = 0; i < (int)input.points_.size(); i++) {
-        const auto &point = input.points_[i];
-        const auto &normal = input.normals_[i];
+    for (int i = 0; i < (int)n_spfh; i++) {
+        const int point_idx = filter_spfh ? spfh_indices[i] : i;
+        const auto &point = input.points_[point_idx];
+        const auto &normal = input.normals_[point_idx];
         std::vector<int> indices;
         std::vector<double> distance2;
         if (kdtree.Search(point, search_param, indices, distance2) > 1) {
@@ -119,31 +143,129 @@ static std::shared_ptr<Feature> ComputeSPFHFeature(
 std::shared_ptr<Feature> ComputeFPFHFeature(
         const geometry::PointCloud &input,
         const geometry::KDTreeSearchParam
-                &search_param /* = geometry::KDTreeSearchParamKNN()*/) {
-    auto feature = std::make_shared<Feature>();
-    feature->Resize(33, (int)input.points_.size());
+                &search_param /* = geometry::KDTreeSearchParamKNN()*/,
+        const utility::optional<std::vector<size_t>>
+                &indices /* = utility::nullopt*/) {
     if (!input.HasNormals()) {
         utility::LogError("Failed because input point cloud has no normal.");
     }
+
+    const bool filter_fpfh = indices.has_value();
+    std::vector<int> fpfh_indices;
+    if (filter_fpfh) {
+        std::vector<bool> mask_fpfh(input.points_.size(), false);
+        for (auto idx : indices.value()) {
+            if (idx < mask_fpfh.size()) {
+                if (!mask_fpfh[idx]) {
+                    mask_fpfh[idx] = true;
+                }
+            } else {
+                utility::LogWarning(
+                        "[ComputeFPFHFeature] contains index {} that is "
+                        "not within the bounds",
+                        idx);
+            }
+        }
+        fpfh_indices.reserve(indices.value().size());
+        for (size_t i = 0; i < mask_fpfh.size(); i++) {
+            if (mask_fpfh[i]) {
+                fpfh_indices.push_back(i);
+            }
+        }
+    }
+
+    const size_t n_fpfh =
+            filter_fpfh ? fpfh_indices.size() : input.points_.size();
+
     geometry::KDTreeFlann kdtree(input);
-    auto spfh = ComputeSPFHFeature(input, kdtree, search_param);
+
+    std::vector<size_t> spfh_indices;
+    std::vector<int> map_point_idx_to_spfh_idx;
+    std::vector<std::vector<int>> map_fpfh_idx_to_indices;
+    std::vector<std::vector<double>> map_fpfh_idx_to_distance2;
+    if (filter_fpfh) {
+        // compute neighbors of the selected points
+        // using vector<uint8_t> as a boolean mask for the parallel loop
+        // since vector<bool> is not thread safe in writing.
+        std::vector<uint8_t> mask_spfh(input.points_.size(), 0);
+        map_fpfh_idx_to_indices = std::vector<std::vector<int>>(n_fpfh);
+        map_fpfh_idx_to_distance2 = std::vector<std::vector<double>>(n_fpfh);
+#pragma omp parallel for schedule(static) \
+        num_threads(utility::EstimateMaxThreads())
+        for (int i = 0; i < (int)n_fpfh; i++) {
+            const auto &point = input.points_[fpfh_indices[i]];
+            std::vector<int> p_indices;
+            std::vector<double> p_distance2;
+            kdtree.Search(point, search_param, p_indices, p_distance2);
+            for (size_t k = 0; k < p_indices.size(); k++) {
+                if (!mask_spfh[p_indices[k]]) {
+                    mask_spfh[p_indices[k]] = 1;
+                }
+            }
+            map_fpfh_idx_to_indices[i] = std::move(p_indices);
+            map_fpfh_idx_to_distance2[i] = std::move(p_distance2);
+        }
+        size_t spfh_indices_reserve_factor;
+        switch (search_param.GetSearchType()) {
+            case geometry::KDTreeSearchParam::SearchType::Knn:
+                spfh_indices_reserve_factor =
+                        ((const geometry::KDTreeSearchParamKNN &)search_param)
+                                .knn_;
+                break;
+            case geometry::KDTreeSearchParam::SearchType::Hybrid:
+                spfh_indices_reserve_factor =
+                        ((const geometry::KDTreeSearchParamHybrid &)
+                                 search_param)
+                                .max_nn_;
+                break;
+            default:
+                spfh_indices_reserve_factor = 30;
+        }
+        spfh_indices.reserve(spfh_indices_reserve_factor * fpfh_indices.size());
+        map_point_idx_to_spfh_idx = std::vector<int>(input.points_.size(), -1);
+        for (size_t i = 0; i < mask_spfh.size(); i++) {
+            if (mask_spfh[i]) {
+                map_point_idx_to_spfh_idx[i] = spfh_indices.size();
+                spfh_indices.push_back(i);
+            }
+        }
+    }
+
+    auto feature = std::make_shared<Feature>();
+    feature->Resize(33, (int)n_fpfh);
+
+    auto spfh = filter_fpfh ? ComputeSPFHFeature(input, kdtree, search_param,
+                                                 spfh_indices)
+                            : ComputeSPFHFeature(input, kdtree, search_param);
     if (spfh == nullptr) {
         utility::LogError("Internal error: SPFH feature is nullptr.");
     }
 #pragma omp parallel for schedule(static) \
         num_threads(utility::EstimateMaxThreads())
-    for (int i = 0; i < (int)input.points_.size(); i++) {
-        const auto &point = input.points_[i];
-        std::vector<int> indices;
-        std::vector<double> distance2;
-        if (kdtree.Search(point, search_param, indices, distance2) > 1) {
+    for (int i = 0; i < (int)n_fpfh; i++) {
+        int i_spfh;
+        std::vector<int> p_indices;
+        std::vector<double> p_distance2;
+        if (filter_fpfh) {
+            i_spfh = map_point_idx_to_spfh_idx[fpfh_indices[i]];
+            p_indices = std::move(map_fpfh_idx_to_indices[i]);
+            p_distance2 = std::move(map_fpfh_idx_to_distance2[i]);
+        } else {
+            i_spfh = i;
+            kdtree.Search(input.points_[i], search_param, p_indices,
+                          p_distance2);
+        }
+        if (p_indices.size() > 1) {
             double sum[3] = {0.0, 0.0, 0.0};
-            for (size_t k = 1; k < indices.size(); k++) {
+            for (size_t k = 1; k < p_indices.size(); k++) {
                 // skip the point itself
-                double dist = distance2[k];
+                double dist = p_distance2[k];
                 if (dist == 0.0) continue;
+                int p_index_k =
+                        filter_fpfh ? map_point_idx_to_spfh_idx[p_indices[k]]
+                                    : p_indices[k];
                 for (int j = 0; j < 33; j++) {
-                    double val = spfh->data_(j, indices[k]) / dist;
+                    double val = spfh->data_(j, p_index_k) / dist;
                     sum[j / 11] += val;
                     feature->data_(j, i) += val;
                 }
@@ -157,10 +279,16 @@ std::shared_ptr<Feature> ComputeFPFHFeature(
                 // Our initial test shows that the full fpfh function in the
                 // paper seems to be better than PCL implementation. Further
                 // test required.
-                feature->data_(j, i) += spfh->data_(j, i);
+                feature->data_(j, i) += spfh->data_(j, i_spfh);
             }
         }
     }
+
+    utility::LogDebug(
+            "[ComputeFPFHFeature] Computed {:d} features from "
+            "input point cloud with {:d} points.",
+            (int)feature->data_.cols(), (int)input.points_.size());
+
     return feature;
 }
 

cpp/open3d/pipelines/registration/Feature.h

@@ -12,6 +12,7 @@
 #include <vector>
 
 #include "open3d/geometry/KDTreeSearchParam.h"
+#include "open3d/utility/Optional.h"
 
 namespace open3d {
 
@@ -59,10 +60,14 @@ class Feature {
 ///
 /// \param input The Input point cloud.
 /// \param search_param KDTree KNN search parameter.
+/// \param indices Indices of the points to compute FPFH features on.
+/// If not set, compute features for the whole point cloud.
 std::shared_ptr<Feature> ComputeFPFHFeature(
         const geometry::PointCloud &input,
         const geometry::KDTreeSearchParam &search_param =
-                geometry::KDTreeSearchParamKNN());
+                geometry::KDTreeSearchParamKNN(),
+        const utility::optional<std::vector<size_t>> &indices =
+                utility::nullopt);
 
 /// \brief Function to find correspondences via 1-nearest neighbor feature
 /// matching. Target is used to construct a nearest neighbor search

cpp/open3d/t/pipelines/kernel/Feature.cpp

@@ -20,19 +20,38 @@ void ComputeFPFHFeature(const core::Tensor &points,
                         const core::Tensor &indices,
                         const core::Tensor &distance2,
                         const core::Tensor &counts,
-                        core::Tensor &fpfhs) {
-    core::AssertTensorShape(fpfhs, {points.GetLength(), 33});
+                        core::Tensor &fpfhs,
+                        const utility::optional<core::Tensor> &mask,
+                        const utility::optional<core::Tensor>
+                                &map_batch_info_idx_to_point_idx) {
+    if (mask.has_value()) {
+        const int64_t size =
+                mask.value().To(core::Int64).Sum({0}).Item<int64_t>();
+        core::AssertTensorShape(fpfhs, {size, 33});
+        core::AssertTensorShape(mask.value(), {points.GetLength()});
+    } else {
+        core::AssertTensorShape(fpfhs, {points.GetLength(), 33});
+    }
+    if (map_batch_info_idx_to_point_idx.has_value()) {
+        core::AssertTensorShape(map_batch_info_idx_to_point_idx.value(),
+                                {counts.GetLength()});
+    }
     const core::Tensor points_d = points.Contiguous();
     const core::Tensor normals_d = normals.Contiguous();
     const core::Tensor counts_d = counts.To(core::Int32);
     if (points_d.IsCPU()) {
         ComputeFPFHFeatureCPU(points_d, normals_d, indices, distance2, counts_d,
-                              fpfhs);
+                              fpfhs, mask, map_batch_info_idx_to_point_idx);
     } else {
         core::CUDAScopedDevice scoped_device(points.GetDevice());
         CUDA_CALL(ComputeFPFHFeatureCUDA, points_d, normals_d, indices,
-                  distance2, counts_d, fpfhs);
+                  distance2, counts_d, fpfhs, mask,
+                  map_batch_info_idx_to_point_idx);
     }
+    utility::LogDebug(
+            "[ComputeFPFHFeature] Computed {:d} features from "
+            "input point cloud with {:d} points.",
+            (int)fpfhs.GetLength(), (int)points.GetLength());
 }
 
 }  // namespace kernel

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

@@ -7,33 +7,46 @@
 
 #include "open3d/core/CUDAUtils.h"
 #include "open3d/core/Tensor.h"
+#include "open3d/utility/Optional.h"
 
 namespace open3d {
 namespace t {
 namespace pipelines {
 namespace kernel {
 
-void ComputeFPFHFeature(const core::Tensor &points,
-                        const core::Tensor &normals,
-                        const core::Tensor &indices,
-                        const core::Tensor &distance2,
-                        const core::Tensor &counts,
-                        core::Tensor &fpfhs);
+void ComputeFPFHFeature(
+        const core::Tensor &points,
+        const core::Tensor &normals,
+        const core::Tensor &indices,
+        const core::Tensor &distance2,
+        const core::Tensor &counts,
+        core::Tensor &fpfhs,
+        const utility::optional<core::Tensor> &mask = utility::nullopt,
+        const utility::optional<core::Tensor> &map_batch_info_idx_to_point_idx =
+                utility::nullopt);
 
-void ComputeFPFHFeatureCPU(const core::Tensor &points,
-                           const core::Tensor &normals,
-                           const core::Tensor &indices,
-                           const core::Tensor &distance2,
-                           const core::Tensor &counts,
-                           core::Tensor &fpfhs);
+void ComputeFPFHFeatureCPU(
+        const core::Tensor &points,
+        const core::Tensor &normals,
+        const core::Tensor &indices,
+        const core::Tensor &distance2,
+        const core::Tensor &counts,
+        core::Tensor &fpfhs,
+        const utility::optional<core::Tensor> &mask = utility::nullopt,
+        const utility::optional<core::Tensor> &map_batch_info_idx_to_point_idx =
+                utility::nullopt);
 
 #ifdef BUILD_CUDA_MODULE
-void ComputeFPFHFeatureCUDA(const core::Tensor &points,
-                            const core::Tensor &normals,
-                            const core::Tensor &indices,
-                            const core::Tensor &distance2,
-                            const core::Tensor &counts,
-                            core::Tensor &fpfhs);
+void ComputeFPFHFeatureCUDA(
+        const core::Tensor &points,
+        const core::Tensor &normals,
+        const core::Tensor &indices,
+        const core::Tensor &distance2,
+        const core::Tensor &counts,
+        core::Tensor &fpfhs,
+        const utility::optional<core::Tensor> &mask = utility::nullopt,
+        const utility::optional<core::Tensor> &map_batch_info_idx_to_point_idx =
+                utility::nullopt);
 #endif
 
 }  // namespace kernel