Open3D (C++ API)  0.18.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"
17 #include "open3d/utility/Helper.h"
19 
20 namespace open3d {
21 namespace core {
22 namespace nns {
23 
25 template <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 };
85 namespace impl {
86 
87 namespace {
88 template <class T, class TIndex, int METRIC>
89 void _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 
97 template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
98 void _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 
200 template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
201 void _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 
318 template <class T, class TIndex, class OUTPUT_ALLOCATOR, int METRIC>
319 void _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 
401 template <class T, class TIndex>
402 std::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 
479 template <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 
573 template <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 
661 template <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