Add select_by_index method to Feature class (#7039)

* Add select_by_index method to Feature class

* better comments

---------

Co-authored-by: Benjamin Ummenhofer <[email protected]>

Author: nicolaloi

CHANGELOG.md

@@ -52,6 +52,7 @@
 -   Fix alpha shape reconstruction if alpha too small for point scale (PR #6998)
 -   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)
 
 ## 0.13
 

cpp/open3d/pipelines/registration/Feature.cpp

@@ -18,6 +18,31 @@ namespace open3d {
 namespace pipelines {
 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);
+    for (size_t i : indices) {
+        mask[i] = !invert;
+    }
+
+    size_t current_col_feature = 0;
+    for (size_t i = 0; i < static_cast<size_t>(data_.cols()); i++) {
+        if (mask[i]) {
+            output->data_.col(current_col_feature) = data_.col(i);
+            current_col_feature++;
+        }
+    }
+
+    utility::LogDebug(
+            "Feature group down sampled from {:d} features to {:d} features.",
+            (int)data_.cols(), (int)output->data_.cols());
+
+    return output;
+}
+
 static Eigen::Vector4d ComputePairFeatures(const Eigen::Vector3d &p1,
                                            const Eigen::Vector3d &n1,
                                            const Eigen::Vector3d &p2,

cpp/open3d/pipelines/registration/Feature.h

@@ -42,6 +42,14 @@ class Feature {
     /// Returns number of points.
     size_t Num() const { return data_.cols(); }
 
+    /// \brief Selects features from \p input Feature group, with indices in \p
+    /// indices, and returns a new Feature group with selected features.
+    ///
+    /// \param indices Indices of features to be selected.
+    /// \param invert Set to `True` to invert the selection of indices.
+    std::shared_ptr<Feature> SelectByIndex(const std::vector<size_t> &indices,
+                                           bool invert = false) const;
+
 public:
     /// Data buffer storing features.
     Eigen::MatrixXd data_;

cpp/pybind/pipelines/registration/feature.cpp

@@ -31,6 +31,10 @@ void pybind_feature_definitions(py::module &m_registration) {
             .def("dimension", &Feature::Dimension,
                  "Returns feature dimensions per point.")
             .def("num", &Feature::Num, "Returns number of points.")
+            .def("select_by_index", &Feature::SelectByIndex,
+                 "Function to select features from input Feature group into "
+                 "output Feature group.",
+                 "indices"_a, "invert"_a = false)
             .def_readwrite("data", &Feature::data_,
                            "``dim x n`` float64 numpy array: Data buffer "
                            "storing features.")
@@ -48,6 +52,11 @@ void pybind_feature_definitions(py::module &m_registration) {
     docstring::ClassMethodDocInject(m_registration, "Feature", "resize",
                                     {{"dim", "Feature dimension per point."},
                                      {"n", "Number of points."}});
+    docstring::ClassMethodDocInject(
+            m_registration, "Feature", "select_by_index",
+            {{"indices", "Indices of features to be selected."},
+             {"invert",
+              "Set to ``True`` to invert the selection of indices."}});
     m_registration.def("compute_fpfh_feature", &ComputeFPFHFeature,
                        "Function to compute FPFH feature for a point cloud",
                        "input"_a, "search_param"_a);

cpp/tests/t/pipelines/registration/Feature.cpp

@@ -25,6 +25,38 @@ INSTANTIATE_TEST_SUITE_P(Feature,
                          FeaturePermuteDevices,
                          testing::ValuesIn(PermuteDevices::TestCases()));
 
+TEST_P(FeaturePermuteDevices, SelectByIndex) {
+    core::Device device = GetParam();
+
+    open3d::geometry::PointCloud pcd_legacy;
+    data::BunnyMesh bunny;
+    open3d::io::ReadPointCloud(bunny.GetPath(), pcd_legacy);
+
+    pcd_legacy.EstimateNormals();
+    // Convert to float64 to avoid precision loss.
+    const auto pcd = t::geometry::PointCloud::FromLegacy(pcd_legacy,
+                                                         core::Float64, device);
+
+    const auto fpfh = pipelines::registration::ComputeFPFHFeature(
+            pcd_legacy, geometry::KDTreeSearchParamHybrid(0.01, 100));
+    const auto fpfh_t =
+            t::pipelines::registration::ComputeFPFHFeature(pcd, 100, 0.01);
+
+    const auto selected_fpfh =
+            fpfh->SelectByIndex({53, 194, 839, 2543, 6391, 29483}, false);
+    const auto selected_indeces_t =
+            core::TensorKey::IndexTensor(core::Tensor::Init<int64_t>(
+                    {53, 194, 839, 2543, 6391, 29483}, device));
+    const auto selected_fpfh_t = fpfh_t.GetItem(selected_indeces_t);
+
+    EXPECT_TRUE(selected_fpfh_t.AllClose(
+            core::eigen_converter::EigenMatrixToTensor(selected_fpfh->data_)
+                    .T()
+                    .To(selected_fpfh_t.GetDevice(),
+                        selected_fpfh_t.GetDtype()),
+            1e-4, 1e-4));
+}
+
 TEST_P(FeaturePermuteDevices, ComputeFPFHFeature) {
     core::Device device = GetParam();