From 8e665814a9fcbc3b8a3e5bde135917f586abed6e Mon Sep 17 00:00:00 2001 From: koide3 <31344317+koide3@users.noreply.github.com> Date: Tue, 24 Sep 2024 18:04:55 +0900 Subject: [PATCH] projective search (#92) --- include/small_gicp/ann/projective_search.hpp | 219 +++++++++++++++++++ 1 file changed, 219 insertions(+) create mode 100644 include/small_gicp/ann/projective_search.hpp diff --git a/include/small_gicp/ann/projective_search.hpp b/include/small_gicp/ann/projective_search.hpp new file mode 100644 index 0000000..3a542ef --- /dev/null +++ b/include/small_gicp/ann/projective_search.hpp @@ -0,0 +1,219 @@ +// SPDX-FileCopyrightText: Copyright 2024 Kenji Koide +// SPDX-License-Identifier: MIT +#pragma once + +#include +#include +#include +#include + +namespace small_gicp { + +/// @brief Equirectangular projection. +struct EquirectangularProjection { +public: + /// @brief Project the point into the normalized image coordinates. (u, v) in ([0, 1], [0, 1]) + Eigen::Vector2d operator()(const Eigen::Vector3d& pt_3d) const { + if (pt_3d.squaredNorm() < 1e-3) { + return Eigen::Vector2d(0.5, 0.5); + } + + const Eigen::Vector3d bearing = pt_3d.normalized(); + const double lat = -std::asin(bearing[1]); + const double lon = std::atan2(bearing[0], bearing[2]); + + return Eigen::Vector2d(lon / (2.0 * M_PI) + 0.5, lat / M_PI + 0.5); + }; +}; + +/// @brief Border clamp mode. Points out of the border are discarded. +struct BorderClamp { +public: + int operator()(int x, int width) const { return x; } +}; + +/// @brief Border repeat mode. Points out of the border are wrapped around. +struct BorderRepeat { +public: + int operator()(int x, int width) const { return x < 0 ? x + width : (x >= width ? x - width : x); } +}; + +/// @brief "Unsafe" projective search. This class does not hold the ownership of the target point cloud. +template +struct UnsafeProjectiveSearch { +public: + /// @brief Constructor. + /// @param width Index map width + /// @param height Index map height + /// @param points Target point cloud + UnsafeProjectiveSearch(int width, int height, const PointCloud& points) : points(points), index_map(height, width), search_window_h(10), search_window_v(5) { + index_map.setConstant(invalid_index); + + Projection project; + for (size_t i = 0; i < traits::size(points); ++i) { + const Eigen::Vector4d pt = traits::point(points, i); + const Eigen::Vector2d uv = project(pt.head<3>()); + const int u = uv[0] * index_map.cols(); + const int v = uv[1] * index_map.rows(); + + if (u < 0 || u >= index_map.cols() || v < 0 || v >= index_map.rows()) { + continue; + } + index_map(v, u) = i; + } + } + + /// @brief Find the nearest neighbor. + /// @param query Query point + /// @param k_indices Index of the nearest neighbor (uninitialized if not found) + /// @param k_sq_dists Squared distance to the nearest neighbor (uninitialized if not found) + /// @param setting KNN search setting + /// @return Number of found neighbors (0 or 1) + size_t nearest_neighbor_search(const Eigen::Vector4d& query, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const { + return knn_search<1>(query, k_indices, k_sq_dists, setting); + } + + /// @brief Find k-nearest neighbors. This method uses dynamic memory allocation. + /// @param query Query point + /// @param k Number of neighbors + /// @param k_indices Indices of neighbors + /// @param k_sq_dists Squared distances to neighbors (sorted in ascending order) + /// @param setting KNN search setting + /// @return Number of found neighbors + size_t knn_search(const Eigen::Vector4d& query, int k, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const { + KnnResult<-1> result(k_indices, k_sq_dists, k); + knn_search(query, result, setting); + return result.num_found(); + } + + /// @brief Find k-nearest neighbors. This method uses fixed and static memory allocation. Might be faster for small k. + /// @param query Query point + /// @param k_indices Indices of neighbors + /// @param k_sq_dists Squared distances to neighbors (sorted in ascending order) + /// @param setting KNN search setting + /// @return Number of found neighbors + template + size_t knn_search(const Eigen::Vector4d& query, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const { + KnnResult result(k_indices, k_sq_dists); + knn_search(query, result, setting); + return result.num_found(); + } + +private: + template + void knn_search(const Eigen::Vector4d& query, Result& result, const KnnSetting& setting) const { + BorderModeH border_h; + BorderModeV border_v; + + Projection project; + const Eigen::Vector2d uv = project(query.head<3>()); + const int u = uv[0] * index_map.cols(); + const int v = uv[1] * index_map.rows(); + + for (int dv = -search_window_v; dv <= search_window_v; dv++) { + const int v_clamped = border_v(v + dv, index_map.rows()); + if (v_clamped < 0 || v_clamped >= index_map.rows()) { + continue; + } + + for (int du = -search_window_h; du <= search_window_h; du++) { + const int u_clamped = border_h(u + du, index_map.cols()); + if (u_clamped < 0 || u_clamped >= index_map.cols()) { + continue; + } + + const auto index = index_map(v_clamped, u_clamped); + if (index == invalid_index) { + continue; + } + + const double sq_dist = (traits::point(points, index) - query).squaredNorm(); + result.push(index, sq_dist); + + if (setting.fulfilled(result)) { + return; + } + } + } + } + +public: + static constexpr std::uint32_t invalid_index = std::numeric_limits::max(); + + const PointCloud& points; + Eigen::Matrix index_map; + + int search_window_h; + int search_window_v; +}; + +/// @brief "Safe" projective search. This class holds the ownership of the target point cloud. +template +struct ProjectiveSearch { +public: + using Ptr = std::shared_ptr>; + using ConstPtr = std::shared_ptr>; + + explicit ProjectiveSearch(int width, int height, std::shared_ptr points) : points(points), search(width, height, *points) {} + + /// @brief Find k-nearest neighbors. This method uses dynamic memory allocation. + /// @param query Query point + /// @param k Number of neighbors + /// @param k_indices Indices of neighbors + /// @param k_sq_dists Squared distances to neighbors (sorted in ascending order) + /// @param setting KNN search setting + /// @return Number of found neighbors + size_t nearest_neighbor_search(const Eigen::Vector4d& query, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const { + return search.nearest_neighbor_search(query, k_indices, k_sq_dists, setting); + } + + /// @brief Find k-nearest neighbors. This method uses dynamic memory allocation. + /// @param query Query point + /// @param k Number of neighbors + /// @param k_indices Indices of neighbors + /// @param k_sq_dists Squared distances to neighbors (sorted in ascending order) + /// @param setting KNN search setting + /// @return Number of found neighbors + size_t knn_search(const Eigen::Vector4d& query, size_t k, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const { + return search.knn_search(query, k, k_indices, k_sq_dists, setting); + } + +public: + const std::shared_ptr points; ///< Points + const UnsafeProjectiveSearch search; ///< Search +}; + +namespace traits { + +template +struct Traits> { + static size_t nearest_neighbor_search( + const UnsafeProjectiveSearch& tree, + const Eigen::Vector4d& point, + size_t* k_indices, + double* k_sq_dists) { + return tree.nearest_neighbor_search(point, k_indices, k_sq_dists); + } + + static size_t + knn_search(const UnsafeProjectiveSearch& tree, const Eigen::Vector4d& point, size_t k, size_t* k_indices, double* k_sq_dists) { + return tree.knn_search(point, k, k_indices, k_sq_dists); + } +}; + +template +struct Traits> { + static size_t + nearest_neighbor_search(const ProjectiveSearch& tree, const Eigen::Vector4d& point, size_t* k_indices, double* k_sq_dists) { + return tree.nearest_neighbor_search(point, k_indices, k_sq_dists); + } + + static size_t + knn_search(const ProjectiveSearch& tree, const Eigen::Vector4d& point, size_t k, size_t* k_indices, double* k_sq_dists) { + return tree.knn_search(point, k, k_indices, k_sq_dists); + } +}; + +} // namespace traits + +} // namespace small_gicp