Open3D (C++ API)  0.12.0
KnnSearch.h
Go to the documentation of this file.
1 // ----------------------------------------------------------------------------
2 // - Open3D: www.open3d.org -
3 // ----------------------------------------------------------------------------
4 // The MIT License (MIT)
5 //
6 // Copyright (c) 2019 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 
27 #pragma once
28 
29 #include <tbb/parallel_for.h>
30 
31 #include <mutex>
32 
33 #include "open3d/core/Atomic.h"
36 
37 namespace open3d {
38 namespace ml {
39 namespace impl {
40 
41 namespace {
42 
44 template <class T, class OUTPUT_ALLOCATOR, int METRIC>
45 void _KnnSearchCPU(int64_t* query_neighbors_row_splits,
46  size_t num_points,
47  const T* const points,
48  size_t num_queries,
49  const T* const queries,
50  const int k,
51  bool ignore_query_point,
52  bool return_distances,
53  OUTPUT_ALLOCATOR& output_allocator) {
54  using namespace open3d::utility;
55 
56  // return empty indices array if there are no points
57  if (num_points == 0 || num_queries == 0) {
58  std::fill(query_neighbors_row_splits,
59  query_neighbors_row_splits + num_queries + 1, 0);
60  int32_t* indices_ptr;
61  output_allocator.AllocIndices(&indices_ptr, 0);
62 
63  T* distances_ptr;
64  output_allocator.AllocDistances(&distances_ptr, 0);
65 
66  return;
67  }
68 
69  Adaptor<T> adaptor(num_points, points);
70 
72  typename SelectNanoflannAdaptor<METRIC, T>::Adaptor_t, Adaptor<T>,
73  3>
74  KDTree_t;
75 
76  KDTree_t index(3, adaptor);
77  index.buildIndex();
78 
79  // temporary storage for the result
80  struct Pair {
81  int32_t i, j;
82  };
83  std::vector<Pair> pairs;
84 
85  auto points_equal = [](const T* const p1, const T* const p2) {
86  return p1[0] == p2[0] && p1[1] == p2[1] && p1[2] == p2[2];
87  };
88 
89  auto distance_fn = [](const T* const p1, const T* const p2) {
90  T dx = p1[0] - p2[0];
91  T dy = p1[1] - p2[1];
92  T dz = p1[2] - p2[2];
93  if (METRIC == L2) {
94  return dx * dx + dy * dy + dz * dz;
95  } else // L1
96  {
97  return std::abs(dx) + std::abs(dy) + std::abs(dz);
98  }
99  };
100 
101  std::mutex pairs_mutex; // protects write access to pairs
102  std::vector<uint32_t> neighbors_count(num_queries, 0);
103 
104  // compute nearest neighbors and store in pairs
105  tbb::parallel_for(
106  tbb::blocked_range<size_t>(0, num_queries),
107  [&](const tbb::blocked_range<size_t>& r) {
108  std::vector<Pair> pairs_private;
109 
110  std::vector<size_t> result_indices(k);
111  std::vector<T> result_distances(k);
112  for (size_t i = r.begin(); i != r.end(); ++i) {
113  size_t num_valid = index.knnSearch(&queries[i * 3], k,
114  result_indices.data(),
115  result_distances.data());
116 
117  int num_neighbors = 0;
118  for (size_t valid_i = 0; valid_i < num_valid; ++valid_i) {
119  size_t idx = result_indices[valid_i];
120  if (ignore_query_point &&
121  points_equal(&queries[i * 3], &points[idx * 3])) {
122  continue;
123  }
124  pairs_private.push_back(Pair{int32_t(i), int32_t(idx)});
125  ++num_neighbors;
126  }
127  neighbors_count[i] = num_neighbors;
128  }
129  {
130  std::lock_guard<std::mutex> lock(pairs_mutex);
131  pairs.insert(pairs.end(), pairs_private.begin(),
132  pairs_private.end());
133  }
134  });
135 
136  query_neighbors_row_splits[0] = 0;
137  InclusivePrefixSum(&neighbors_count[0],
138  &neighbors_count[neighbors_count.size()],
139  query_neighbors_row_splits + 1);
140 
141  int32_t* neighbors_indices_ptr;
142  output_allocator.AllocIndices(&neighbors_indices_ptr, pairs.size());
143  T* distances_ptr;
144  if (return_distances)
145  output_allocator.AllocDistances(&distances_ptr, pairs.size());
146  else
147  output_allocator.AllocDistances(&distances_ptr, 0);
148 
149  std::fill(neighbors_count.begin(), neighbors_count.end(), 0);
150 
151  // fill output index and distance arrays
152  tbb::parallel_for(tbb::blocked_range<size_t>(0, pairs.size()),
153  [&](const tbb::blocked_range<size_t>& r) {
154  for (size_t i = r.begin(); i != r.end(); ++i) {
155  Pair pair = pairs[i];
156 
157  int64_t idx =
158  query_neighbors_row_splits[pair.i] +
160  &neighbors_count[pair.i], 1);
161  neighbors_indices_ptr[idx] = pair.j;
162 
163  if (return_distances) {
164  T dist = distance_fn(&points[pair.j * 3],
165  &queries[pair.i * 3]);
166  distances_ptr[idx] = dist;
167  }
168  }
169  });
170 }
171 
172 } // namespace
173 
221 template <class T, class OUTPUT_ALLOCATOR>
222 void KnnSearchCPU(int64_t* query_neighbors_row_splits,
223  size_t num_points,
224  const T* const points,
225  size_t num_queries,
226  const T* const queries,
227  const int k,
228  const Metric metric,
229  const bool ignore_query_point,
230  const bool return_distances,
231  OUTPUT_ALLOCATOR& output_allocator) {
232 #define FN_PARAMETERS \
233  query_neighbors_row_splits, num_points, points, num_queries, queries, k, \
234  ignore_query_point, return_distances, output_allocator
235 
236 #define CALL_TEMPLATE(METRIC) \
237  if (METRIC == metric) \
238  _KnnSearchCPU<T, OUTPUT_ALLOCATOR, METRIC>(FN_PARAMETERS);
239 
240 #define CALL_TEMPLATE2 \
241  CALL_TEMPLATE(L1) \
242  CALL_TEMPLATE(L2)
243 
245 
246 #undef CALL_TEMPLATE
247 #undef CALL_TEMPLATE2
248 
249 #undef FN_PARAMETERS
250 }
251 
252 } // namespace impl
253 } // namespace ml
254 } // namespace open3d
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:67
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
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:398
#define CALL_TEMPLATE2
int points
Definition: FilePCD.cpp:73
Definition: NanoFlannIndex.h:45
Definition: NeighborSearchCommon.h:38
Definition: PinholeCameraIntrinsic.cpp:35
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition: Atomic.h:40
Definition: Console.cpp:51
void KnnSearchCPU(int64_t *query_neighbors_row_splits, size_t num_points, const T *const points, size_t num_queries, const T *const queries, const int k, const Metric metric, const bool ignore_query_point, const bool return_distances, OUTPUT_ALLOCATOR &output_allocator)
Definition: KnnSearch.h:222