Skip to content

Commit 8e66581

Browse files
koide3koide3
authored
projective search (#92)
1 parent 4e8e745 commit 8e66581

File tree

1 file changed

+219
-0
lines changed

1 file changed

+219
-0
lines changed

‎include/small_gicp/ann/projective_search.hpp

+219
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
// SPDX-FileCopyrightText: Copyright 2024 Kenji Koide
2+
// SPDX-License-Identifier: MIT
3+
#pragma once
4+
5+
#include <cmath>
6+
#include <Eigen/Core>
7+
#include <small_gicp/points/traits.hpp>
8+
#include <small_gicp/ann/knn_result.hpp>
9+
10+
namespace small_gicp {
11+
12+
/// @brief Equirectangular projection.
13+
struct EquirectangularProjection {
14+
public:
15+
/// @brief Project the point into the normalized image coordinates. (u, v) in ([0, 1], [0, 1])
16+
Eigen::Vector2d operator()(const Eigen::Vector3d& pt_3d) const {
17+
if (pt_3d.squaredNorm() < 1e-3) {
18+
return Eigen::Vector2d(0.5, 0.5);
19+
}
20+
21+
const Eigen::Vector3d bearing = pt_3d.normalized();
22+
const double lat = -std::asin(bearing[1]);
23+
const double lon = std::atan2(bearing[0], bearing[2]);
24+
25+
return Eigen::Vector2d(lon / (2.0 * M_PI) + 0.5, lat / M_PI + 0.5);
26+
};
27+
};
28+
29+
/// @brief Border clamp mode. Points out of the border are discarded.
30+
struct BorderClamp {
31+
public:
32+
int operator()(int x, int width) const { return x; }
33+
};
34+
35+
/// @brief Border repeat mode. Points out of the border are wrapped around.
36+
struct BorderRepeat {
37+
public:
38+
int operator()(int x, int width) const { return x < 0 ? x + width : (x >= width ? x - width : x); }
39+
};
40+
41+
/// @brief "Unsafe" projective search. This class does not hold the ownership of the target point cloud.
42+
template <typename PointCloud, typename Projection = EquirectangularProjection, typename BorderModeH = BorderRepeat, typename BorderModeV = BorderClamp>
43+
struct UnsafeProjectiveSearch {
44+
public:
45+
/// @brief Constructor.
46+
/// @param width Index map width
47+
/// @param height Index map height
48+
/// @param points Target point cloud
49+
UnsafeProjectiveSearch(int width, int height, const PointCloud& points) : points(points), index_map(height, width), search_window_h(10), search_window_v(5) {
50+
index_map.setConstant(invalid_index);
51+
52+
Projection project;
53+
for (size_t i = 0; i < traits::size(points); ++i) {
54+
const Eigen::Vector4d pt = traits::point(points, i);
55+
const Eigen::Vector2d uv = project(pt.head<3>());
56+
const int u = uv[0] * index_map.cols();
57+
const int v = uv[1] * index_map.rows();
58+
59+
if (u < 0 || u >= index_map.cols() || v < 0 || v >= index_map.rows()) {
60+
continue;
61+
}
62+
index_map(v, u) = i;
63+
}
64+
}
65+
66+
/// @brief Find the nearest neighbor.
67+
/// @param query Query point
68+
/// @param k_indices Index of the nearest neighbor (uninitialized if not found)
69+
/// @param k_sq_dists Squared distance to the nearest neighbor (uninitialized if not found)
70+
/// @param setting KNN search setting
71+
/// @return Number of found neighbors (0 or 1)
72+
size_t nearest_neighbor_search(const Eigen::Vector4d& query, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const {
73+
return knn_search<1>(query, k_indices, k_sq_dists, setting);
74+
}
75+
76+
/// @brief Find k-nearest neighbors. This method uses dynamic memory allocation.
77+
/// @param query Query point
78+
/// @param k Number of neighbors
79+
/// @param k_indices Indices of neighbors
80+
/// @param k_sq_dists Squared distances to neighbors (sorted in ascending order)
81+
/// @param setting KNN search setting
82+
/// @return Number of found neighbors
83+
size_t knn_search(const Eigen::Vector4d& query, int k, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const {
84+
KnnResult<-1> result(k_indices, k_sq_dists, k);
85+
knn_search(query, result, setting);
86+
return result.num_found();
87+
}
88+
89+
/// @brief Find k-nearest neighbors. This method uses fixed and static memory allocation. Might be faster for small k.
90+
/// @param query Query point
91+
/// @param k_indices Indices of neighbors
92+
/// @param k_sq_dists Squared distances to neighbors (sorted in ascending order)
93+
/// @param setting KNN search setting
94+
/// @return Number of found neighbors
95+
template <int N>
96+
size_t knn_search(const Eigen::Vector4d& query, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const {
97+
KnnResult<N> result(k_indices, k_sq_dists);
98+
knn_search(query, result, setting);
99+
return result.num_found();
100+
}
101+
102+
private:
103+
template <typename Result>
104+
void knn_search(const Eigen::Vector4d& query, Result& result, const KnnSetting& setting) const {
105+
BorderModeH border_h;
106+
BorderModeV border_v;
107+
108+
Projection project;
109+
const Eigen::Vector2d uv = project(query.head<3>());
110+
const int u = uv[0] * index_map.cols();
111+
const int v = uv[1] * index_map.rows();
112+
113+
for (int dv = -search_window_v; dv <= search_window_v; dv++) {
114+
const int v_clamped = border_v(v + dv, index_map.rows());
115+
if (v_clamped < 0 || v_clamped >= index_map.rows()) {
116+
continue;
117+
}
118+
119+
for (int du = -search_window_h; du <= search_window_h; du++) {
120+
const int u_clamped = border_h(u + du, index_map.cols());
121+
if (u_clamped < 0 || u_clamped >= index_map.cols()) {
122+
continue;
123+
}
124+
125+
const auto index = index_map(v_clamped, u_clamped);
126+
if (index == invalid_index) {
127+
continue;
128+
}
129+
130+
const double sq_dist = (traits::point(points, index) - query).squaredNorm();
131+
result.push(index, sq_dist);
132+
133+
if (setting.fulfilled(result)) {
134+
return;
135+
}
136+
}
137+
}
138+
}
139+
140+
public:
141+
static constexpr std::uint32_t invalid_index = std::numeric_limits<std::uint32_t>::max();
142+
143+
const PointCloud& points;
144+
Eigen::Matrix<std::uint32_t, -1, -1> index_map;
145+
146+
int search_window_h;
147+
int search_window_v;
148+
};
149+
150+
/// @brief "Safe" projective search. This class holds the ownership of the target point cloud.
151+
template <typename PointCloud, typename Projection = EquirectangularProjection, typename BorderModeH = BorderRepeat, typename BorderModeV = BorderClamp>
152+
struct ProjectiveSearch {
153+
public:
154+
using Ptr = std::shared_ptr<ProjectiveSearch<PointCloud, Projection>>;
155+
using ConstPtr = std::shared_ptr<const ProjectiveSearch<PointCloud, Projection>>;
156+
157+
explicit ProjectiveSearch(int width, int height, std::shared_ptr<const PointCloud> points) : points(points), search(width, height, *points) {}
158+
159+
/// @brief Find k-nearest neighbors. This method uses dynamic memory allocation.
160+
/// @param query Query point
161+
/// @param k Number of neighbors
162+
/// @param k_indices Indices of neighbors
163+
/// @param k_sq_dists Squared distances to neighbors (sorted in ascending order)
164+
/// @param setting KNN search setting
165+
/// @return Number of found neighbors
166+
size_t nearest_neighbor_search(const Eigen::Vector4d& query, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const {
167+
return search.nearest_neighbor_search(query, k_indices, k_sq_dists, setting);
168+
}
169+
170+
/// @brief Find k-nearest neighbors. This method uses dynamic memory allocation.
171+
/// @param query Query point
172+
/// @param k Number of neighbors
173+
/// @param k_indices Indices of neighbors
174+
/// @param k_sq_dists Squared distances to neighbors (sorted in ascending order)
175+
/// @param setting KNN search setting
176+
/// @return Number of found neighbors
177+
size_t knn_search(const Eigen::Vector4d& query, size_t k, size_t* k_indices, double* k_sq_dists, const KnnSetting& setting = KnnSetting()) const {
178+
return search.knn_search(query, k, k_indices, k_sq_dists, setting);
179+
}
180+
181+
public:
182+
const std::shared_ptr<const PointCloud> points; ///< Points
183+
const UnsafeProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV> search; ///< Search
184+
};
185+
186+
namespace traits {
187+
188+
template <typename PointCloud, typename Projection, typename BorderModeH, typename BorderModeV>
189+
struct Traits<UnsafeProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV>> {
190+
static size_t nearest_neighbor_search(
191+
const UnsafeProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV>& tree,
192+
const Eigen::Vector4d& point,
193+
size_t* k_indices,
194+
double* k_sq_dists) {
195+
return tree.nearest_neighbor_search(point, k_indices, k_sq_dists);
196+
}
197+
198+
static size_t
199+
knn_search(const UnsafeProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV>& tree, const Eigen::Vector4d& point, size_t k, size_t* k_indices, double* k_sq_dists) {
200+
return tree.knn_search(point, k, k_indices, k_sq_dists);
201+
}
202+
};
203+
204+
template <typename PointCloud, typename Projection, typename BorderModeH, typename BorderModeV>
205+
struct Traits<ProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV>> {
206+
static size_t
207+
nearest_neighbor_search(const ProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV>& tree, const Eigen::Vector4d& point, size_t* k_indices, double* k_sq_dists) {
208+
return tree.nearest_neighbor_search(point, k_indices, k_sq_dists);
209+
}
210+
211+
static size_t
212+
knn_search(const ProjectiveSearch<PointCloud, Projection, BorderModeH, BorderModeV>& tree, const Eigen::Vector4d& point, size_t k, size_t* k_indices, double* k_sq_dists) {
213+
return tree.knn_search(point, k, k_indices, k_sq_dists);
214+
}
215+
};
216+
217+
} // namespace traits
218+
219+
} // namespace small_gicp

0 commit comments

Comments
 (0)