Open3D (C++ API)  0.16.0
NanoFlannImpl.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// The MIT License (MIT)
5//
6// Copyright (c) 2018-2021 www.open3d.org
7//
8// Permission is hereby granted, free of charge, to any person obtaining a copy
9// of this software and associated documentation files (the "Software"), to deal
10// in the Software without restriction, including without limitation the rights
11// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12// copies of the Software, and to permit persons to whom the Software is
13// furnished to do so, subject to the following conditions:
14//
15// The above copyright notice and this permission notice shall be included in
16// all copies or substantial portions of the Software.
17//
18// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
23// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
24// IN THE SOFTWARE.
25// ----------------------------------------------------------------------------
26#pragma once
27
28#include <tbb/parallel_for.h>
29
30#include <algorithm>
31#include <mutex>
32#include <nanoflann.hpp>
33
34#include "open3d/core/Atomic.h"
38
39namespace open3d {
40namespace core {
41namespace nns {
42
44template <int METRIC, class TReal, class TIndex>
47 struct DataAdaptor {
48 DataAdaptor(size_t dataset_size,
49 int dimension,
50 const TReal *const data_ptr)
51 : dataset_size_(dataset_size),
52 dimension_(dimension),
53 data_ptr_(data_ptr) {}
54
55 inline size_t kdtree_get_point_count() const { return dataset_size_; }
56
57 inline TReal kdtree_get_pt(const size_t idx, const size_t dim) const {
58 return data_ptr_[idx * dimension_ + dim];
59 }
60
61 template <class BBOX>
62 bool kdtree_get_bbox(BBOX &) const {
63 return false;
64 }
65
66 size_t dataset_size_ = 0;
67 int dimension_ = 0;
68 const TReal *const data_ptr_;
69 };
70
72 template <int M, typename fake = void>
74
75 template <typename fake>
76 struct SelectNanoflannAdaptor<L2, fake> {
77 typedef nanoflann::L2_Adaptor<TReal, DataAdaptor, TReal> adaptor_t;
78 };
79
80 template <typename fake>
81 struct SelectNanoflannAdaptor<L1, fake> {
82 typedef nanoflann::L1_Adaptor<TReal, DataAdaptor, TReal> adaptor_t;
83 };
84
86 typedef nanoflann::KDTreeSingleIndexAdaptor<
89 -1,
90 TIndex>
92
93 NanoFlannIndexHolder(size_t dataset_size,
94 int dimension,
95 const TReal *data_ptr) {
96 adaptor_.reset(new DataAdaptor(dataset_size, dimension, data_ptr));
97 index_.reset(new KDTree_t(dimension, *adaptor_.get()));
98 index_->buildIndex();
99 }
100
101 std::unique_ptr<KDTree_t> index_;
102 std::unique_ptr<DataAdaptor> adaptor_;
103};
104namespace impl {
105
106namespace {
107template <class T, class TIndex, int METRIC>
108void _BuildKdTree(size_t num_points,
109 const T *const points,
110 size_t dimension,
111 NanoFlannIndexHolderBase **holder) {
112 *holder = new NanoFlannIndexHolder<METRIC, T, TIndex>(num_points, dimension,
113 points);
114}
115
116template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
117void _KnnSearchCPU(NanoFlannIndexHolderBase *holder,
118 int64_t *query_neighbors_row_splits,
119 size_t num_points,
120 const T *const points,
121 size_t num_queries,
122 const T *const queries,
123 const size_t dimension,
124 int knn,
125 bool ignore_query_point,
126 bool return_distances,
127 OUTPUT_ALLOCATOR &output_allocator) {
128 // return empty indices array if there are no points
129 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
130 std::fill(query_neighbors_row_splits,
131 query_neighbors_row_splits + num_queries + 1, 0);
132 TIndex *indices_ptr;
133 output_allocator.AllocIndices(&indices_ptr, 0);
134
135 T *distances_ptr;
136 output_allocator.AllocDistances(&distances_ptr, 0);
137 return;
138 }
139
140 auto points_equal = [](const T *const p1, const T *const p2,
141 size_t dimension) {
142 std::vector<T> p1_vec(p1, p1 + dimension);
143 std::vector<T> p2_vec(p2, p2 + dimension);
144 return p1_vec == p2_vec;
145 };
146
147 std::vector<std::vector<TIndex>> neighbors_indices(num_queries);
148 std::vector<std::vector<T>> neighbors_distances(num_queries);
149 std::vector<uint32_t> neighbors_count(num_queries, 0);
150
151 // cast NanoFlannIndexHolder
152 auto holder_ =
153 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
154
155 tbb::parallel_for(
156 tbb::blocked_range<size_t>(0, num_queries),
157 [&](const tbb::blocked_range<size_t> &r) {
158 std::vector<TIndex> result_indices(knn);
159 std::vector<T> result_distances(knn);
160 for (size_t i = r.begin(); i != r.end(); ++i) {
161 size_t num_valid = holder_->index_->knnSearch(
162 &queries[i * dimension], knn, result_indices.data(),
163 result_distances.data());
164
165 int num_neighbors = 0;
166 for (size_t valid_i = 0; valid_i < num_valid; ++valid_i) {
167 TIndex idx = result_indices[valid_i];
168 if (ignore_query_point &&
169 points_equal(&queries[i * dimension],
170 &points[idx * dimension], dimension)) {
171 continue;
172 }
173 neighbors_indices[i].push_back(idx);
174 if (return_distances) {
175 T dist = result_distances[valid_i];
176 neighbors_distances[i].push_back(dist);
177 }
178 ++num_neighbors;
179 }
180 neighbors_count[i] = num_neighbors;
181 }
182 });
183
184 query_neighbors_row_splits[0] = 0;
185 utility::InclusivePrefixSum(neighbors_count.data(),
186 neighbors_count.data() + neighbors_count.size(),
187 query_neighbors_row_splits + 1);
188
189 int64_t num_indices = query_neighbors_row_splits[num_queries];
190
191 TIndex *indices_ptr;
192 output_allocator.AllocIndices(&indices_ptr, num_indices);
193 T *distances_ptr;
194 if (return_distances)
195 output_allocator.AllocDistances(&distances_ptr, num_indices);
196 else
197 output_allocator.AllocDistances(&distances_ptr, 0);
198
199 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
200
201 // fill output index and distance arrays
202 tbb::parallel_for(tbb::blocked_range<size_t>(0, num_queries),
203 [&](const tbb::blocked_range<size_t> &r) {
204 for (size_t i = r.begin(); i != r.end(); ++i) {
205 int64_t start_idx = query_neighbors_row_splits[i];
206 std::copy(neighbors_indices[i].begin(),
207 neighbors_indices[i].end(),
208 &indices_ptr[start_idx]);
209
210 if (return_distances) {
211 std::copy(neighbors_distances[i].begin(),
212 neighbors_distances[i].end(),
213 &distances_ptr[start_idx]);
214 }
215 }
216 });
217}
218
219template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
220void _RadiusSearchCPU(NanoFlannIndexHolderBase *holder,
221 int64_t *query_neighbors_row_splits,
222 size_t num_points,
223 const T *const points,
224 size_t num_queries,
225 const T *const queries,
226 const size_t dimension,
227 const T *const radii,
228 bool ignore_query_point,
229 bool return_distances,
230 bool normalize_distances,
231 bool sort,
232 OUTPUT_ALLOCATOR &output_allocator) {
233 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
234 std::fill(query_neighbors_row_splits,
235 query_neighbors_row_splits + num_queries + 1, 0);
236 TIndex *indices_ptr;
237 output_allocator.AllocIndices(&indices_ptr, 0);
238
239 T *distances_ptr;
240 output_allocator.AllocDistances(&distances_ptr, 0);
241 return;
242 }
243
244 auto points_equal = [](const T *const p1, const T *const p2,
245 size_t dimension) {
246 std::vector<T> p1_vec(p1, p1 + dimension);
247 std::vector<T> p2_vec(p2, p2 + dimension);
248 return p1_vec == p2_vec;
249 };
250
251 std::vector<std::vector<TIndex>> neighbors_indices(num_queries);
252 std::vector<std::vector<T>> neighbors_distances(num_queries);
253 std::vector<uint32_t> neighbors_count(num_queries, 0);
254
255 nanoflann::SearchParams params;
256 params.sorted = sort;
257
258 auto holder_ =
259 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
260 tbb::parallel_for(
261 tbb::blocked_range<size_t>(0, num_queries),
262 [&](const tbb::blocked_range<size_t> &r) {
263 std::vector<std::pair<TIndex, T>> search_result;
264 for (size_t i = r.begin(); i != r.end(); ++i) {
265 T radius = radii[i];
266 if (METRIC == L2) {
267 radius = radius * radius;
268 }
269
270 holder_->index_->radiusSearch(&queries[i * dimension],
271 radius, search_result,
272 params);
273
274 int num_neighbors = 0;
275 for (const auto &idx_dist : search_result) {
276 if (ignore_query_point &&
277 points_equal(&queries[i * dimension],
278 &points[idx_dist.first * dimension],
279 dimension)) {
280 continue;
281 }
282 neighbors_indices[i].push_back(idx_dist.first);
283 if (return_distances) {
284 neighbors_distances[i].push_back(idx_dist.second);
285 }
286 ++num_neighbors;
287 }
288 neighbors_count[i] = num_neighbors;
289 }
290 });
291
292 query_neighbors_row_splits[0] = 0;
293 utility::InclusivePrefixSum(neighbors_count.data(),
294 neighbors_count.data() + neighbors_count.size(),
295 query_neighbors_row_splits + 1);
296
297 int64_t num_indices = query_neighbors_row_splits[num_queries];
298
299 TIndex *indices_ptr;
300 output_allocator.AllocIndices(&indices_ptr, num_indices);
301 T *distances_ptr;
302 if (return_distances)
303 output_allocator.AllocDistances(&distances_ptr, num_indices);
304 else
305 output_allocator.AllocDistances(&distances_ptr, 0);
306
307 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
308
309 // fill output index and distance arrays
310 tbb::parallel_for(
311 tbb::blocked_range<size_t>(0, num_queries),
312 [&](const tbb::blocked_range<size_t> &r) {
313 for (size_t i = r.begin(); i != r.end(); ++i) {
314 int64_t start_idx = query_neighbors_row_splits[i];
315 std::copy(neighbors_indices[i].begin(),
316 neighbors_indices[i].end(),
317 &indices_ptr[start_idx]);
318 if (return_distances) {
319 std::transform(neighbors_distances[i].begin(),
320 neighbors_distances[i].end(),
321 &distances_ptr[start_idx], [&](T dist) {
322 T d = dist;
323 if (normalize_distances) {
324 if (METRIC == L2) {
325 d /= (radii[i] * radii[i]);
326 } else {
327 d /= radii[i];
328 }
329 }
330 return d;
331 });
332 }
333 }
334 });
335}
336
337template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
338void _HybridSearchCPU(NanoFlannIndexHolderBase *holder,
339 size_t num_points,
340 const T *const points,
341 size_t num_queries,
342 const T *const queries,
343 const size_t dimension,
344 const T radius,
345 const int max_knn,
346 bool ignore_query_point,
347 bool return_distances,
348 OUTPUT_ALLOCATOR &output_allocator) {
349 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
350 TIndex *indices_ptr, *counts_ptr;
351 output_allocator.AllocIndices(&indices_ptr, 0);
352 output_allocator.AllocCounts(&counts_ptr, 0);
353
354 T *distances_ptr;
355 output_allocator.AllocDistances(&distances_ptr, 0);
356 return;
357 }
358
359 T radius_squared = radius * radius;
360 TIndex *indices_ptr, *counts_ptr;
361 T *distances_ptr;
362
363 size_t num_indices = static_cast<size_t>(max_knn) * num_queries;
364 output_allocator.AllocIndices(&indices_ptr, num_indices);
365 output_allocator.AllocDistances(&distances_ptr, num_indices);
366 output_allocator.AllocCounts(&counts_ptr, num_queries);
367
368 nanoflann::SearchParams params;
369 params.sorted = true;
370
371 auto holder_ =
372 static_cast<NanoFlannIndexHolder<METRIC, T, TIndex> *>(holder);
373 tbb::parallel_for(
374 tbb::blocked_range<size_t>(0, num_queries),
375 [&](const tbb::blocked_range<size_t> &r) {
376 std::vector<std::pair<TIndex, T>> ret_matches;
377 for (size_t i = r.begin(); i != r.end(); ++i) {
378 size_t num_results = holder_->index_->radiusSearch(
379 &queries[i * dimension], radius_squared,
380 ret_matches, params);
381 ret_matches.resize(num_results);
382
383 TIndex count_i = static_cast<TIndex>(num_results);
384 count_i = count_i < max_knn ? count_i : max_knn;
385 counts_ptr[i] = count_i;
386
387 int neighbor_idx = 0;
388 for (auto it = ret_matches.begin();
389 it < ret_matches.end() && neighbor_idx < max_knn;
390 it++, neighbor_idx++) {
391 indices_ptr[i * max_knn + neighbor_idx] = it->first;
392 distances_ptr[i * max_knn + neighbor_idx] = it->second;
393 }
394
395 while (neighbor_idx < max_knn) {
396 indices_ptr[i * max_knn + neighbor_idx] = -1;
397 distances_ptr[i * max_knn + neighbor_idx] = 0;
398 neighbor_idx += 1;
399 }
400 }
401 });
402}
403
404} // namespace
405
420template <class T, class TIndex>
421std::unique_ptr<NanoFlannIndexHolderBase> BuildKdTree(size_t num_points,
422 const T *const points,
423 size_t dimension,
424 const Metric metric) {
425 NanoFlannIndexHolderBase *holder = nullptr;
426#define FN_PARAMETERS num_points, points, dimension, &holder
427
428#define CALL_TEMPLATE(METRIC) \
429 if (METRIC == metric) { \
430 _BuildKdTree<T, TIndex, METRIC>(FN_PARAMETERS); \
431 }
432
433#define CALL_TEMPLATE2 \
434 CALL_TEMPLATE(L1) \
435 CALL_TEMPLATE(L2)
436
438
439#undef CALL_TEMPLATE
440#undef CALL_TEMPLATE2
441
442#undef FN_PARAMETERS
443 return std::unique_ptr<NanoFlannIndexHolderBase>(holder);
444}
445
498template <class T, class TIndex, class OUTPUT_ALLOCATOR>
500 int64_t *query_neighbors_row_splits,
501 size_t num_points,
502 const T *const points,
503 size_t num_queries,
504 const T *const queries,
505 const size_t dimension,
506 int knn,
507 const Metric metric,
508 bool ignore_query_point,
509 bool return_distances,
510 OUTPUT_ALLOCATOR &output_allocator) {
511#define FN_PARAMETERS \
512 holder, query_neighbors_row_splits, num_points, points, num_queries, \
513 queries, dimension, knn, ignore_query_point, return_distances, \
514 output_allocator
515
516#define CALL_TEMPLATE(METRIC) \
517 if (METRIC == metric) { \
518 _KnnSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
519 }
520
521#define CALL_TEMPLATE2 \
522 CALL_TEMPLATE(L1) \
523 CALL_TEMPLATE(L2)
524
526
527#undef CALL_TEMPLATE
528#undef CALL_TEMPLATE2
529
530#undef FN_PARAMETERS
531}
532
592template <class T, class TIndex, class OUTPUT_ALLOCATOR>
594 int64_t *query_neighbors_row_splits,
595 size_t num_points,
596 const T *const points,
597 size_t num_queries,
598 const T *const queries,
599 const size_t dimension,
600 const T *const radii,
601 const Metric metric,
602 bool ignore_query_point,
603 bool return_distances,
604 bool normalize_distances,
605 bool sort,
606 OUTPUT_ALLOCATOR &output_allocator) {
607#define FN_PARAMETERS \
608 holder, query_neighbors_row_splits, num_points, points, num_queries, \
609 queries, dimension, radii, ignore_query_point, return_distances, \
610 normalize_distances, sort, output_allocator
611
612#define CALL_TEMPLATE(METRIC) \
613 if (METRIC == metric) { \
614 _RadiusSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
615 }
616
617#define CALL_TEMPLATE2 \
618 CALL_TEMPLATE(L1) \
619 CALL_TEMPLATE(L2)
620
622
623#undef CALL_TEMPLATE
624#undef CALL_TEMPLATE2
625
626#undef FN_PARAMETERS
627}
628
680template <class T, class TIndex, class OUTPUT_ALLOCATOR>
682 size_t num_points,
683 const T *const points,
684 size_t num_queries,
685 const T *const queries,
686 const size_t dimension,
687 const T radius,
688 const int max_knn,
689 const Metric metric,
690 bool ignore_query_point,
691 bool return_distances,
692 OUTPUT_ALLOCATOR &output_allocator) {
693#define FN_PARAMETERS \
694 holder, num_points, points, num_queries, queries, dimension, radius, \
695 max_knn, ignore_query_point, return_distances, output_allocator
696
697#define CALL_TEMPLATE(METRIC) \
698 if (METRIC == metric) { \
699 _HybridSearchCPU<T, TIndex, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
700 }
701
702#define CALL_TEMPLATE2 \
703 CALL_TEMPLATE(L1) \
704 CALL_TEMPLATE(L2)
705
707
708#undef CALL_TEMPLATE
709#undef CALL_TEMPLATE2
710
711#undef FN_PARAMETERS
712}
713
714} // namespace impl
715} // namespace nns
716} // namespace core
717} // namespace open3d
#define CALL_TEMPLATE2
int points
Definition: FilePCD.cpp:73
std::unique_ptr< NanoFlannIndexHolderBase > BuildKdTree(size_t num_points, const T *const points, size_t dimension, const Metric metric)
Definition: NanoFlannImpl.h:421
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:593
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:681
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:499
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
@ L1
Definition: NeighborSearchCommon.h:38
@ L2
Definition: NeighborSearchCommon.h:38
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:90
Definition: PinholeCameraIntrinsic.cpp:35
This class is the Adaptor for connecting Open3D Tensor and NanoFlann.
Definition: NanoFlannImpl.h:47
const TReal *const data_ptr_
Definition: NanoFlannImpl.h:68
size_t dataset_size_
Definition: NanoFlannImpl.h:66
int dimension_
Definition: NanoFlannImpl.h:67
TReal kdtree_get_pt(const size_t idx, const size_t dim) const
Definition: NanoFlannImpl.h:57
DataAdaptor(size_t dataset_size, int dimension, const TReal *const data_ptr)
Definition: NanoFlannImpl.h:48
bool kdtree_get_bbox(BBOX &) const
Definition: NanoFlannImpl.h:62
size_t kdtree_get_point_count() const
Definition: NanoFlannImpl.h:55
nanoflann::L1_Adaptor< TReal, DataAdaptor, TReal > adaptor_t
Definition: NanoFlannImpl.h:82
nanoflann::L2_Adaptor< TReal, DataAdaptor, TReal > adaptor_t
Definition: NanoFlannImpl.h:77
Adaptor Selector.
Definition: NanoFlannImpl.h:73
Base struct for NanoFlann index holder.
Definition: NeighborSearchCommon.h:72
NanoFlann Index Holder.
Definition: NanoFlannImpl.h:45
std::unique_ptr< KDTree_t > index_
Definition: NanoFlannImpl.h:101
nanoflann::KDTreeSingleIndexAdaptor< typename SelectNanoflannAdaptor< METRIC >::adaptor_t, DataAdaptor, -1, TIndex > KDTree_t
typedef for KDtree.
Definition: NanoFlannImpl.h:91
std::unique_ptr< DataAdaptor > adaptor_
Definition: NanoFlannImpl.h:102
NanoFlannIndexHolder(size_t dataset_size, int dimension, const TReal *data_ptr)
Definition: NanoFlannImpl.h:93