Open3D (C++ API)  0.15.1
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
44
46template <int METRIC, class TReal, class TIndex>
49 struct DataAdaptor {
50 DataAdaptor(size_t dataset_size,
51 int dimension,
52 const TReal *const data_ptr)
53 : dataset_size_(dataset_size),
54 dimension_(dimension),
55 data_ptr_(data_ptr) {}
56
57 inline size_t kdtree_get_point_count() const { return dataset_size_; }
58
59 inline TReal kdtree_get_pt(const size_t idx, const size_t dim) const {
60 return data_ptr_[idx * dimension_ + dim];
61 }
62
63 template <class BBOX>
64 bool kdtree_get_bbox(BBOX &) const {
65 return false;
66 }
67
68 size_t dataset_size_ = 0;
69 int dimension_ = 0;
70 const TReal *const data_ptr_;
71 };
72
74 template <int M, typename fake = void>
76
77 template <typename fake>
78 struct SelectNanoflannAdaptor<L2, fake> {
79 typedef nanoflann::L2_Adaptor<TReal, DataAdaptor, TReal> adaptor_t;
80 };
81
82 template <typename fake>
83 struct SelectNanoflannAdaptor<L1, fake> {
84 typedef nanoflann::L1_Adaptor<TReal, DataAdaptor, TReal> adaptor_t;
85 };
86
88 typedef nanoflann::KDTreeSingleIndexAdaptor<
91 -1,
92 TIndex>
94
95 NanoFlannIndexHolder(size_t dataset_size,
96 int dimension,
97 const TReal *data_ptr) {
98 adaptor_.reset(new DataAdaptor(dataset_size, dimension, data_ptr));
99 index_.reset(new KDTree_t(dimension, *adaptor_.get()));
100 index_->buildIndex();
101 }
102
103 std::unique_ptr<KDTree_t> index_;
104 std::unique_ptr<DataAdaptor> adaptor_;
105};
106namespace impl {
107
108namespace {
109template <class T, int METRIC>
110void _BuildKdTree(size_t num_points,
111 const T *const points,
112 size_t dimension,
113 NanoFlannIndexHolderBase **holder) {
114 *holder = new NanoFlannIndexHolder<METRIC, T, index_t>(num_points,
115 dimension, points);
116}
117
118template <class T, class OUTPUT_ALLOCATOR, int METRIC>
119void _KnnSearchCPU(NanoFlannIndexHolderBase *holder,
120 int64_t *query_neighbors_row_splits,
121 size_t num_points,
122 const T *const points,
123 size_t num_queries,
124 const T *const queries,
125 const size_t dimension,
126 int knn,
127 bool ignore_query_point,
128 bool return_distances,
129 OUTPUT_ALLOCATOR &output_allocator) {
130 // return empty indices array if there are no points
131 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
132 std::fill(query_neighbors_row_splits,
133 query_neighbors_row_splits + num_queries + 1, 0);
134 index_t *indices_ptr;
135 output_allocator.AllocIndices(&indices_ptr, 0);
136
137 T *distances_ptr;
138 output_allocator.AllocDistances(&distances_ptr, 0);
139 return;
140 }
141
142 auto points_equal = [](const T *const p1, const T *const p2,
143 size_t dimension) {
144 std::vector<T> p1_vec(p1, p1 + dimension);
145 std::vector<T> p2_vec(p2, p2 + dimension);
146 return p1_vec == p2_vec;
147 };
148
149 std::vector<std::vector<index_t>> neighbors_indices(num_queries);
150 std::vector<std::vector<T>> neighbors_distances(num_queries);
151 std::vector<uint32_t> neighbors_count(num_queries, 0);
152
153 // cast NanoFlannIndexHolder
154 auto holder_ =
155 static_cast<NanoFlannIndexHolder<METRIC, T, index_t> *>(holder);
156
157 tbb::parallel_for(
158 tbb::blocked_range<size_t>(0, num_queries),
159 [&](const tbb::blocked_range<size_t> &r) {
160 std::vector<index_t> result_indices(knn);
161 std::vector<T> result_distances(knn);
162 for (size_t i = r.begin(); i != r.end(); ++i) {
163 size_t num_valid = holder_->index_->knnSearch(
164 &queries[i * dimension], knn, result_indices.data(),
165 result_distances.data());
166
167 int num_neighbors = 0;
168 for (size_t valid_i = 0; valid_i < num_valid; ++valid_i) {
169 index_t idx = result_indices[valid_i];
170 if (ignore_query_point &&
171 points_equal(&queries[i * dimension],
172 &points[idx * dimension], dimension)) {
173 continue;
174 }
175 neighbors_indices[i].push_back(idx);
176 if (return_distances) {
177 T dist = result_distances[valid_i];
178 neighbors_distances[i].push_back(dist);
179 }
180 ++num_neighbors;
181 }
182 neighbors_count[i] = num_neighbors;
183 }
184 });
185
186 query_neighbors_row_splits[0] = 0;
187 utility::InclusivePrefixSum(neighbors_count.data(),
188 neighbors_count.data() + neighbors_count.size(),
189 query_neighbors_row_splits + 1);
190
191 int64_t num_indices = query_neighbors_row_splits[num_queries];
192
193 index_t *indices_ptr;
194 output_allocator.AllocIndices(&indices_ptr, num_indices);
195 T *distances_ptr;
196 if (return_distances)
197 output_allocator.AllocDistances(&distances_ptr, num_indices);
198 else
199 output_allocator.AllocDistances(&distances_ptr, 0);
200
201 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
202
203 // fill output index and distance arrays
204 tbb::parallel_for(tbb::blocked_range<size_t>(0, num_queries),
205 [&](const tbb::blocked_range<size_t> &r) {
206 for (size_t i = r.begin(); i != r.end(); ++i) {
207 int64_t start_idx = query_neighbors_row_splits[i];
208 std::copy(neighbors_indices[i].begin(),
209 neighbors_indices[i].end(),
210 &indices_ptr[start_idx]);
211
212 if (return_distances) {
213 std::copy(neighbors_distances[i].begin(),
214 neighbors_distances[i].end(),
215 &distances_ptr[start_idx]);
216 }
217 }
218 });
219}
220
221template <class T, class OUTPUT_ALLOCATOR, int METRIC>
222void _RadiusSearchCPU(NanoFlannIndexHolderBase *holder,
223 int64_t *query_neighbors_row_splits,
224 size_t num_points,
225 const T *const points,
226 size_t num_queries,
227 const T *const queries,
228 const size_t dimension,
229 const T *const radii,
230 bool ignore_query_point,
231 bool return_distances,
232 bool normalize_distances,
233 bool sort,
234 OUTPUT_ALLOCATOR &output_allocator) {
235 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
236 std::fill(query_neighbors_row_splits,
237 query_neighbors_row_splits + num_queries + 1, 0);
238 index_t *indices_ptr;
239 output_allocator.AllocIndices(&indices_ptr, 0);
240
241 T *distances_ptr;
242 output_allocator.AllocDistances(&distances_ptr, 0);
243 return;
244 }
245
246 auto points_equal = [](const T *const p1, const T *const p2,
247 size_t dimension) {
248 std::vector<T> p1_vec(p1, p1 + dimension);
249 std::vector<T> p2_vec(p2, p2 + dimension);
250 return p1_vec == p2_vec;
251 };
252
253 std::vector<std::vector<index_t>> neighbors_indices(num_queries);
254 std::vector<std::vector<T>> neighbors_distances(num_queries);
255 std::vector<uint32_t> neighbors_count(num_queries, 0);
256
257 nanoflann::SearchParams params;
258 params.sorted = sort;
259
260 auto holder_ =
261 static_cast<NanoFlannIndexHolder<METRIC, T, index_t> *>(holder);
262 tbb::parallel_for(
263 tbb::blocked_range<size_t>(0, num_queries),
264 [&](const tbb::blocked_range<size_t> &r) {
265 std::vector<std::pair<index_t, T>> search_result;
266 for (size_t i = r.begin(); i != r.end(); ++i) {
267 T radius = radii[i];
268 if (METRIC == L2) {
269 radius = radius * radius;
270 }
271
272 holder_->index_->radiusSearch(&queries[i * dimension],
273 radius, search_result,
274 params);
275
276 int num_neighbors = 0;
277 for (const auto &idx_dist : search_result) {
278 if (ignore_query_point &&
279 points_equal(&queries[i * dimension],
280 &points[idx_dist.first * dimension],
281 dimension)) {
282 continue;
283 }
284 neighbors_indices[i].push_back(idx_dist.first);
285 if (return_distances) {
286 neighbors_distances[i].push_back(idx_dist.second);
287 }
288 ++num_neighbors;
289 }
290 neighbors_count[i] = num_neighbors;
291 }
292 });
293
294 query_neighbors_row_splits[0] = 0;
295 utility::InclusivePrefixSum(neighbors_count.data(),
296 neighbors_count.data() + neighbors_count.size(),
297 query_neighbors_row_splits + 1);
298
299 int64_t num_indices = query_neighbors_row_splits[num_queries];
300
301 index_t *indices_ptr;
302 output_allocator.AllocIndices(&indices_ptr, num_indices);
303 T *distances_ptr;
304 if (return_distances)
305 output_allocator.AllocDistances(&distances_ptr, num_indices);
306 else
307 output_allocator.AllocDistances(&distances_ptr, 0);
308
309 std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
310
311 // fill output index and distance arrays
312 tbb::parallel_for(
313 tbb::blocked_range<size_t>(0, num_queries),
314 [&](const tbb::blocked_range<size_t> &r) {
315 for (size_t i = r.begin(); i != r.end(); ++i) {
316 int64_t start_idx = query_neighbors_row_splits[i];
317 std::copy(neighbors_indices[i].begin(),
318 neighbors_indices[i].end(),
319 &indices_ptr[start_idx]);
320 if (return_distances) {
321 std::transform(neighbors_distances[i].begin(),
322 neighbors_distances[i].end(),
323 &distances_ptr[start_idx], [&](T dist) {
324 T d = dist;
325 if (normalize_distances) {
326 if (METRIC == L2) {
327 d /= (radii[i] * radii[i]);
328 } else {
329 d /= radii[i];
330 }
331 }
332 return d;
333 });
334 }
335 }
336 });
337}
338
339template <class T, class OUTPUT_ALLOCATOR, int METRIC>
340void _HybridSearchCPU(NanoFlannIndexHolderBase *holder,
341 size_t num_points,
342 const T *const points,
343 size_t num_queries,
344 const T *const queries,
345 const size_t dimension,
346 const T radius,
347 const int max_knn,
348 bool ignore_query_point,
349 bool return_distances,
350 OUTPUT_ALLOCATOR &output_allocator) {
351 if (num_queries == 0 || num_points == 0 || holder == nullptr) {
352 index_t *indices_ptr, *counts_ptr;
353 output_allocator.AllocIndices(&indices_ptr, 0);
354 output_allocator.AllocCounts(&counts_ptr, 0);
355
356 T *distances_ptr;
357 output_allocator.AllocDistances(&distances_ptr, 0);
358 return;
359 }
360
361 T radius_squared = radius * radius;
362 index_t *indices_ptr, *counts_ptr;
363 T *distances_ptr;
364
365 size_t num_indices = static_cast<size_t>(max_knn) * num_queries;
366 output_allocator.AllocIndices(&indices_ptr, num_indices);
367 output_allocator.AllocDistances(&distances_ptr, num_indices);
368 output_allocator.AllocCounts(&counts_ptr, num_queries);
369
370 nanoflann::SearchParams params;
371 params.sorted = true;
372
373 auto holder_ =
374 static_cast<NanoFlannIndexHolder<METRIC, T, index_t> *>(holder);
375 tbb::parallel_for(
376 tbb::blocked_range<size_t>(0, num_queries),
377 [&](const tbb::blocked_range<size_t> &r) {
378 std::vector<std::pair<index_t, T>> ret_matches;
379 for (size_t i = r.begin(); i != r.end(); ++i) {
380 size_t num_results = holder_->index_->radiusSearch(
381 &queries[i * dimension], radius_squared,
382 ret_matches, params);
383 ret_matches.resize(num_results);
384
385 index_t count_i = static_cast<index_t>(num_results);
386 count_i = count_i < max_knn ? count_i : max_knn;
387 counts_ptr[i] = count_i;
388
389 int neighbor_idx = 0;
390 for (auto it = ret_matches.begin();
391 it < ret_matches.end() && neighbor_idx < max_knn;
392 it++, neighbor_idx++) {
393 indices_ptr[i * max_knn + neighbor_idx] = it->first;
394 distances_ptr[i * max_knn + neighbor_idx] = it->second;
395 }
396
397 while (neighbor_idx < max_knn) {
398 indices_ptr[i * max_knn + neighbor_idx] = -1;
399 distances_ptr[i * max_knn + neighbor_idx] = 0;
400 neighbor_idx += 1;
401 }
402 }
403 });
404}
405
406} // namespace
407
422template <class T>
423std::unique_ptr<NanoFlannIndexHolderBase> BuildKdTree(size_t num_points,
424 const T *const points,
425 size_t dimension,
426 const Metric metric) {
427 NanoFlannIndexHolderBase *holder = nullptr;
428#define FN_PARAMETERS num_points, points, dimension, &holder
429
430#define CALL_TEMPLATE(METRIC) \
431 if (METRIC == metric) { \
432 _BuildKdTree<T, METRIC>(FN_PARAMETERS); \
433 }
434
435#define CALL_TEMPLATE2 \
436 CALL_TEMPLATE(L1) \
437 CALL_TEMPLATE(L2)
438
440
441#undef CALL_TEMPLATE
442#undef CALL_TEMPLATE2
443
444#undef FN_PARAMETERS
445 return std::unique_ptr<NanoFlannIndexHolderBase>(holder);
446}
447
500template <class T, class OUTPUT_ALLOCATOR>
502 int64_t *query_neighbors_row_splits,
503 size_t num_points,
504 const T *const points,
505 size_t num_queries,
506 const T *const queries,
507 const size_t dimension,
508 int knn,
509 const Metric metric,
510 bool ignore_query_point,
511 bool return_distances,
512 OUTPUT_ALLOCATOR &output_allocator) {
513#define FN_PARAMETERS \
514 holder, query_neighbors_row_splits, num_points, points, num_queries, \
515 queries, dimension, knn, ignore_query_point, return_distances, \
516 output_allocator
517
518#define CALL_TEMPLATE(METRIC) \
519 if (METRIC == metric) { \
520 _KnnSearchCPU<T, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
521 }
522
523#define CALL_TEMPLATE2 \
524 CALL_TEMPLATE(L1) \
525 CALL_TEMPLATE(L2)
526
528
529#undef CALL_TEMPLATE
530#undef CALL_TEMPLATE2
531
532#undef FN_PARAMETERS
533}
534
594template <class T, class OUTPUT_ALLOCATOR>
596 int64_t *query_neighbors_row_splits,
597 size_t num_points,
598 const T *const points,
599 size_t num_queries,
600 const T *const queries,
601 const size_t dimension,
602 const T *const radii,
603 const Metric metric,
604 bool ignore_query_point,
605 bool return_distances,
606 bool normalize_distances,
607 bool sort,
608 OUTPUT_ALLOCATOR &output_allocator) {
609#define FN_PARAMETERS \
610 holder, query_neighbors_row_splits, num_points, points, num_queries, \
611 queries, dimension, radii, ignore_query_point, return_distances, \
612 normalize_distances, sort, output_allocator
613
614#define CALL_TEMPLATE(METRIC) \
615 if (METRIC == metric) { \
616 _RadiusSearchCPU<T, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
617 }
618
619#define CALL_TEMPLATE2 \
620 CALL_TEMPLATE(L1) \
621 CALL_TEMPLATE(L2)
622
624
625#undef CALL_TEMPLATE
626#undef CALL_TEMPLATE2
627
628#undef FN_PARAMETERS
629}
630
682template <class T, class OUTPUT_ALLOCATOR>
684 size_t num_points,
685 const T *const points,
686 size_t num_queries,
687 const T *const queries,
688 const size_t dimension,
689 const T radius,
690 const int max_knn,
691 const Metric metric,
692 bool ignore_query_point,
693 bool return_distances,
694 OUTPUT_ALLOCATOR &output_allocator) {
695#define FN_PARAMETERS \
696 holder, num_points, points, num_queries, queries, dimension, radius, \
697 max_knn, ignore_query_point, return_distances, output_allocator
698
699#define CALL_TEMPLATE(METRIC) \
700 if (METRIC == metric) { \
701 _HybridSearchCPU<T, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS); \
702 }
703
704#define CALL_TEMPLATE2 \
705 CALL_TEMPLATE(L1) \
706 CALL_TEMPLATE(L2)
707
709
710#undef CALL_TEMPLATE
711#undef CALL_TEMPLATE2
712
713#undef FN_PARAMETERS
714}
715
716} // namespace impl
717} // namespace nns
718} // namespace core
719} // 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:423
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:501
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:595
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:683
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
@ L1
Definition: NeighborSearchCommon.h:38
@ L2
Definition: NeighborSearchCommon.h:38
int32_t index_t
Definition: NanoFlannImpl.h:43
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t int32_t
Definition: K4aPlugin.cpp:414
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:67
Definition: PinholeCameraIntrinsic.cpp:35
This class is the Adaptor for connecting Open3D Tensor and NanoFlann.
Definition: NanoFlannImpl.h:49
const TReal *const data_ptr_
Definition: NanoFlannImpl.h:70
size_t dataset_size_
Definition: NanoFlannImpl.h:68
int dimension_
Definition: NanoFlannImpl.h:69
TReal kdtree_get_pt(const size_t idx, const size_t dim) const
Definition: NanoFlannImpl.h:59
DataAdaptor(size_t dataset_size, int dimension, const TReal *const data_ptr)
Definition: NanoFlannImpl.h:50
bool kdtree_get_bbox(BBOX &) const
Definition: NanoFlannImpl.h:64
size_t kdtree_get_point_count() const
Definition: NanoFlannImpl.h:57
nanoflann::L1_Adaptor< TReal, DataAdaptor, TReal > adaptor_t
Definition: NanoFlannImpl.h:84
nanoflann::L2_Adaptor< TReal, DataAdaptor, TReal > adaptor_t
Definition: NanoFlannImpl.h:79
Adaptor Selector.
Definition: NanoFlannImpl.h:75
Base struct for NanoFlann index holder.
Definition: NeighborSearchCommon.h:72
NanoFlann Index Holder.
Definition: NanoFlannImpl.h:47
std::unique_ptr< KDTree_t > index_
Definition: NanoFlannImpl.h:103
nanoflann::KDTreeSingleIndexAdaptor< typename SelectNanoflannAdaptor< METRIC >::adaptor_t, DataAdaptor, -1, TIndex > KDTree_t
typedef for KDtree.
Definition: NanoFlannImpl.h:93
std::unique_ptr< DataAdaptor > adaptor_
Definition: NanoFlannImpl.h:104
NanoFlannIndexHolder(size_t dataset_size, int dimension, const TReal *data_ptr)
Definition: NanoFlannImpl.h:95