From 583eaa3351f9411f2d2f1faf7ac17479cc11975b Mon Sep 17 00:00:00 2001 From: Nicola Loi Date: Mon, 30 Dec 2024 00:22:11 +0100 Subject: [PATCH 1/2] Add optional indices arg for fast computation of a small subset of FPFH features --- CHANGELOG.md | 1 + cpp/open3d/pipelines/registration/Feature.cpp | 178 +++++++++++++++--- cpp/open3d/pipelines/registration/Feature.h | 7 +- cpp/open3d/t/pipelines/kernel/Feature.cpp | 27 ++- cpp/open3d/t/pipelines/kernel/Feature.h | 49 +++-- cpp/open3d/t/pipelines/kernel/FeatureImpl.h | 115 +++++++++-- .../t/pipelines/registration/Feature.cpp | 134 ++++++++++--- cpp/open3d/t/pipelines/registration/Feature.h | 5 +- cpp/pybind/pipelines/registration/feature.cpp | 11 +- .../t/pipelines/registration/feature.cpp | 12 +- .../t/pipelines/registration/Feature.cpp | 40 +++- 11 files changed, 470 insertions(+), 109 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ad1f4f43b17..0b0bff57be8 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/cpp/open3d/pipelines/registration/Feature.cpp b/cpp/open3d/pipelines/registration/Feature.cpp index c43adbd585f..ac23e432240 100644 --- a/cpp/open3d/pipelines/registration/Feature.cpp +++ b/cpp/open3d/pipelines/registration/Feature.cpp @@ -21,23 +21,38 @@ namespace registration { std::shared_ptr Feature::SelectByIndex( const std::vector &indices, bool invert /* = false */) const { auto output = std::make_shared(); - output->Resize(data_.rows(), indices.size()); std::vector mask = std::vector(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(data_.cols()); i++) { + output->Resize(data_.rows(), n_output_features); + + for (size_t i = 0, current_col_feature = 0; + i < static_cast(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 ComputeSPFHFeature( const geometry::PointCloud &input, const geometry::KDTreeFlann &kdtree, - const geometry::KDTreeSearchParam &search_param) { + const geometry::KDTreeSearchParam &search_param, + const utility::optional> &indices = + utility::nullopt) { + const bool filter_spfh = indices.has_value(); + const auto spfh_indices = indices.value_or(std::vector()); + + const size_t n_spfh = + filter_spfh ? spfh_indices.size() : input.points_.size(); auto feature = std::make_shared(); - 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 indices; std::vector distance2; if (kdtree.Search(point, search_param, indices, distance2) > 1) { @@ -119,31 +143,129 @@ static std::shared_ptr ComputeSPFHFeature( std::shared_ptr ComputeFPFHFeature( const geometry::PointCloud &input, const geometry::KDTreeSearchParam - &search_param /* = geometry::KDTreeSearchParamKNN()*/) { - auto feature = std::make_shared(); - feature->Resize(33, (int)input.points_.size()); + &search_param /* = geometry::KDTreeSearchParamKNN()*/, + const utility::optional> + &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 fpfh_indices; + if (filter_fpfh) { + std::vector 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 spfh_indices; + std::vector map_point_idx_to_spfh_idx; + std::vector> map_fpfh_idx_to_indices; + std::vector> map_fpfh_idx_to_distance2; + if (filter_fpfh) { + // compute neighbors of the selected points + // using vector as a boolean mask for the parallel loop + // since vector is not thread safe in writing. + std::vector mask_spfh(input.points_.size(), 0); + map_fpfh_idx_to_indices = std::vector>(n_fpfh); + map_fpfh_idx_to_distance2 = std::vector>(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 p_indices; + std::vector 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(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->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 indices; - std::vector distance2; - if (kdtree.Search(point, search_param, indices, distance2) > 1) { + for (int i = 0; i < (int)n_fpfh; i++) { + int i_spfh; + std::vector p_indices; + std::vector 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 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; } diff --git a/cpp/open3d/pipelines/registration/Feature.h b/cpp/open3d/pipelines/registration/Feature.h index f4d9630688d..c7bafc021d5 100644 --- a/cpp/open3d/pipelines/registration/Feature.h +++ b/cpp/open3d/pipelines/registration/Feature.h @@ -12,6 +12,7 @@ #include #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 ComputeFPFHFeature( const geometry::PointCloud &input, const geometry::KDTreeSearchParam &search_param = - geometry::KDTreeSearchParamKNN()); + geometry::KDTreeSearchParamKNN(), + const utility::optional> &indices = + utility::nullopt); /// \brief Function to find correspondences via 1-nearest neighbor feature /// matching. Target is used to construct a nearest neighbor search diff --git a/cpp/open3d/t/pipelines/kernel/Feature.cpp b/cpp/open3d/t/pipelines/kernel/Feature.cpp index bd4014ecb5d..6601ef5432f 100644 --- a/cpp/open3d/t/pipelines/kernel/Feature.cpp +++ b/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 &mask, + const utility::optional + &map_batch_info_idx_to_point_idx) { + if (mask.has_value()) { + const int64_t size = + mask.value().To(core::Int64).Sum({0}).Item(); + 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 diff --git a/cpp/open3d/t/pipelines/kernel/Feature.h b/cpp/open3d/t/pipelines/kernel/Feature.h index 6dd4a405342..90ddfb0ae6e 100644 --- a/cpp/open3d/t/pipelines/kernel/Feature.h +++ b/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 &mask = utility::nullopt, + const utility::optional &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 &mask = utility::nullopt, + const utility::optional &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 &mask = utility::nullopt, + const utility::optional &map_batch_info_idx_to_point_idx = + utility::nullopt); #endif } // namespace kernel diff --git a/cpp/open3d/t/pipelines/kernel/FeatureImpl.h b/cpp/open3d/t/pipelines/kernel/FeatureImpl.h index bd00f29fb77..9a56e7b72e4 100644 --- a/cpp/open3d/t/pipelines/kernel/FeatureImpl.h +++ b/cpp/open3d/t/pipelines/kernel/FeatureImpl.h @@ -114,11 +114,69 @@ void ComputeFPFHFeatureCPU const core::Tensor &indices, const core::Tensor &distance2, const core::Tensor &counts, - core::Tensor &fpfhs) { + core::Tensor &fpfhs, + const utility::optional &mask, + const utility::optional + &map_batch_info_idx_to_point_idx) { const core::Dtype dtype = points.GetDtype(); - const int64_t n = points.GetLength(); + const core::Device device = points.GetDevice(); + const int64_t n_points = points.GetLength(); - core::Tensor spfhs = fpfhs.Clone(); + const bool filter_fpfh = + mask.has_value() & map_batch_info_idx_to_point_idx.has_value(); + if (mask.has_value() ^ map_batch_info_idx_to_point_idx.has_value()) { + utility::LogError( + "Parameters mask and map_batch_info_idx_to_point_idx must be " + "both " + "provided or both not provided."); + } + if (filter_fpfh) { + if (mask.value().GetShape()[0] != n_points) { + utility::LogError( + "Parameter mask was provided, but its size {:d} should" + "be equal to the number of points {:d}.", + (int)mask.value().GetShape()[0], n_points); + } + if (map_batch_info_idx_to_point_idx.value().GetShape()[0] != + counts.GetShape()[0]) { + utility::LogError( + "Parameter map_batch_info_idx_to_point_idx was provided, " + "but its size" + "{:d} should be equal to the size of counts {:d}.", + (int)map_batch_info_idx_to_point_idx.value().GetShape()[0], + (int)counts.GetShape()[0]); + } + } + + core::Tensor map_spfh_info_idx_to_point_idx = + map_batch_info_idx_to_point_idx.value_or( + core::Tensor::Empty({0}, core::Int64, device)); + + const core::Tensor map_fpfh_idx_to_point_idx = + filter_fpfh ? mask.value().NonZero().GetItem( + {core::TensorKey::Index(0)}) + : core::Tensor::Empty({0}, core::Int64, device); + + const int32_t n_fpfh = + filter_fpfh ? map_fpfh_idx_to_point_idx.GetLength() : n_points; + const int32_t n_spfh = + filter_fpfh ? map_spfh_info_idx_to_point_idx.GetLength() : n_points; + + core::Tensor spfhs = + core::Tensor::Zeros({n_spfh, 33}, dtype, fpfhs.GetDevice()); + + core::Tensor map_point_idx_to_spfh_idx; + if (filter_fpfh) { + map_point_idx_to_spfh_idx = core::Tensor::Full( + {n_points}, -1, core::Int64, fpfhs.GetDevice()); + map_point_idx_to_spfh_idx.IndexSet( + {map_spfh_info_idx_to_point_idx}, + core::Tensor::Arange(0, n_spfh, 1, core::Int64, + fpfhs.GetDevice())); + } else { + map_point_idx_to_spfh_idx = + core::Tensor::Empty({0}, core::Int64, fpfhs.GetDevice()); + } // Check the nns type (knn = hybrid = false, radius = true). // The nns radius search mode will resulting a prefix sum 1D tensor. @@ -139,11 +197,22 @@ void ComputeFPFHFeatureCPU const int32_t *counts_ptr = counts.GetDataPtr(); scalar_t *spfhs_ptr = spfhs.GetDataPtr(); scalar_t *fpfhs_ptr = fpfhs.GetDataPtr(); + const int64_t *map_spfh_info_idx_to_point_idx_ptr = + map_spfh_info_idx_to_point_idx.GetDataPtr(); + const int64_t *map_fpfh_idx_to_point_idx_ptr = + map_fpfh_idx_to_point_idx.GetDataPtr(); + const int64_t *map_point_idx_to_spfh_idx_ptr = + map_point_idx_to_spfh_idx.GetDataPtr(); - // Compute SPFH features for each point. + // Compute SPFH features for the points. core::ParallelFor( - points.GetDevice(), n, [=] OPEN3D_DEVICE(int64_t workload_idx) { - int64_t idx = 3 * workload_idx; + points.GetDevice(), n_spfh, + [=] OPEN3D_DEVICE(int64_t workload_idx) { + int64_t workload_point_idx = + filter_fpfh ? map_spfh_info_idx_to_point_idx_ptr + [workload_idx] + : workload_idx; + int64_t idx = 3 * workload_point_idx; const scalar_t *point = points_ptr + idx; const scalar_t *normal = normals_ptr + idx; @@ -178,28 +247,36 @@ void ComputeFPFHFeatureCPU } }); - // Compute FPFH features for each point. + // Compute FPFH features for the points. core::ParallelFor( - points.GetDevice(), n, [=] OPEN3D_DEVICE(int64_t workload_idx) { + points.GetDevice(), n_fpfh, + [=] OPEN3D_DEVICE(int64_t workload_idx) { + int64_t workload_spfh_idx = + filter_fpfh ? map_point_idx_to_spfh_idx_ptr + [map_fpfh_idx_to_point_idx_ptr + [workload_idx]] + : workload_idx; const int indice_size = - is_radius_search ? (counts_ptr[workload_idx + 1] - - counts_ptr[workload_idx]) - : counts_ptr[workload_idx]; - + is_radius_search + ? (counts_ptr[workload_spfh_idx + 1] - + counts_ptr[workload_spfh_idx]) + : counts_ptr[workload_spfh_idx]; if (indice_size > 1) { scalar_t sum[3] = {0.0, 0.0, 0.0}; for (int i = 1; i < indice_size; i++) { const int idx = is_radius_search - ? i + counts_ptr[workload_idx] - : workload_idx * nn_size + i; + ? i + counts_ptr[workload_spfh_idx] + : workload_spfh_idx * nn_size + i; const scalar_t dist = distance2_ptr[idx]; if (dist == 0.0) continue; - + const int32_t spfh_idx = + filter_fpfh ? map_point_idx_to_spfh_idx_ptr + [indices_ptr[idx]] + : indices_ptr[idx]; for (int j = 0; j < 33; j++) { const scalar_t val = - spfhs_ptr[indices_ptr[idx] * 33 + j] / - dist; + spfhs_ptr[spfh_idx * 33 + j] / dist; sum[j / 11] += val; fpfhs_ptr[workload_idx * 33 + j] += val; } @@ -210,12 +287,12 @@ void ComputeFPFHFeatureCPU for (int j = 0; j < 33; j++) { fpfhs_ptr[workload_idx * 33 + j] *= sum[j / 11]; fpfhs_ptr[workload_idx * 33 + j] += - spfhs_ptr[workload_idx * 33 + j]; + spfhs_ptr[workload_spfh_idx * 33 + j]; } } }); }); -} // namespace kernel +} } // namespace kernel } // namespace pipelines diff --git a/cpp/open3d/t/pipelines/registration/Feature.cpp b/cpp/open3d/t/pipelines/registration/Feature.cpp index 8fc78a535ef..5db1170abe9 100644 --- a/cpp/open3d/t/pipelines/registration/Feature.cpp +++ b/cpp/open3d/t/pipelines/registration/Feature.cpp @@ -18,9 +18,11 @@ namespace t { namespace pipelines { namespace registration { -core::Tensor ComputeFPFHFeature(const geometry::PointCloud &input, - const utility::optional max_nn, - const utility::optional radius) { +core::Tensor ComputeFPFHFeature( + const geometry::PointCloud &input, + const utility::optional max_nn, + const utility::optional radius, + const utility::optional &indices) { core::AssertTensorDtypes(input.GetPointPositions(), {core::Float64, core::Float32}); if (max_nn.has_value() && max_nn.value() <= 3) { @@ -37,44 +39,106 @@ core::Tensor ComputeFPFHFeature(const geometry::PointCloud &input, const core::Dtype dtype = input.GetPointPositions().GetDtype(); const core::Device device = input.GetPointPositions().GetDevice(); - // Compute nearest neighbors and squared distances. - core::Tensor indices, distance2, counts; core::nns::NearestNeighborSearch tree(input.GetPointPositions(), core::Int32); + bool tree_set = false; + + core::Tensor mask_fpfh_points; + core::Tensor mask_required_points; + if (indices.has_value()) { + mask_fpfh_points = + core::Tensor::Zeros({num_points}, core::Bool, device); + mask_required_points = + core::Tensor::Zeros({num_points}, core::Bool, device); + core::Tensor indices_tmp, distance2_tmp, counts_tmp; + mask_fpfh_points.IndexSet({indices.value()}, + core::Tensor::Ones({1}, core::Bool, device)); + const core::Tensor query_point_positions = + input.GetPointPositions().IndexGet({indices.value()}); + if (radius.has_value() && max_nn.has_value()) { + tree_set = tree.HybridIndex(radius.value()); + if (!tree_set) { + utility::LogError("Building HybridIndex failed."); + } + std::tie(indices_tmp, distance2_tmp, counts_tmp) = + tree.HybridSearch(query_point_positions, radius.value(), + max_nn.value()); + } else if (!radius.has_value() && max_nn.has_value()) { + tree_set = tree.KnnIndex(); + if (!tree_set) { + utility::LogError("Building KnnIndex failed."); + } + std::tie(indices_tmp, distance2_tmp) = + tree.KnnSearch(query_point_positions, max_nn.value()); + } else if (radius.has_value() && !max_nn.has_value()) { + tree_set = tree.FixedRadiusIndex(radius.value()); + if (!tree_set) { + utility::LogError("Building RadiusIndex failed."); + } + std::tie(indices_tmp, distance2_tmp, counts_tmp) = + tree.FixedRadiusSearch(query_point_positions, + radius.value()); + } else { + utility::LogError("Both max_nn and radius are none."); + } + + indices_tmp = indices_tmp.To(core::Int64).View({-1}); + mask_required_points.IndexSet( + {indices_tmp}, core::Tensor::Ones({1}, core::Bool, device)); + + } else { + mask_fpfh_points = core::Tensor::Zeros({0}, core::Bool, device); + mask_required_points = core::Tensor::Zeros({0}, core::Bool, device); + } + + const core::Tensor query_point_positions = + mask_required_points.GetShape()[0] > 0 + ? input.GetPointPositions().IndexGet({mask_required_points}) + : input.GetPointPositions(); + + // Compute nearest neighbors and squared distances. + core::Tensor p_indices, p_distance2, p_counts; + if (radius.has_value() && max_nn.has_value()) { - bool check = tree.HybridIndex(radius.value()); - if (!check) { - utility::LogError("Building HybridIndex failed."); + if (!tree_set) { + tree_set = tree.HybridIndex(radius.value()); + if (!tree_set) { + utility::LogError("Building HybridIndex failed."); + } } - std::tie(indices, distance2, counts) = tree.HybridSearch( - input.GetPointPositions(), radius.value(), max_nn.value()); + std::tie(p_indices, p_distance2, p_counts) = tree.HybridSearch( + query_point_positions, radius.value(), max_nn.value()); utility::LogDebug( "Use HybridSearch [max_nn: {} | radius {}] for computing FPFH " "feature.", max_nn.value(), radius.value()); } else if (!radius.has_value() && max_nn.has_value()) { - bool check = tree.KnnIndex(); - if (!check) { - utility::LogError("Building KnnIndex failed."); + if (!tree_set) { + tree_set = tree.KnnIndex(); + if (!tree_set) { + utility::LogError("Building KnnIndex failed."); + } } - std::tie(indices, distance2) = - tree.KnnSearch(input.GetPointPositions(), max_nn.value()); + std::tie(p_indices, p_distance2) = + tree.KnnSearch(query_point_positions, max_nn.value()); // Make counts full with min(max_nn, num_points). const int fill_value = max_nn.value() > num_points ? num_points : max_nn.value(); - counts = core::Tensor::Full({num_points}, fill_value, core::Int32, - device); + p_counts = core::Tensor::Full({query_point_positions.GetLength()}, + fill_value, core::Int32, device); utility::LogDebug( "Use KNNSearch [max_nn: {}] for computing FPFH feature.", max_nn.value()); } else if (radius.has_value() && !max_nn.has_value()) { - bool check = tree.FixedRadiusIndex(radius.value()); - if (!check) { - utility::LogError("Building RadiusIndex failed."); + if (!tree_set) { + tree_set = tree.FixedRadiusIndex(radius.value()); + if (!tree_set) { + utility::LogError("Building RadiusIndex failed."); + } } - std::tie(indices, distance2, counts) = tree.FixedRadiusSearch( - input.GetPointPositions(), radius.value()); + std::tie(p_indices, p_distance2, p_counts) = + tree.FixedRadiusSearch(query_point_positions, radius.value()); utility::LogDebug( "Use RadiusSearch [radius: {}] for computing FPFH feature.", radius.value()); @@ -82,12 +146,26 @@ core::Tensor ComputeFPFHFeature(const geometry::PointCloud &input, utility::LogError("Both max_nn and radius are none."); } - const int64_t size = input.GetPointPositions().GetLength(); - - core::Tensor fpfh = core::Tensor::Zeros({size, 33}, dtype, device); - pipelines::kernel::ComputeFPFHFeature(input.GetPointPositions(), - input.GetPointNormals(), indices, - distance2, counts, fpfh); + core::Tensor fpfh; + if (indices.has_value()) { + const auto mask_fpfh_points_indices = + mask_fpfh_points.NonZero().GetItem({core::TensorKey::Index(0)}); + const auto map_batch_info_idx_to_point_idx = + mask_required_points.NonZero().GetItem( + {core::TensorKey::Index(0)}); + fpfh = core::Tensor::Zeros({mask_fpfh_points_indices.GetLength(), 33}, + dtype, device); + pipelines::kernel::ComputeFPFHFeature( + input.GetPointPositions(), input.GetPointNormals(), p_indices, + p_distance2, p_counts, fpfh, mask_fpfh_points, + map_batch_info_idx_to_point_idx); + } else { + const int64_t size = input.GetPointPositions().GetLength(); + fpfh = core::Tensor::Zeros({size, 33}, dtype, device); + pipelines::kernel::ComputeFPFHFeature( + input.GetPointPositions(), input.GetPointNormals(), p_indices, + p_distance2, p_counts, fpfh); + } return fpfh; } diff --git a/cpp/open3d/t/pipelines/registration/Feature.h b/cpp/open3d/t/pipelines/registration/Feature.h index 65a88408833..6b544d71334 100644 --- a/cpp/open3d/t/pipelines/registration/Feature.h +++ b/cpp/open3d/t/pipelines/registration/Feature.h @@ -30,12 +30,15 @@ namespace registration { /// 100]. /// \param radius [optional] Neighbor search radius parameter. [Recommended ~5x /// voxel size]. +/// \param indices [optional] Tensor with the indices of the points to compute " +/// FPFH features on. [Default = None]. /// \return A Tensor of FPFH feature of the input point cloud with /// shape {N, 33}, data type and device same as input. core::Tensor ComputeFPFHFeature( const geometry::PointCloud &input, const utility::optional max_nn = 100, - const utility::optional radius = utility::nullopt); + const utility::optional radius = utility::nullopt, + const utility::optional &indices = utility::nullopt); /// \brief Function to find correspondences via 1-nearest neighbor feature /// matching. Target is used to construct a nearest neighbor search diff --git a/cpp/pybind/pipelines/registration/feature.cpp b/cpp/pybind/pipelines/registration/feature.cpp index 5ab1ce3cd2c..8f1697a6927 100644 --- a/cpp/pybind/pipelines/registration/feature.cpp +++ b/cpp/pybind/pipelines/registration/feature.cpp @@ -7,6 +7,7 @@ #include "open3d/pipelines/registration/Feature.h" +#include "open3d/geometry/KDTreeSearchParam.h" #include "open3d/geometry/PointCloud.h" #include "pybind/docstring.h" #include "pybind/pipelines/registration/registration.h" @@ -63,11 +64,15 @@ void pybind_feature_definitions(py::module &m_registration) { "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); + "input"_a, "search_param"_a, "indices"_a = py::none()); docstring::FunctionDocInject( m_registration, "compute_fpfh_feature", - {{"input", "The Input point cloud."}, - {"search_param", "KDTree KNN search parameter."}}); + { + {"input", "The Input point cloud."}, + {"search_param", "KDTree KNN search parameter."}, + {"indices", + "Indices to select points for feature computation."}, + }); m_registration.def( "correspondences_from_features", &CorrespondencesFromFeatures, diff --git a/cpp/pybind/t/pipelines/registration/feature.cpp b/cpp/pybind/t/pipelines/registration/feature.cpp index b7e9e1d5ca5..8b4631084d7 100644 --- a/cpp/pybind/t/pipelines/registration/feature.cpp +++ b/cpp/pybind/t/pipelines/registration/feature.cpp @@ -23,8 +23,11 @@ void pybind_feature_definitions(py::module &m_registration) { R"(Function to compute FPFH feature for a point cloud. It uses KNN search (Not recommended to use on GPU) if only max_nn parameter is provided, Radius search (Not recommended to use on GPU) if only radius -parameter is provided, and Hybrid search (Recommended) if both are provided.)", - "input"_a, "max_nn"_a = 100, "radius"_a = py::none()); +parameter is provided, and Hybrid search (Recommended) if both are provided. +If indices is provided, the function will compute FPFH features only on the +selected points.)", + "input"_a, "max_nn"_a = 100, "radius"_a = py::none(), + "indices"_a = py::none()); docstring::FunctionDocInject( m_registration, "compute_fpfh_feature", {{"input", @@ -34,7 +37,10 @@ parameter is provided, and Hybrid search (Recommended) if both are provided.)", "100]"}, {"radius", "[optional] Neighbor search radius parameter. [Recommended ~5x " - "voxel size]"}}); + "voxel size]"}, + {"indices", + "[optional] Tensor with the indices of the points to compute " + "FPFH features on. [Default = None]"}}); m_registration.def( "correspondences_from_features", &CorrespondencesFromFeatures, diff --git a/cpp/tests/t/pipelines/registration/Feature.cpp b/cpp/tests/t/pipelines/registration/Feature.cpp index d4adc0958b3..a3741a3ed8a 100644 --- a/cpp/tests/t/pipelines/registration/Feature.cpp +++ b/cpp/tests/t/pipelines/registration/Feature.cpp @@ -42,19 +42,21 @@ TEST_P(FeaturePermuteDevices, SelectByIndex) { 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( - {53, 194, 839, 2543, 6391, 29483}, device)); - const auto selected_fpfh_t = fpfh_t.GetItem(selected_indeces_t); + const std::vector indices = {53, 6391, 194, 31037, 15936, + 345, 6839, 2543, 29483}; + const auto selected_fpfh = fpfh->SelectByIndex(indices, false); + std::vector sorted_indices(indices.begin(), indices.end()); + std::sort(sorted_indices.begin(), sorted_indices.end()); + const auto indices_t = core::TensorKey::IndexTensor(core::Tensor( + sorted_indices, {(int)sorted_indices.size()}, core::Int64, device)); + const auto selected_fpfh_t = fpfh_t.GetItem(indices_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)); + 1e-6, 1e-6)); } TEST_P(FeaturePermuteDevices, ComputeFPFHFeature) { @@ -79,6 +81,30 @@ TEST_P(FeaturePermuteDevices, ComputeFPFHFeature) { .T() .To(fpfh_t.GetDevice(), fpfh_t.GetDtype()), 1e-4, 1e-4)); + + const std::vector indices = {4, 27403, 103, 9172, 5728, 839, + 12943, 28, 9374, 17837, 7390, 473, + 11836, 26362, 3046, 35027, 5738}; + const auto selected_fpfh_by_index = fpfh->SelectByIndex(indices); + const auto computed_fpfh_by_index = + pipelines::registration::ComputeFPFHFeature( + pcd_legacy, geometry::KDTreeSearchParamHybrid(0.01, 100), + indices); + + EXPECT_TRUE(selected_fpfh_by_index->data_.isApprox( + computed_fpfh_by_index->data_, 1e-4)); + + std::vector sorted_indices(indices.begin(), indices.end()); + std::sort(sorted_indices.begin(), sorted_indices.end()); + const auto indices_t = core::Tensor( + sorted_indices, {(int)sorted_indices.size()}, core::Int64, device); + const auto selected_fpfh_t_by_index = fpfh_t.IndexGet({indices_t}); + const auto computed_fpfh_t_by_index = + t::pipelines::registration::ComputeFPFHFeature(pcd, 100, 0.01, + indices_t); + + EXPECT_TRUE(selected_fpfh_t_by_index.AllClose(computed_fpfh_t_by_index, + 1e-4, 1e-4)); } TEST_P(FeaturePermuteDevices, CorrespondencesFromFeatures) { From 6dcd9d9a93e245b9aaf37b4cc5b07d644db5b14f Mon Sep 17 00:00:00 2001 From: Nicola Loi Date: Mon, 30 Dec 2024 11:26:29 +0100 Subject: [PATCH 2/2] fix for Windows and Mac CI --- cpp/open3d/pipelines/registration/Feature.cpp | 4 ++-- cpp/open3d/t/pipelines/kernel/FeatureImpl.h | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/open3d/pipelines/registration/Feature.cpp b/cpp/open3d/pipelines/registration/Feature.cpp index ac23e432240..f42fed61df8 100644 --- a/cpp/open3d/pipelines/registration/Feature.cpp +++ b/cpp/open3d/pipelines/registration/Feature.cpp @@ -185,9 +185,9 @@ std::shared_ptr ComputeFPFHFeature( std::vector> map_fpfh_idx_to_distance2; if (filter_fpfh) { // compute neighbors of the selected points - // using vector as a boolean mask for the parallel loop + // using vector as a boolean mask for the parallel loop // since vector is not thread safe in writing. - std::vector mask_spfh(input.points_.size(), 0); + std::vector mask_spfh(input.points_.size(), 0); map_fpfh_idx_to_indices = std::vector>(n_fpfh); map_fpfh_idx_to_distance2 = std::vector>(n_fpfh); #pragma omp parallel for schedule(static) \ diff --git a/cpp/open3d/t/pipelines/kernel/FeatureImpl.h b/cpp/open3d/t/pipelines/kernel/FeatureImpl.h index 9a56e7b72e4..657f9878e14 100644 --- a/cpp/open3d/t/pipelines/kernel/FeatureImpl.h +++ b/cpp/open3d/t/pipelines/kernel/FeatureImpl.h @@ -123,7 +123,7 @@ void ComputeFPFHFeatureCPU const int64_t n_points = points.GetLength(); const bool filter_fpfh = - mask.has_value() & map_batch_info_idx_to_point_idx.has_value(); + mask.has_value() && map_batch_info_idx_to_point_idx.has_value(); if (mask.has_value() ^ map_batch_info_idx_to_point_idx.has_value()) { utility::LogError( "Parameters mask and map_batch_info_idx_to_point_idx must be "