Open3D (C++ API)  0.19.0
Loading...
Searching...
No Matches
NanoFlannImpl.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// Copyright (c) 2018-2024 www.open3d.org
5// SPDX-License-Identifier: MIT
6// ----------------------------------------------------------------------------
7#pragma once
8
9#include <tbb/parallel_for.h>
10
11#include <algorithm>
12#include <mutex>
13#include <nanoflann.hpp>
14
15#include "open3d/core/Atomic.h"
19
20namespace open3d {
21namespace core {
22namespace nns {
23
25template <int METRIC, class TReal, class TIndex>
28 struct DataAdaptor {
29 DataAdaptor(size_t dataset_size,
30 int dimension,
31 const TReal *const data_ptr)
32 : dataset_size_(dataset_size),
33 dimension_(dimension),
34 data_ptr_(data_ptr) {}
35
36 inline size_t kdtree_get_point_count() const { return dataset_size_; }
37
38 inline TReal kdtree_get_pt(const size_t idx, const size_t dim) const {
39 return data_ptr_[idx * dimension_ + dim];
40 }
41
42 template <class BBOX>
43 bool kdtree_get_bbox(BBOX &) const {
44 return false;
45 }
46
47 size_t dataset_size_ = 0;
48 int dimension_ = 0;
49 const TReal *const data_ptr_;
50 };
51
53 template <int M, typename fake = void>
55
56 template <typename fake>
57 struct SelectNanoflannAdaptor<L2, fake> {
58 typedef nanoflann::L2_Adaptor<TReal, DataAdaptor, TReal> adaptor_t;
59 };
60
61 template <typename fake>
62 struct SelectNanoflannAdaptor<L1, fake> {
63 typedef nanoflann::L1_Adaptor<TReal, DataAdaptor, TReal> adaptor_t;
64 };
65
67 typedef nanoflann::KDTreeSingleIndexAdaptor<
70 -1,
71 TIndex>
73
74 NanoFlannIndexHolder(size_t dataset_size,
75 int dimension,
76 const TReal *data_ptr) {
77 adaptor_.reset(new DataAdaptor(dataset_size, dimension, data_ptr));
78 index_.reset(new KDTree_t(dimension, *adaptor_.get()));
79 index_->buildIndex();
80 }
81
82 std::unique_ptr<KDTree_t> index_;
83 std::unique_ptr<DataAdaptor> adaptor_;
84};
85namespace impl {
86
87namespace {
88template <class T, class TIndex, int METRIC>
89void _BuildKdTree(size_t num_points,
90 const T *const points,
91 size_t dimension,
92 NanoFlannIndexHolderBase **holder) {
93 *holder = new NanoFlannIndexHolder<METRIC, T, TIndex>(num_points, dimension,
94 points);
95}
96
97template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
98void _KnnSearchCPU(NanoFlannIndexHolderBase *holder,
99 int64_t *query_neighbors_row_splits,
100 size_t num_points,
101 const T *const points,
102 size_t num_queries,
103 const T *const queries,
104 const size_t dimension,
105 int knn,
106 bool ignore_query_point,
107 bool return_distances,
108 OUTPUT_ALLOCATOR &output_allocator) {
109 // return empty indices array if there are no points
110 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
111 std::fill(query_neighbors_row_splits,
112 query_neighbors_row_splits + num_queries + 1, 0);
113 TIndex *indices_ptr;
114 output_allocator.AllocIndices(&indices_ptr, 0);
115
116 T *distances_ptr;
117 output_allocator.AllocDistances(&distances_ptr, 0);
118 return;
119 }
120
121 std::vector<std::vector<TIndex>> neighbors_indices(num_queries);
122 std::vector<std::vector<T>> neighbors_distances(num_queries);
123 std::vector<uint32_t> neighbors_count(num_queries, 0);
124
125 // cast NanoFlannIndexHolder
126 auto holder_ =
127 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
128
129 tbb::parallel_for(
130 tbb::blocked_range<size_t>(0, num_queries),
131 [&](const tbb::blocked_range<size_t> &r) {
132 std::vector<TIndex> result_indices(knn);
133 std::vector<T> result_distances(knn);
134 for (size_t i = r.begin(); i != r.end(); ++i) {
135 size_t num_valid = holder_->index_->knnSearch(
136 &queries[i * dimension], knn, result_indices.data(),
137 result_distances.data());
138
139 int num_neighbors = 0;
140 for (size_t valid_i = 0; valid_i < num_valid; ++valid_i) {
141 TIndex idx = result_indices[valid_i];
142 if (ignore_query_point &&
143 std::equal(&queries[i * dimension],
144 &queries[i * dimension] + dimension,
145 &points[idx * dimension])) {
146 continue;
147 }
148 neighbors_indices[i].push_back(idx);
149 if (return_distances) {
150 T dist = result_distances[valid_i];
151 neighbors_distances[i].push_back(dist);
152 }
153 ++num_neighbors;
154 }
155 neighbors_count[i] = num_neighbors;
156 }
157 });
158
159 query_neighbors_row_splits[0] = 0;
160 utility::InclusivePrefixSum(neighbors_count.data(),
161 neighbors_count.data() + neighbors_count.size(),
162 query_neighbors_row_splits + 1);
163
164 int64_t num_indices = query_neighbors_row_splits[num_queries];
165
166 TIndex *indices_ptr;
167 output_allocator.AllocIndices(&indices_ptr, num_indices);
168 T *distances_ptr;
169 if (return_distances)
170 output_allocator.AllocDistances(&distances_ptr, num_indices);
171 else
172 output_allocator.AllocDistances(&distances_ptr, 0);
173
174 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
175
176 // fill output index and distance arrays
177 tbb::parallel_for(tbb::blocked_range<size_t>(0, num_queries),
178 [&](const tbb::blocked_range<size_t> &r) {
179 for (size_t i = r.begin(); i != r.end(); ++i) {
180 int64_t start_idx = query_neighbors_row_splits[i];
181 std::copy(neighbors_indices[i].begin(),
182 neighbors_indices[i].end(),
183 &indices_ptr[start_idx]);
184
185 if (return_distances) {
186 std::copy(neighbors_distances[i].begin(),
187 neighbors_distances[i].end(),
188 &distances_ptr[start_idx]);
189 }
190 }
191 });
192}
193
194template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
195void _RadiusSearchCPU(NanoFlannIndexHolderBase *holder,
196 int64_t *query_neighbors_row_splits,
197 size_t num_points,
198 const T *const points,
199 size_t num_queries,
200 const T *const queries,
201 const size_t dimension,
202 const T *const radii,
203 bool ignore_query_point,
204 bool return_distances,
205 bool normalize_distances,
206 bool sort,
207 OUTPUT_ALLOCATOR &output_allocator) {
208 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
209 std::fill(query_neighbors_row_splits,
210 query_neighbors_row_splits + num_queries + 1, 0);
211 TIndex *indices_ptr;
212 output_allocator.AllocIndices(&indices_ptr, 0);
213
214 T *distances_ptr;
215 output_allocator.AllocDistances(&distances_ptr, 0);
216 return;
217 }
218
219 std::vector<std::vector<TIndex>> neighbors_indices(num_queries);
220 std::vector<std::vector<T>> neighbors_distances(num_queries);
221 std::vector<uint32_t> neighbors_count(num_queries, 0);
222
223 nanoflann::SearchParameters params;
224 params.sorted = sort;
225
226 auto holder_ =
227 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
228 tbb::parallel_for(
229 tbb::blocked_range<size_t>(0, num_queries),
230 [&](const tbb::blocked_range<size_t> &r) {
231 std::vector<nanoflann::ResultItem<TIndex, T>> search_result;
232 for (size_t i = r.begin(); i != r.end(); ++i) {
233 T radius = radii[i];
234 if (METRIC == L2) {
235 radius = radius * radius;
236 }
237
238 holder_->index_->radiusSearch(&queries[i * dimension],
239 radius, search_result,
240 params);
241
242 int num_neighbors = 0;
243 for (const auto &idx_dist : search_result) {
244 if (ignore_query_point &&
245 std::equal(&queries[i * dimension],
246 &queries[i * dimension] + dimension,
247 &points[idx_dist.first * dimension])) {
248 continue;
249 }
250 neighbors_indices[i].push_back(idx_dist.first);
251 if (return_distances) {
252 neighbors_distances[i].push_back(idx_dist.second);
253 }
254 ++num_neighbors;
255 }
256 neighbors_count[i] = num_neighbors;
257 }
258 });
259
260 query_neighbors_row_splits[0] = 0;
261 utility::InclusivePrefixSum(neighbors_count.data(),
262 neighbors_count.data() + neighbors_count.size(),
263 query_neighbors_row_splits + 1);
264
265 int64_t num_indices = query_neighbors_row_splits[num_queries];
266
267 TIndex *indices_ptr;
268 output_allocator.AllocIndices(&indices_ptr, num_indices);
269 T *distances_ptr;
270 if (return_distances)
271 output_allocator.AllocDistances(&distances_ptr, num_indices);
272 else
273 output_allocator.AllocDistances(&distances_ptr, 0);
274
275 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
276
277 // fill output index and distance arrays
278 tbb::parallel_for(
279 tbb::blocked_range<size_t>(0, num_queries),
280 [&](const tbb::blocked_range<size_t> &r) {
281 for (size_t i = r.begin(); i != r.end(); ++i) {
282 int64_t start_idx = query_neighbors_row_splits[i];
283 std::copy(neighbors_indices[i].begin(),
284 neighbors_indices[i].end(),
285 &indices_ptr[start_idx]);
286 if (return_distances) {
287 std::transform(neighbors_distances[i].begin(),
288 neighbors_distances[i].end(),
289 &distances_ptr[start_idx], [&](T dist) {
290 T d = dist;
291 if (normalize_distances) {
292 if (METRIC == L2) {
293 d /= (radii[i] * radii[i]);
294 } else {
295 d /= radii[i];
296 }
297 }
298 return d;
299 });
300 }
301 }
302 });
303}
304
305template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
306void _HybridSearchCPU(NanoFlannIndexHolderBase *holder,
307 size_t num_points,
308 const T *const points,
309 size_t num_queries,
310 const T *const queries,
311 const size_t dimension,
312 const T radius,
313 const int max_knn,
314 bool ignore_query_point,
315 bool return_distances,
316 OUTPUT_ALLOCATOR &output_allocator) {
317 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
318 TIndex *indices_ptr, *counts_ptr;
319 output_allocator.AllocIndices(&indices_ptr, 0);
320 output_allocator.AllocCounts(&counts_ptr, 0);
321
322 T *distances_ptr;
323 output_allocator.AllocDistances(&distances_ptr, 0);
324 return;
325 }
326
327 T radius_squared = radius * radius;
328 TIndex *indices_ptr, *counts_ptr;
329 T *distances_ptr;
330
331 size_t num_indices = static_cast<size_t>(max_knn) * num_queries;
332 output_allocator.AllocIndices(&indices_ptr, num_indices);
333 output_allocator.AllocDistances(&distances_ptr, num_indices);
334 output_allocator.AllocCounts(&counts_ptr, num_queries);
335
336 nanoflann::SearchParameters params;
337 params.sorted = true;
338
339 auto holder_ =
340 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
341 tbb::parallel_for(
342 tbb::blocked_range<size_t>(0, num_queries),
343 [&](const tbb::blocked_range<size_t> &r) {
344 std::vector<nanoflann::ResultItem<TIndex, T>> ret_matches;
345 for (size_t i = r.begin(); i != r.end(); ++i) {
346 size_t num_results = holder_->index_->radiusSearch(
347 &queries[i * dimension], radius_squared,
348 ret_matches, params);
349 ret_matches.resize(num_results);
350
351 TIndex count_i = static_cast<TIndex>(num_results);
352 count_i = count_i < max_knn ? count_i : max_knn;
353 counts_ptr[i] = count_i;
354
355 int neighbor_idx = 0;
356 for (auto it = ret_matches.begin();
357 it < ret_matches.end() && neighbor_idx < max_knn;
358 it++, neighbor_idx++) {
359 indices_ptr[i * max_knn + neighbor_idx] = it->first;
360 distances_ptr[i * max_knn + neighbor_idx] = it->second;
361 }
362
363 while (neighbor_idx < max_knn) {
364 indices_ptr[i * max_knn + neighbor_idx] = -1;
365 distances_ptr[i * max_knn + neighbor_idx] = 0;
366 neighbor_idx += 1;
367 }
368 }
369 });
370}
371
372} // namespace
373
388template <class T, class TIndex>
389std::unique_ptr<NanoFlannIndexHolderBase> BuildKdTree(size_t num_points,
390 const T *const points,
391 size_t dimension,
392 const Metric metric) {
393 NanoFlannIndexHolderBase *holder = nullptr;
394#define FN_PARAMETERS num_points, points, dimension, &holder
395
396#define CALL_TEMPLATE(METRIC) \
397 if (METRIC == metric) { \
398 _BuildKdTree<T, TIndex, METRIC>(FN_PARAMETERS); \
399 }
400
401#define CALL_TEMPLATE2 \
402 CALL_TEMPLATE(L1) \
403 CALL_TEMPLATE(L2)
404
406
407#undef CALL_TEMPLATE
408#undef CALL_TEMPLATE2
409
410#undef FN_PARAMETERS
411 return std::unique_ptr<NanoFlannIndexHolderBase>(holder);
412}
413
466template <class T, class TIndex, class OUTPUT_ALLOCATOR>
468 int64_t *query_neighbors_row_splits,
469 size_t num_points,
470 const T *const points,
471 size_t num_queries,
472 const T *const queries,
473 const size_t dimension,
474 int knn,
475 const Metric metric,
476 bool ignore_query_point,
477 bool return_distances,
478 OUTPUT_ALLOCATOR &output_allocator) {
479#define FN_PARAMETERS \
480 holder, query_neighbors_row_splits, num_points, points, num_queries, \
481 queries, dimension, knn, ignore_query_point, return_distances, \
482 output_allocator
483
484#define CALL_TEMPLATE(METRIC) \
485 if (METRIC == metric) { \
486 _KnnSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
487 }
488
489#define CALL_TEMPLATE2 \
490 CALL_TEMPLATE(L1) \
491 CALL_TEMPLATE(L2)
492
494
495#undef CALL_TEMPLATE
496#undef CALL_TEMPLATE2
497
498#undef FN_PARAMETERS
499}
500
560template <class T, class TIndex, class OUTPUT_ALLOCATOR>
562 int64_t *query_neighbors_row_splits,
563 size_t num_points,
564 const T *const points,
565 size_t num_queries,
566 const T *const queries,
567 const size_t dimension,
568 const T *const radii,
569 const Metric metric,
570 bool ignore_query_point,
571 bool return_distances,
572 bool normalize_distances,
573 bool sort,
574 OUTPUT_ALLOCATOR &output_allocator) {
575#define FN_PARAMETERS \
576 holder, query_neighbors_row_splits, num_points, points, num_queries, \
577 queries, dimension, radii, ignore_query_point, return_distances, \
578 normalize_distances, sort, output_allocator
579
580#define CALL_TEMPLATE(METRIC) \
581 if (METRIC == metric) { \
582 _RadiusSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
583 }
584
585#define CALL_TEMPLATE2 \
586 CALL_TEMPLATE(L1) \
587 CALL_TEMPLATE(L2)
588
590
591#undef CALL_TEMPLATE
592#undef CALL_TEMPLATE2
593
594#undef FN_PARAMETERS
595}
596
648template <class T, class TIndex, class OUTPUT_ALLOCATOR>
650 size_t num_points,
651 const T *const points,
652 size_t num_queries,
653 const T *const queries,
654 const size_t dimension,
655 const T radius,
656 const int max_knn,
657 const Metric metric,
658 bool ignore_query_point,
659 bool return_distances,
660 OUTPUT_ALLOCATOR &output_allocator) {
661#define FN_PARAMETERS \
662 holder, num_points, points, num_queries, queries, dimension, radius, \
663 max_knn, ignore_query_point, return_distances, output_allocator
664
665#define CALL_TEMPLATE(METRIC) \
666 if (METRIC == metric) { \
667 _HybridSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
668 }
669
670#define CALL_TEMPLATE2 \
671 CALL_TEMPLATE(L1) \
672 CALL_TEMPLATE(L2)
673
675
676#undef CALL_TEMPLATE
677#undef CALL_TEMPLATE2
678
679#undef FN_PARAMETERS
680}
681
682} // namespace impl
683} // namespace nns
684} // namespace core
685} // namespace open3d
#define CALL_TEMPLATE2(METRIC)
int points
Definition FilePCD.cpp:54
std::unique_ptr< NanoFlannIndexHolderBase > BuildKdTree(size_t num_points, const T *const points, size_t dimension, const Metric metric)
Definition NanoFlannImpl.h:389
void RadiusSearchCPU(NanoFlannIndexHolderBase *holder, int64_t *query_neighbors_row_splits, size_t num_points, const T *const points, size_t num_queries, const T *const queries, const size_t dimension, const T *const radii, const Metric metric, bool ignore_query_point, bool return_distances, bool normalize_distances, bool sort, OUTPUT_ALLOCATOR &output_allocator)
Definition NanoFlannImpl.h:561
void HybridSearchCPU(NanoFlannIndexHolderBase *holder, size_t num_points, const T *const points, size_t num_queries, const T *const queries, const size_t dimension, const T radius, const int max_knn, const Metric metric, bool ignore_query_point, bool return_distances, OUTPUT_ALLOCATOR &output_allocator)
Definition NanoFlannImpl.h:649
void KnnSearchCPU(NanoFlannIndexHolderBase *holder, int64_t *query_neighbors_row_splits, size_t num_points, const T *const points, size_t num_queries, const T *const queries, const size_t dimension, int knn, const Metric metric, bool ignore_query_point, bool return_distances, OUTPUT_ALLOCATOR &output_allocator)
Definition NanoFlannImpl.h:467
Metric
Supported metrics.
Definition NeighborSearchCommon.h:19
@ L1
Definition NeighborSearchCommon.h:19
@ L2
Definition NeighborSearchCommon.h:19
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition ParallelScan.h:43
Definition PinholeCameraIntrinsic.cpp:16
This class is the Adaptor for connecting Open3D Tensor and NanoFlann.
Definition NanoFlannImpl.h:28
const TReal *const data_ptr_
Definition NanoFlannImpl.h:49
size_t dataset_size_
Definition NanoFlannImpl.h:47
int dimension_
Definition NanoFlannImpl.h:48
TReal kdtree_get_pt(const size_t idx, const size_t dim) const
Definition NanoFlannImpl.h:38
DataAdaptor(size_t dataset_size, int dimension, const TReal *const data_ptr)
Definition NanoFlannImpl.h:29
bool kdtree_get_bbox(BBOX &) const
Definition NanoFlannImpl.h:43
size_t kdtree_get_point_count() const
Definition NanoFlannImpl.h:36
nanoflann::L1_Adaptor< TReal, DataAdaptor, TReal > adaptor_t
Definition NanoFlannImpl.h:63
nanoflann::L2_Adaptor< TReal, DataAdaptor, TReal > adaptor_t
Definition NanoFlannImpl.h:58
Adaptor Selector.
Definition NanoFlannImpl.h:54
Base struct for NanoFlann index holder.
Definition NeighborSearchCommon.h:53
NanoFlann Index Holder.
Definition NanoFlannImpl.h:26
std::unique_ptr< KDTree_t > index_
Definition NanoFlannImpl.h:82
nanoflann::KDTreeSingleIndexAdaptor< typename SelectNanoflannAdaptor< METRIC >::adaptor_t, DataAdaptor, -1, TIndex > KDTree_t
typedef for KDtree.
Definition NanoFlannImpl.h:72
std::unique_ptr< DataAdaptor > adaptor_
Definition NanoFlannImpl.h:83
NanoFlannIndexHolder(size_t dataset_size, int dimension, const TReal *data_ptr)
Definition NanoFlannImpl.h:74