Open3D (C++ API)  0.12.0
FixedRadiusSearch.h
Go to the documentation of this file.
1 // ----------------------------------------------------------------------------
2 // - Open3D: www.open3d.org -
3 // ----------------------------------------------------------------------------
4 // The MIT License (MIT)
5 //
6 // Copyright (c) 2020 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 <set>
32 
33 #include "open3d/core/Atomic.h"
35 #include "open3d/utility/Eigen.h"
36 #include "open3d/utility/Helper.h"
38 
39 namespace open3d {
40 namespace ml {
41 namespace impl {
42 
43 namespace {
44 
77 template <class TReal, class TIndex>
78 void BuildSpatialHashTableCPU(const size_t num_points,
79  const TReal* const points,
80  const TReal radius,
81  const size_t points_row_splits_size,
82  const int64_t* points_row_splits,
83  const TIndex* hash_table_splits,
84  const size_t hash_table_cell_splits_size,
85  TIndex* hash_table_cell_splits,
86  TIndex* hash_table_index) {
87  using namespace open3d::utility;
88  typedef MiniVec<TReal, 3> Vec3_t;
89  // typedef Eigen::Array<TReal, 3, 1> Vec3_t;
90 
91  const int batch_size = points_row_splits_size - 1;
92  const TReal voxel_size = 2 * radius;
93  const TReal inv_voxel_size = 1 / voxel_size;
94 
95  memset(&hash_table_cell_splits[0], 0,
96  sizeof(TIndex) * hash_table_cell_splits_size);
97 
98  // compute number of points that map to each hash
99  for (int i = 0; i < batch_size; ++i) {
100  const size_t hash_table_size =
101  hash_table_splits[i + 1] - hash_table_splits[i];
102  const size_t first_cell_idx = hash_table_splits[i];
103  tbb::parallel_for(
104  tbb::blocked_range<int64_t>(points_row_splits[i],
105  points_row_splits[i + 1]),
106  [&](const tbb::blocked_range<int64_t>& r) {
107  for (int64_t i = r.begin(); i != r.end(); ++i) {
108  Vec3_t pos(points + 3 * i);
109 
110  auto voxel_index =
111  ComputeVoxelIndex(pos, inv_voxel_size);
112  size_t hash =
113  SpatialHash(voxel_index) % hash_table_size;
114 
115  // note the +1 because we want the first element to be 0
117  &hash_table_cell_splits[first_cell_idx + hash +
118  1],
119  1);
120  }
121  });
122  }
123  InclusivePrefixSum(&hash_table_cell_splits[0],
124  &hash_table_cell_splits[hash_table_cell_splits_size],
125  &hash_table_cell_splits[0]);
126 
127  std::vector<uint32_t> count_tmp(hash_table_cell_splits_size - 1, 0);
128 
129  // now compute the indices for hash_table_index
130  for (int i = 0; i < batch_size; ++i) {
131  const size_t hash_table_size =
132  hash_table_splits[i + 1] - hash_table_splits[i];
133  const size_t first_cell_idx = hash_table_splits[i];
134  tbb::parallel_for(
135  tbb::blocked_range<size_t>(points_row_splits[i],
136  points_row_splits[i + 1]),
137  [&](const tbb::blocked_range<size_t>& r) {
138  for (size_t i = r.begin(); i != r.end(); ++i) {
139  Vec3_t pos(points + 3 * i);
140 
141  auto voxel_index =
142  ComputeVoxelIndex(pos, inv_voxel_size);
143  size_t hash =
144  SpatialHash(voxel_index) % hash_table_size;
145 
146  hash_table_index
147  [hash_table_cell_splits[hash + first_cell_idx] +
149  &count_tmp[hash + first_cell_idx],
150  1)] = i;
151  }
152  });
153  }
154 }
155 
173 template <int METRIC, class TDerived, int VECSIZE>
174 Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> NeighborsDist(
175  const Eigen::ArrayBase<TDerived>& p,
176  const Eigen::Array<typename TDerived::Scalar, VECSIZE, 3>& points) {
177  typedef Eigen::Array<typename TDerived::Scalar, VECSIZE, 1> VecN_t;
178  VecN_t dist;
179 
180  dist.setZero();
181  if (METRIC == Linf) {
182  dist = (points.rowwise() - p.transpose()).abs().rowwise().maxCoeff();
183  } else if (METRIC == L1) {
184  dist = (points.rowwise() - p.transpose()).abs().rowwise().sum();
185  } else {
186  dist = (points.rowwise() - p.transpose()).square().rowwise().sum();
187  }
188  return dist;
189 }
190 
193 template <class T,
194  class OUTPUT_ALLOCATOR,
195  int METRIC,
196  bool IGNORE_QUERY_POINT,
197  bool RETURN_DISTANCES>
198 void _FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
199  size_t num_points,
200  const T* const points,
201  size_t num_queries,
202  const T* const queries,
203  const T radius,
204  const size_t points_row_splits_size,
205  const int64_t* const points_row_splits,
206  const size_t queries_row_splits_size,
207  const int64_t* const queries_row_splits,
208  const uint32_t* const hash_table_splits,
209  const size_t hash_table_cell_splits_size,
210  const uint32_t* const hash_table_cell_splits,
211  const uint32_t* const hash_table_index,
212  OUTPUT_ALLOCATOR& output_allocator) {
213  using namespace open3d::utility;
214 
215  // number of elements for vectorization
216  const int VECSIZE = 8;
217  typedef MiniVec<T, 3> Vec3_t;
218  typedef Eigen::Array<T, VECSIZE, 1> Vec_t;
219  typedef Eigen::Array<int32_t, VECSIZE, 1> Veci_t;
220 
221  const int batch_size = points_row_splits_size - 1;
222 
223  // return empty output arrays if there are no points
224  if (num_points == 0 || num_queries == 0) {
225  std::fill(query_neighbors_row_splits,
226  query_neighbors_row_splits + num_queries + 1, 0);
227  int32_t* indices_ptr;
228  output_allocator.AllocIndices(&indices_ptr, 0);
229 
230  T* distances_ptr;
231  output_allocator.AllocDistances(&distances_ptr, 0);
232 
233  return;
234  }
235 
236  // use squared radius for L2 to avoid sqrt
237  const T threshold = (METRIC == L2 ? radius * radius : radius);
238 
239  const T voxel_size = 2 * radius;
240  const T inv_voxel_size = 1 / voxel_size;
241 
242  // counts the number of indices we have to return. This is the number of all
243  // neighbors we find.
244  size_t num_indices = 0;
245 
246  // count the number of neighbors for all query points and update num_indices
247  // and populate query_neighbors_row_splits with the number of neighbors
248  // for each query point
249  for (int i = 0; i < batch_size; ++i) {
250  const size_t hash_table_size =
251  hash_table_splits[i + 1] - hash_table_splits[i];
252  const size_t first_cell_idx = hash_table_splits[i];
253  tbb::parallel_for(
254  tbb::blocked_range<size_t>(queries_row_splits[i],
255  queries_row_splits[i + 1]),
256  [&](const tbb::blocked_range<size_t>& r) {
257  size_t num_indices_local = 0;
258  for (size_t i = r.begin(); i != r.end(); ++i) {
259  size_t neighbors_count = 0;
260 
261  Vec3_t pos(queries + i * 3);
262 
263  std::set<size_t> bins_to_visit;
264 
265  auto voxel_index =
266  ComputeVoxelIndex(pos, inv_voxel_size);
267  size_t hash =
268  SpatialHash(voxel_index) % hash_table_size;
269 
270  bins_to_visit.insert(first_cell_idx + hash);
271 
272  for (int dz = -1; dz <= 1; dz += 2)
273  for (int dy = -1; dy <= 1; dy += 2)
274  for (int dx = -1; dx <= 1; dx += 2) {
275  Vec3_t p =
276  pos + radius * Vec3_t(T(dx), T(dy),
277  T(dz));
278  voxel_index = ComputeVoxelIndex(
279  p, inv_voxel_size);
280  hash = SpatialHash(voxel_index) %
281  hash_table_size;
282  bins_to_visit.insert(first_cell_idx + hash);
283  }
284 
285  Eigen::Array<T, VECSIZE, 3> xyz;
286  int vec_i = 0;
287 
288  for (size_t bin : bins_to_visit) {
289  size_t begin_idx = hash_table_cell_splits[bin];
290  size_t end_idx = hash_table_cell_splits[bin + 1];
291 
292  for (size_t j = begin_idx; j < end_idx; ++j) {
293  uint32_t idx = hash_table_index[j];
294  if (IGNORE_QUERY_POINT) {
295  if (points[idx * 3 + 0] == pos[0] &&
296  points[idx * 3 + 1] == pos[1] &&
297  points[idx * 3 + 2] == pos[2])
298  continue;
299  }
300  xyz(vec_i, 0) = points[idx * 3 + 0];
301  xyz(vec_i, 1) = points[idx * 3 + 1];
302  xyz(vec_i, 2) = points[idx * 3 + 2];
303  ++vec_i;
304  if (VECSIZE == vec_i) {
305  Eigen::Array<T, 3, 1> pos_arr(
306  pos[0], pos[1], pos[2]);
307  Vec_t dist =
308  NeighborsDist<METRIC>(pos_arr, xyz);
309  Eigen::Array<bool, VECSIZE, 1> test_result =
310  dist <= threshold;
311  neighbors_count += test_result.count();
312  vec_i = 0;
313  }
314  }
315  }
316  // process the tail
317  if (vec_i) {
318  Eigen::Array<T, 3, 1> pos_arr(pos[0], pos[1],
319  pos[2]);
320  Eigen::Array<T, VECSIZE, 1> dist =
321  NeighborsDist<METRIC>(pos_arr, xyz);
322  Eigen::Array<bool, VECSIZE, 1> test_result =
323  dist <= threshold;
324  for (int k = 0; k < vec_i; ++k) {
325  neighbors_count += int(test_result(k));
326  }
327  vec_i = 0;
328  }
329  num_indices_local += neighbors_count;
330  // note the +1
331  query_neighbors_row_splits[i + 1] = neighbors_count;
332  }
333 
334  core::AtomicFetchAddRelaxed((uint64_t*)&num_indices,
335  num_indices_local);
336  });
337  }
338 
339  // Allocate output arrays
340  // output for the indices to the neighbors
341  int32_t* indices_ptr;
342  output_allocator.AllocIndices(&indices_ptr, num_indices);
343 
344  // output for the distances
345  T* distances_ptr;
346  if (RETURN_DISTANCES)
347  output_allocator.AllocDistances(&distances_ptr, num_indices);
348  else
349  output_allocator.AllocDistances(&distances_ptr, 0);
350 
351  query_neighbors_row_splits[0] = 0;
352  InclusivePrefixSum(query_neighbors_row_splits + 1,
353  query_neighbors_row_splits + num_queries + 1,
354  query_neighbors_row_splits + 1);
355 
356  // now populate the indices_ptr and distances_ptr array
357  for (int i = 0; i < batch_size; ++i) {
358  const size_t hash_table_size =
359  hash_table_splits[i + 1] - hash_table_splits[i];
360  const size_t first_cell_idx = hash_table_splits[i];
361  tbb::parallel_for(
362  tbb::blocked_range<size_t>(queries_row_splits[i],
363  queries_row_splits[i + 1]),
364  [&](const tbb::blocked_range<size_t>& r) {
365  for (size_t i = r.begin(); i != r.end(); ++i) {
366  size_t neighbors_count = 0;
367 
368  size_t indices_offset = query_neighbors_row_splits[i];
369 
370  Vec3_t pos(queries[i * 3 + 0], queries[i * 3 + 1],
371  queries[i * 3 + 2]);
372 
373  std::set<size_t> bins_to_visit;
374 
375  auto voxel_index =
376  ComputeVoxelIndex(pos, inv_voxel_size);
377  size_t hash =
378  SpatialHash(voxel_index) % hash_table_size;
379 
380  bins_to_visit.insert(first_cell_idx + hash);
381 
382  for (int dz = -1; dz <= 1; dz += 2)
383  for (int dy = -1; dy <= 1; dy += 2)
384  for (int dx = -1; dx <= 1; dx += 2) {
385  Vec3_t p =
386  pos + radius * Vec3_t(T(dx), T(dy),
387  T(dz));
388  voxel_index = ComputeVoxelIndex(
389  p, inv_voxel_size);
390  hash = SpatialHash(voxel_index) %
391  hash_table_size;
392  bins_to_visit.insert(first_cell_idx + hash);
393  }
394 
395  Eigen::Array<T, VECSIZE, 3> xyz;
396  Veci_t idx_vec;
397  int vec_i = 0;
398 
399  for (size_t bin : bins_to_visit) {
400  size_t begin_idx = hash_table_cell_splits[bin];
401  size_t end_idx = hash_table_cell_splits[bin + 1];
402 
403  for (size_t j = begin_idx; j < end_idx; ++j) {
404  uint32_t idx = hash_table_index[j];
405  if (IGNORE_QUERY_POINT) {
406  if (points[idx * 3 + 0] == pos[0] &&
407  points[idx * 3 + 1] == pos[1] &&
408  points[idx * 3 + 2] == pos[2])
409  continue;
410  }
411  xyz(vec_i, 0) = points[idx * 3 + 0];
412  xyz(vec_i, 1) = points[idx * 3 + 1];
413  xyz(vec_i, 2) = points[idx * 3 + 2];
414  idx_vec(vec_i) = idx;
415  ++vec_i;
416  if (VECSIZE == vec_i) {
417  Eigen::Array<T, 3, 1> pos_arr(
418  pos[0], pos[1], pos[2]);
419  Eigen::Array<T, VECSIZE, 1> dist =
420  NeighborsDist<METRIC>(pos_arr, xyz);
421  Eigen::Array<bool, VECSIZE, 1> test_result =
422  dist <= threshold;
423  for (int k = 0; k < vec_i; ++k) {
424  if (test_result(k)) {
425  indices_ptr[indices_offset +
426  neighbors_count] =
427  idx_vec[k];
428  if (RETURN_DISTANCES) {
429  distances_ptr[indices_offset +
430  neighbors_count] =
431  dist[k];
432  }
433  }
434  neighbors_count += int(test_result(k));
435  }
436  vec_i = 0;
437  }
438  }
439  }
440  // process the tail
441  if (vec_i) {
442  Eigen::Array<T, 3, 1> pos_arr(pos[0], pos[1],
443  pos[2]);
444  Eigen::Array<T, VECSIZE, 1> dist =
445  NeighborsDist<METRIC>(pos_arr, xyz);
446  Eigen::Array<bool, VECSIZE, 1> test_result =
447  dist <= threshold;
448  for (int k = 0; k < vec_i; ++k) {
449  if (test_result(k)) {
450  indices_ptr[indices_offset +
451  neighbors_count] = idx_vec[k];
452  if (RETURN_DISTANCES) {
453  distances_ptr[indices_offset +
454  neighbors_count] =
455  dist[k];
456  }
457  }
458  neighbors_count += int(test_result(k));
459  }
460  vec_i = 0;
461  }
462  }
463  });
464  }
465 }
466 
467 } // namespace
468 
547 template <class T, class OUTPUT_ALLOCATOR>
548 void FixedRadiusSearchCPU(int64_t* query_neighbors_row_splits,
549  const size_t num_points,
550  const T* const points,
551  const size_t num_queries,
552  const T* const queries,
553  const T radius,
554  const size_t points_row_splits_size,
555  const int64_t* const points_row_splits,
556  const size_t queries_row_splits_size,
557  const int64_t* const queries_row_splits,
558  const uint32_t* const hash_table_splits,
559  const size_t hash_table_cell_splits_size,
560  const uint32_t* const hash_table_cell_splits,
561  const uint32_t* const hash_table_index,
562  const Metric metric,
563  const bool ignore_query_point,
564  const bool return_distances,
565  OUTPUT_ALLOCATOR& output_allocator) {
566  // Dispatch all template parameter combinations
567 
568 #define FN_PARAMETERS \
569  query_neighbors_row_splits, num_points, points, num_queries, queries, \
570  radius, points_row_splits_size, points_row_splits, \
571  queries_row_splits_size, queries_row_splits, hash_table_splits, \
572  hash_table_cell_splits_size, hash_table_cell_splits, \
573  hash_table_index, output_allocator
574 
575 #define CALL_TEMPLATE(METRIC, IGNORE_QUERY_POINT, RETURN_DISTANCES) \
576  if (METRIC == metric && IGNORE_QUERY_POINT == ignore_query_point && \
577  RETURN_DISTANCES == return_distances) \
578  _FixedRadiusSearchCPU<T, OUTPUT_ALLOCATOR, METRIC, IGNORE_QUERY_POINT, \
579  RETURN_DISTANCES>(FN_PARAMETERS);
580 
581 #define CALL_TEMPLATE2(METRIC) \
582  CALL_TEMPLATE(METRIC, true, true) \
583  CALL_TEMPLATE(METRIC, true, false) \
584  CALL_TEMPLATE(METRIC, false, true) \
585  CALL_TEMPLATE(METRIC, false, false)
586 
587 #define CALL_TEMPLATE3 \
588  CALL_TEMPLATE2(L1) \
589  CALL_TEMPLATE2(L2) \
590  CALL_TEMPLATE2(Linf)
591 
593 
594 #undef CALL_TEMPLATE
595 #undef CALL_TEMPLATE2
596 #undef CALL_TEMPLATE3
597 #undef FN_PARAMETERS
598 }
599 
600 } // namespace impl
601 } // namespace ml
602 } // namespace open3d
void FixedRadiusSearchCPU(int64_t *query_neighbors_row_splits, const size_t num_points, const T *const points, const size_t num_queries, const T *const queries, const T radius, const size_t points_row_splits_size, const int64_t *const points_row_splits, const size_t queries_row_splits_size, const int64_t *const queries_row_splits, const uint32_t *const hash_table_splits, const size_t hash_table_cell_splits_size, const uint32_t *const hash_table_cell_splits, const uint32_t *const hash_table_index, const Metric metric, const bool ignore_query_point, const bool return_distances, OUTPUT_ALLOCATOR &output_allocator)
Definition: FixedRadiusSearch.h:548
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 timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c k4a_image_t image_handle uint8_t image_handle image_handle image_handle image_handle uint32_t
Definition: K4aPlugin.cpp:557
void InclusivePrefixSum(const Tin *first, const Tin *last, Tout *out)
Definition: ParallelScan.h:67
void BuildSpatialHashTableCPU(const torch::Tensor &points, double radius, const torch::Tensor &points_row_splits, const std::vector< uint32_t > &hash_table_splits, torch::Tensor &hash_table_index, torch::Tensor &hash_table_cell_splits)
Definition: BuildSpatialHashTableOpKernel.cpp:33
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
HOST_DEVICE utility::MiniVec< int, 3 > ComputeVoxelIndex(const TVecf &pos, const typename TVecf::Scalar_t &inv_voxel_size)
Definition: NeighborSearchCommon.h:61
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
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 uint64_t
Definition: K4aPlugin.cpp:352
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 timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c int
Definition: K4aPlugin.cpp:479
#define CALL_TEMPLATE3
int points
Definition: FilePCD.cpp:73
Definition: NeighborSearchCommon.h:38
Definition: PinholeCameraIntrinsic.cpp:35
HOST_DEVICE size_t SpatialHash(int x, int y, int z)
Spatial hashing function for integer coordinates.
Definition: NeighborSearchCommon.h:47
uint32_t AtomicFetchAddRelaxed(uint32_t *address, uint32_t val)
Definition: Atomic.h:40
Definition: Console.cpp:51
Definition: MiniVec.h:43