Open3D (C++ API)  0.17.0
NanoFlannImpl.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// Copyright (c) 2018-2023 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 auto points_equal = [](const T *const p1, const T *const p2,
122 size_t dimension) {
123 std::vector<T> p1_vec(p1, p1 + dimension);
124 std::vector<T> p2_vec(p2, p2 + dimension);
125 return p1_vec == p2_vec;
126 };
127
128 std::vector<std::vector<TIndex>> neighbors_indices(num_queries);
129 std::vector<std::vector<T>> neighbors_distances(num_queries);
130 std::vector<uint32_t> neighbors_count(num_queries, 0);
131
132 // cast NanoFlannIndexHolder
133 auto holder_ =
134 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
135
136 tbb::parallel_for(
137 tbb::blocked_range<size_t>(0, num_queries),
138 [&](const tbb::blocked_range<size_t> &r) {
139 std::vector<TIndex> result_indices(knn);
140 std::vector<T> result_distances(knn);
141 for (size_t i = r.begin(); i != r.end(); ++i) {
142 size_t num_valid = holder_->index_->knnSearch(
143 &queries[i * dimension], knn, result_indices.data(),
144 result_distances.data());
145
146 int num_neighbors = 0;
147 for (size_t valid_i = 0; valid_i < num_valid; ++valid_i) {
148 TIndex idx = result_indices[valid_i];
149 if (ignore_query_point &&
150 points_equal(&queries[i * dimension],
151 &points[idx * dimension], dimension)) {
152 continue;
153 }
154 neighbors_indices[i].push_back(idx);
155 if (return_distances) {
156 T dist = result_distances[valid_i];
157 neighbors_distances[i].push_back(dist);
158 }
159 ++num_neighbors;
160 }
161 neighbors_count[i] = num_neighbors;
162 }
163 });
164
165 query_neighbors_row_splits[0] = 0;
166 utility::InclusivePrefixSum(neighbors_count.data(),
167 neighbors_count.data() + neighbors_count.size(),
168 query_neighbors_row_splits + 1);
169
170 int64_t num_indices = query_neighbors_row_splits[num_queries];
171
172 TIndex *indices_ptr;
173 output_allocator.AllocIndices(&indices_ptr, num_indices);
174 T *distances_ptr;
175 if (return_distances)
176 output_allocator.AllocDistances(&distances_ptr, num_indices);
177 else
178 output_allocator.AllocDistances(&distances_ptr, 0);
179
180 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
181
182 // fill output index and distance arrays
183 tbb::parallel_for(tbb::blocked_range<size_t>(0, num_queries),
184 [&](const tbb::blocked_range<size_t> &r) {
185 for (size_t i = r.begin(); i != r.end(); ++i) {
186 int64_t start_idx = query_neighbors_row_splits[i];
187 std::copy(neighbors_indices[i].begin(),
188 neighbors_indices[i].end(),
189 &indices_ptr[start_idx]);
190
191 if (return_distances) {
192 std::copy(neighbors_distances[i].begin(),
193 neighbors_distances[i].end(),
194 &distances_ptr[start_idx]);
195 }
196 }
197 });
198}
199
200template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
201void _RadiusSearchCPU(NanoFlannIndexHolderBase *holder,
202 int64_t *query_neighbors_row_splits,
203 size_t num_points,
204 const T *const points,
205 size_t num_queries,
206 const T *const queries,
207 const size_t dimension,
208 const T *const radii,
209 bool ignore_query_point,
210 bool return_distances,
211 bool normalize_distances,
212 bool sort,
213 OUTPUT_ALLOCATOR &output_allocator) {
214 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
215 std::fill(query_neighbors_row_splits,
216 query_neighbors_row_splits + num_queries + 1, 0);
217 TIndex *indices_ptr;
218 output_allocator.AllocIndices(&indices_ptr, 0);
219
220 T *distances_ptr;
221 output_allocator.AllocDistances(&distances_ptr, 0);
222 return;
223 }
224
225 auto points_equal = [](const T *const p1, const T *const p2,
226 size_t dimension) {
227 std::vector<T> p1_vec(p1, p1 + dimension);
228 std::vector<T> p2_vec(p2, p2 + dimension);
229 return p1_vec == p2_vec;
230 };
231
232 std::vector<std::vector<TIndex>> neighbors_indices(num_queries);
233 std::vector<std::vector<T>> neighbors_distances(num_queries);
234 std::vector<uint32_t> neighbors_count(num_queries, 0);
235
236 nanoflann::SearchParameters params;
237 params.sorted = sort;
238
239 auto holder_ =
240 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
241 tbb::parallel_for(
242 tbb::blocked_range<size_t>(0, num_queries),
243 [&](const tbb::blocked_range<size_t> &r) {
244 std::vector<nanoflann::ResultItem<TIndex, T>> search_result;
245 for (size_t i = r.begin(); i != r.end(); ++i) {
246 T radius = radii[i];
247 if (METRIC == L2) {
248 radius = radius * radius;
249 }
250
251 holder_->index_->radiusSearch(&queries[i * dimension],
252 radius, search_result,
253 params);
254
255 int num_neighbors = 0;
256 for (const auto &idx_dist : search_result) {
257 if (ignore_query_point &&
258 points_equal(&queries[i * dimension],
259 &points[idx_dist.first * dimension],
260 dimension)) {
261 continue;
262 }
263 neighbors_indices[i].push_back(idx_dist.first);
264 if (return_distances) {
265 neighbors_distances[i].push_back(idx_dist.second);
266 }
267 ++num_neighbors;
268 }
269 neighbors_count[i] = num_neighbors;
270 }
271 });
272
273 query_neighbors_row_splits[0] = 0;
274 utility::InclusivePrefixSum(neighbors_count.data(),
275 neighbors_count.data() + neighbors_count.size(),
276 query_neighbors_row_splits + 1);
277
278 int64_t num_indices = query_neighbors_row_splits[num_queries];
279
280 TIndex *indices_ptr;
281 output_allocator.AllocIndices(&indices_ptr, num_indices);
282 T *distances_ptr;
283 if (return_distances)
284 output_allocator.AllocDistances(&distances_ptr, num_indices);
285 else
286 output_allocator.AllocDistances(&distances_ptr, 0);
287
288 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
289
290 // fill output index and distance arrays
291 tbb::parallel_for(
292 tbb::blocked_range<size_t>(0, num_queries),
293 [&](const tbb::blocked_range<size_t> &r) {
294 for (size_t i = r.begin(); i != r.end(); ++i) {
295 int64_t start_idx = query_neighbors_row_splits[i];
296 std::copy(neighbors_indices[i].begin(),
297 neighbors_indices[i].end(),
298 &indices_ptr[start_idx]);
299 if (return_distances) {
300 std::transform(neighbors_distances[i].begin(),
301 neighbors_distances[i].end(),
302 &distances_ptr[start_idx], [&](T dist) {
303 T d = dist;
304 if (normalize_distances) {
305 if (METRIC == L2) {
306 d /= (radii[i] * radii[i]);
307 } else {
308 d /= radii[i];
309 }
310 }
311 return d;
312 });
313 }
314 }
315 });
316}
317
318template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
319void _HybridSearchCPU(NanoFlannIndexHolderBase *holder,
320 size_t num_points,
321 const T *const points,
322 size_t num_queries,
323 const T *const queries,
324 const size_t dimension,
325 const T radius,
326 const int max_knn,
327 bool ignore_query_point,
328 bool return_distances,
329 OUTPUT_ALLOCATOR &output_allocator) {
330 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
331 TIndex *indices_ptr, *counts_ptr;
332 output_allocator.AllocIndices(&indices_ptr, 0);
333 output_allocator.AllocCounts(&counts_ptr, 0);
334
335 T *distances_ptr;
336 output_allocator.AllocDistances(&distances_ptr, 0);
337 return;
338 }
339
340 T radius_squared = radius * radius;
341 TIndex *indices_ptr, *counts_ptr;
342 T *distances_ptr;
343
344 size_t num_indices = static_cast<size_t>(max_knn) * num_queries;
345 output_allocator.AllocIndices(&indices_ptr, num_indices);
346 output_allocator.AllocDistances(&distances_ptr, num_indices);
347 output_allocator.AllocCounts(&counts_ptr, num_queries);
348
349 nanoflann::SearchParameters params;
350 params.sorted = true;
351
352 auto holder_ =
353 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
354 tbb::parallel_for(
355 tbb::blocked_range<size_t>(0, num_queries),
356 [&](const tbb::blocked_range<size_t> &r) {
357 std::vector<nanoflann::ResultItem<TIndex, T>> ret_matches;
358 for (size_t i = r.begin(); i != r.end(); ++i) {
359 size_t num_results = holder_->index_->radiusSearch(
360 &queries[i * dimension], radius_squared,
361 ret_matches, params);
362 ret_matches.resize(num_results);
363
364 TIndex count_i = static_cast<TIndex>(num_results);
365 count_i = count_i < max_knn ? count_i : max_knn;
366 counts_ptr[i] = count_i;
367
368 int neighbor_idx = 0;
369 for (auto it = ret_matches.begin();
370 it < ret_matches.end() && neighbor_idx < max_knn;
371 it++, neighbor_idx++) {
372 indices_ptr[i * max_knn + neighbor_idx] = it->first;
373 distances_ptr[i * max_knn + neighbor_idx] = it->second;
374 }
375
376 while (neighbor_idx < max_knn) {
377 indices_ptr[i * max_knn + neighbor_idx] = -1;
378 distances_ptr[i * max_knn + neighbor_idx] = 0;
379 neighbor_idx += 1;
380 }
381 }
382 });
383}
384
385} // namespace
386
401template <class T, class TIndex>
402std::unique_ptr<NanoFlannIndexHolderBase> BuildKdTree(size_t num_points,
403 const T *const points,
404 size_t dimension,
405 const Metric metric) {
406 NanoFlannIndexHolderBase *holder = nullptr;
407#define FN_PARAMETERS num_points, points, dimension, &holder
408
409#define CALL_TEMPLATE(METRIC) \
410 if (METRIC == metric) { \
411 _BuildKdTree<T, TIndex, METRIC>(FN_PARAMETERS); \
412 }
413
414#define CALL_TEMPLATE2 \
415 CALL_TEMPLATE(L1) \
416 CALL_TEMPLATE(L2)
417
419
420#undef CALL_TEMPLATE
421#undef CALL_TEMPLATE2
422
423#undef FN_PARAMETERS
424 return std::unique_ptr<NanoFlannIndexHolderBase>(holder);
425}
426
479template <class T, class TIndex, class OUTPUT_ALLOCATOR>
481 int64_t *query_neighbors_row_splits,
482 size_t num_points,
483 const T *const points,
484 size_t num_queries,
485 const T *const queries,
486 const size_t dimension,
487 int knn,
488 const Metric metric,
489 bool ignore_query_point,
490 bool return_distances,
491 OUTPUT_ALLOCATOR &output_allocator) {
492#define FN_PARAMETERS \
493 holder, query_neighbors_row_splits, num_points, points, num_queries, \
494 queries, dimension, knn, ignore_query_point, return_distances, \
495 output_allocator
496
497#define CALL_TEMPLATE(METRIC) \
498 if (METRIC == metric) { \
499 _KnnSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
500 }
501
502#define CALL_TEMPLATE2 \
503 CALL_TEMPLATE(L1) \
504 CALL_TEMPLATE(L2)
505
507
508#undef CALL_TEMPLATE
509#undef CALL_TEMPLATE2
510
511#undef FN_PARAMETERS
512}
513
573template <class T, class TIndex, class OUTPUT_ALLOCATOR>
575 int64_t *query_neighbors_row_splits,
576 size_t num_points,
577 const T *const points,
578 size_t num_queries,
579 const T *const queries,
580 const size_t dimension,
581 const T *const radii,
582 const Metric metric,
583 bool ignore_query_point,
584 bool return_distances,
585 bool normalize_distances,
586 bool sort,
587 OUTPUT_ALLOCATOR &output_allocator) {
588#define FN_PARAMETERS \
589 holder, query_neighbors_row_splits, num_points, points, num_queries, \
590 queries, dimension, radii, ignore_query_point, return_distances, \
591 normalize_distances, sort, output_allocator
592
593#define CALL_TEMPLATE(METRIC) \
594 if (METRIC == metric) { \
595 _RadiusSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
596 }
597
598#define CALL_TEMPLATE2 \
599 CALL_TEMPLATE(L1) \
600 CALL_TEMPLATE(L2)
601
603
604#undef CALL_TEMPLATE
605#undef CALL_TEMPLATE2
606
607#undef FN_PARAMETERS
608}
609
661template <class T, class TIndex, class OUTPUT_ALLOCATOR>
663 size_t num_points,
664 const T *const points,
665 size_t num_queries,
666 const T *const queries,
667 const size_t dimension,
668 const T radius,
669 const int max_knn,
670 const Metric metric,
671 bool ignore_query_point,
672 bool return_distances,
673 OUTPUT_ALLOCATOR &output_allocator) {
674#define FN_PARAMETERS \
675 holder, num_points, points, num_queries, queries, dimension, radius, \
676 max_knn, ignore_query_point, return_distances, output_allocator
677
678#define CALL_TEMPLATE(METRIC) \
679 if (METRIC == metric) { \
680 _HybridSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
681 }
682
683#define CALL_TEMPLATE2 \
684 CALL_TEMPLATE(L1) \
685 CALL_TEMPLATE(L2)
686
688
689#undef CALL_TEMPLATE
690#undef CALL_TEMPLATE2
691
692#undef FN_PARAMETERS
693}
694
695} // namespace impl
696} // namespace nns
697} // namespace core
698} // namespace open3d
#define CALL_TEMPLATE2
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:402
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:574
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:662
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:480
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:71
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