Open3D (C++ API)  0.18.0+252c867
VoxelPooling.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 
8 #pragma once
9 
10 #include <tbb/task_group.h>
11 
12 #include <Eigen/Core>
13 #include <unordered_map>
14 
15 #include "open3d/utility/Helper.h"
16 
17 namespace open3d {
18 namespace ml {
19 namespace impl {
20 
22 
23 template <class TReal,
24  class TFeat,
25  AccumulationFn POS_FN,
26  AccumulationFn FEAT_FN>
27 class Accumulator {
28 public:
30  : count_(0),
31  min_sqr_dist_to_center_(std::numeric_limits<TReal>::max()),
32  position_(0, 0, 0) {
33  static_assert(POS_FN != MAX, "MAX is not allowed for point positions");
34  static_assert(FEAT_FN != CENTER,
35  "CENTER is not allowed for feature vectors");
36  }
37 
38  template <class Derived, class Derived2, class Derived3>
39  inline void AddPoint(const Eigen::MatrixBase<Derived>& pos,
40  const Eigen::MatrixBase<Derived2>& voxel_center,
41  const Eigen::ArrayBase<Derived3>& feat) {
42  bool new_nearest_neighbor = false;
43  TReal sqr_d = 0;
44  if (POS_FN == NEAREST_NEIGHBOR || FEAT_FN == NEAREST_NEIGHBOR) {
45  sqr_d = (voxel_center - pos).squaredNorm();
46  if (sqr_d < min_sqr_dist_to_center_) {
47  new_nearest_neighbor = true;
48  min_sqr_dist_to_center_ = sqr_d;
49  }
50  }
51  if (POS_FN == AVERAGE) {
52  position_ += pos.array();
53  } else if (POS_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
54  position_ = pos;
55  } else if (POS_FN == CENTER) {
56  if (count_ == 0) position_ = voxel_center;
57  }
58 
59  if (count_ == 0) {
60  features_.resizeLike(feat);
61  features_.setZero();
62  }
63  if (FEAT_FN == AVERAGE) {
64  features_ += feat;
65  } else if (FEAT_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
66  features_ = feat;
67  } else if (FEAT_FN == MAX) {
68  features_ = features_.max(feat);
69  }
70  ++count_;
71  }
72 
73  inline Eigen::Array<TReal, 3, 1> Position() const {
74  if (POS_FN == AVERAGE) {
75  return position_ / count_;
76  } else // if( POS_FN == NEAREST_NEIGHBOR || POS_FN == CENTER )
77  {
78  return position_;
79  }
80  }
81 
82  inline Eigen::Array<TFeat, Eigen::Dynamic, 1> Features() const {
83  if (FEAT_FN == AVERAGE) {
84  return features_ / count_;
85  } else // if( FEAT_FN == NEAREST_NEIGHBOR || FEAT_FN == MAX )
86  {
87  return features_;
88  }
89  }
90 
91  inline int Count() const { return count_; }
92 
93 private:
94  int count_;
95  TReal min_sqr_dist_to_center_;
96  Eigen::Array<TReal, 3, 1> position_;
97  Eigen::Array<TFeat, Eigen::Dynamic, 1> features_;
98 }; // namespace
99 
100 template <class TReal,
101  class TFeat,
102  AccumulationFn POS_FN,
103  AccumulationFn FEAT_FN>
105 public:
107  : count_(0),
108  min_sqr_dist_to_center_(std::numeric_limits<TReal>::max()),
109  position_(0, 0, 0) {
110  static_assert(POS_FN != MAX, "MAX is not allowed for point positions");
111  static_assert(FEAT_FN != CENTER,
112  "CENTER is not allowed for feature vectors");
113  }
114 
115  template <class Derived, class Derived2, class Derived3>
116  inline void AddPoint(const Eigen::MatrixBase<Derived>& pos,
117  const Eigen::MatrixBase<Derived2>& voxel_center,
118  const Eigen::ArrayBase<Derived3>& feat,
119  const size_t idx) {
120  bool new_nearest_neighbor = false;
121  TReal sqr_d = 0;
122  if (POS_FN == NEAREST_NEIGHBOR || FEAT_FN == NEAREST_NEIGHBOR) {
123  sqr_d = (voxel_center - pos).squaredNorm();
124  if (sqr_d < min_sqr_dist_to_center_) {
125  new_nearest_neighbor = true;
126  min_sqr_dist_to_center_ = sqr_d;
127  }
128  }
129 
130  if (POS_FN == AVERAGE) {
131  position_ += pos.array();
132  } else if (POS_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
133  position_ = pos;
134  } else if (POS_FN == CENTER) {
135  if (count_ == 0) position_ = voxel_center;
136  }
137 
138  if (count_ == 0) {
139  features_.resizeLike(feat);
140  features_.setZero();
141  if (FEAT_FN == NEAREST_NEIGHBOR) {
142  features_ = feat;
143  index_.resize(1);
144  index_(0) = idx;
145  ++count_;
146  return;
147  } else if (FEAT_FN == MAX) {
148  features_ = feat;
149  index_.resizeLike(feat);
150  index_ = idx;
151  ++count_;
152  return;
153  }
154  }
155  if (FEAT_FN == AVERAGE) {
156  features_ += feat;
157  } else if (FEAT_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
158  features_ = feat;
159  index_(0) = idx;
160  } else if (FEAT_FN == MAX) {
161  for (int i = 0; i < features_.rows(); ++i) {
162  if (feat(i) > features_(i)) {
163  features_(i) = feat(i);
164  index_(i) = idx;
165  }
166  }
167  }
168  ++count_;
169  }
170 
171  inline Eigen::Array<TReal, 3, 1> Position() const {
172  if (POS_FN == AVERAGE) {
173  return position_ / count_;
174  } else // if( POS_FN == NEAREST_NEIGHBOR || POS_FN == CENTER )
175  {
176  return position_;
177  }
178  }
179 
180  inline Eigen::Array<TFeat, Eigen::Dynamic, 1> Features() const {
181  if (FEAT_FN == AVERAGE) {
182  return features_ / count_;
183  } else // if( FEAT_FN == NEAREST_NEIGHBOR || FEAT_FN == MAX )
184  {
185  return features_;
186  }
187  }
188 
189  inline int Count() const { return count_; }
190 
191  inline Eigen::Array<size_t, Eigen::Dynamic, 1> Index() const {
192  return index_;
193  }
194 
195 private:
196  int count_;
197  TReal min_sqr_dist_to_center_;
198  Eigen::Array<TReal, 3, 1> position_;
199  Eigen::Array<TFeat, Eigen::Dynamic, 1> features_;
200  Eigen::Array<size_t, Eigen::Dynamic, 1> index_;
201 };
202 
204 template <class T>
205 bool CheckVoxelSize(std::string& err,
206  const size_t num_positions,
207  const T* const positions,
208  const T voxel_size) {
209  typedef Eigen::Array<double, 3, 1> Vec3_t;
210  if (num_positions == 0) {
211  return true;
212  }
213 
214  Vec3_t bb_min, bb_max;
215  bb_min << positions[0], positions[1], positions[2];
216  bb_max = bb_min;
217 
218  Vec3_t voxel_size3(voxel_size, voxel_size, voxel_size);
219 
220  for (size_t i = 1; i < num_positions; ++i) {
221  Vec3_t pos(positions[i * 3 + 0], positions[i * 3 + 1],
222  positions[i * 3 + 2]);
223  bb_min = bb_min.min(pos);
224  bb_max = bb_max.max(pos);
225  }
226 
227  // the min and max bounding box shall be a multiple of the voxel size
228  bb_min /= voxel_size3;
229  bb_min = bb_min.floor() * voxel_size3;
230  bb_max /= voxel_size3;
231  bb_max = bb_max.ceil() * voxel_size3;
232 
233  if (voxel_size * double(std::numeric_limits<int>::max()) <
234  bb_max.maxCoeff() ||
235  voxel_size * double(std::numeric_limits<int>::min()) >
236  bb_min.maxCoeff()) {
237  err = "voxel_size is too small\n";
238  return false;
239  }
240  return true;
241 }
242 
248 template <class TDerived>
249 Eigen::Vector3i ComputeVoxelIndex(
250  const Eigen::ArrayBase<TDerived>& pos,
251  const typename TDerived::Scalar& inv_voxel_size) {
252  typedef typename TDerived::Scalar Scalar_t;
253  Eigen::Array<Scalar_t, 3, 1> ref_coord = pos * inv_voxel_size;
254 
255  Eigen::Vector3i voxel_index;
256  voxel_index = ref_coord.floor().template cast<int>();
257  return voxel_index;
258 }
259 
260 // implementation for VoxelPooling with template parameter for the accumulator.
261 template <class TReal, class TFeat, class ACCUMULATOR, class OUTPUT_ALLOCATOR>
262 void _VoxelPooling(size_t num_inp,
263  const TReal* const inp_positions,
264  int in_channels,
265  const TFeat* inp_features,
266  TReal voxel_size,
267  OUTPUT_ALLOCATOR& output_allocator) {
268  if (num_inp == 0) {
269  TReal* out_pos_ptr;
270  TFeat* out_feat_ptr;
271  output_allocator.AllocPooledPositions(&out_pos_ptr, 0);
272  output_allocator.AllocPooledFeatures(&out_feat_ptr, 0, in_channels);
273  return;
274  }
275 
276  typedef Eigen::Array<TReal, 3, 1> Vec3_t;
277  typedef Eigen::Array<TFeat, Eigen::Dynamic, 1> FeatureVec_t;
278 
279  std::unordered_map<Eigen::Vector3i, ACCUMULATOR,
281  voxelindex_to_accpoint;
282 
283  Vec3_t voxel_center;
284  Eigen::Vector3i voxel_index;
285  TReal inv_voxel_size = 1 / voxel_size;
286  TReal half_voxel_size = 0.5 * voxel_size;
287  for (size_t i = 0; i < num_inp; ++i) {
288  Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
289 
290  voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
291 
292  voxel_center << voxel_index(0) * voxel_size + half_voxel_size,
293  voxel_index(1) * voxel_size + half_voxel_size,
294  voxel_index(2) * voxel_size + half_voxel_size;
295 
296  Eigen::Map<const FeatureVec_t> feat(inp_features + in_channels * i,
297  in_channels);
298  voxelindex_to_accpoint[voxel_index].AddPoint(
299  pos.matrix(), voxel_center.matrix(), feat);
300  }
301 
302  const size_t num_out = voxelindex_to_accpoint.size();
303 
304  TReal* out_pos_ptr;
305  TFeat* out_feat_ptr;
306  output_allocator.AllocPooledPositions(&out_pos_ptr, num_out);
307  output_allocator.AllocPooledFeatures(&out_feat_ptr, num_out, in_channels);
308 
309  size_t i = 0;
310  for (const auto point : voxelindex_to_accpoint) {
311  Vec3_t pos = point.second.Position();
312  Eigen::Map<Vec3_t> out_pos(out_pos_ptr + i * 3);
313  out_pos = pos;
314 
315  Eigen::Map<FeatureVec_t> out_feat(out_feat_ptr + i * in_channels,
316  in_channels);
317  out_feat = point.second.Features();
318  ++i;
319  }
320 }
321 
322 // implementation for VoxelPoolingBackprop with template parameter for the
323 // accumulator.
324 template <class TReal, class TFeat, class ACCUMULATOR, AccumulationFn FEAT_FN>
325 void _VoxelPoolingBackprop(TFeat* features_backprop,
326  size_t num_inp,
327  const TReal* const inp_positions,
328  int in_channels,
329  const TFeat* const inp_features,
330  size_t num_pooled,
331  const TReal* const pooled_positions,
332  const TFeat* const pooled_features_gradient,
333  TReal voxel_size) {
334  if (num_inp == 0) {
335  return;
336  }
337  memset(features_backprop, 0, sizeof(TFeat) * num_inp * in_channels);
338 
339  typedef Eigen::Array<TReal, 3, 1> Vec3_t;
340  typedef Eigen::Array<TFeat, Eigen::Dynamic, 1> FeatureVec_t;
341 
342  Vec3_t voxel_size3(voxel_size, voxel_size, voxel_size);
343 
344  // create the two hash maps in parallel
345  tbb::task_group task_group;
346 
347  std::unordered_map<Eigen::Vector3i, ACCUMULATOR,
349  voxelindex_to_accpoint;
350 
351  task_group.run([&] {
352  Vec3_t voxel_center;
353  Eigen::Vector3i voxel_index;
354  TReal inv_voxel_size = 1 / voxel_size;
355  TReal half_voxel_size = 0.5 * voxel_size;
356  for (size_t i = 0; i < num_inp; ++i) {
357  Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
358 
359  voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
360 
361  voxel_center << voxel_index(0) * voxel_size + half_voxel_size,
362  voxel_index(1) * voxel_size + half_voxel_size,
363  voxel_index(2) * voxel_size + half_voxel_size;
364 
365  Eigen::Map<const FeatureVec_t> feat(inp_features + in_channels * i,
366  in_channels);
367  voxelindex_to_accpoint[voxel_index].AddPoint(
368  pos.matrix(), voxel_center.matrix(), feat, i);
369  }
370  });
371 
372  std::unordered_map<Eigen::Vector3i, size_t,
374  voxelindex_to_gradindex;
375 
376  task_group.run([&] {
377  Eigen::Vector3i voxel_index;
378  TReal inv_voxel_size = 1 / voxel_size;
379  for (size_t i = 0; i < num_pooled; ++i) {
380  Eigen::Map<const Vec3_t> pos(pooled_positions + i * 3);
381 
382  voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
383 
384  voxelindex_to_gradindex[voxel_index] = i;
385  }
386  });
387 
388  task_group.wait();
389 
390  if (FEAT_FN == AVERAGE) {
391  Eigen::Vector3i voxel_index;
392  TReal inv_voxel_size = 1 / voxel_size;
393  for (size_t i = 0; i < num_inp; ++i) {
394  Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
395 
396  voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
397 
398  Eigen::Map<FeatureVec_t> feat_bp(
399  features_backprop + in_channels * i, in_channels);
400 
401  size_t grad_idx = voxelindex_to_gradindex[voxel_index];
402  int count = voxelindex_to_accpoint[voxel_index].Count();
403  Eigen::Map<const FeatureVec_t> grad(
404  pooled_features_gradient + in_channels * grad_idx,
405  in_channels);
406  feat_bp = grad / count;
407  }
408  }
409 
410  if (FEAT_FN == NEAREST_NEIGHBOR) {
411  for (const auto point : voxelindex_to_accpoint) {
412  size_t idx = point.second.Index()(0);
413  Eigen::Map<FeatureVec_t> feat_bp(
414  features_backprop + in_channels * idx, in_channels);
415 
416  size_t grad_idx = voxelindex_to_gradindex[point.first];
417  Eigen::Map<const FeatureVec_t> grad(
418  pooled_features_gradient + in_channels * grad_idx,
419  in_channels);
420  feat_bp = grad;
421  }
422  }
423 
424  if (FEAT_FN == MAX) {
425  for (const auto point : voxelindex_to_accpoint) {
426  size_t grad_idx = voxelindex_to_gradindex[point.first];
427  Eigen::Map<const FeatureVec_t> grad(
428  pooled_features_gradient + in_channels * grad_idx,
429  in_channels);
430  for (int i = 0; i < in_channels; ++i) {
431  size_t idx = point.second.Index()(i);
432  Eigen::Map<FeatureVec_t> feat_bp(
433  features_backprop + in_channels * idx, in_channels);
434  feat_bp(i) = grad(i);
435  }
436  }
437  }
438 }
439 
484 template <class TReal, class TFeat, class OUTPUT_ALLOCATOR>
485 void VoxelPooling(size_t num_inp,
486  const TReal* const inp_positions,
487  int in_channels,
488  const TFeat* inp_features,
489  TReal voxel_size,
490  OUTPUT_ALLOCATOR& output_allocator,
491  AccumulationFn position_fn,
492  AccumulationFn feature_fn) {
493 #define CALL_TEMPLATE(POS_FN, FEAT_FN) \
494  if (POS_FN == position_fn && FEAT_FN == feature_fn) { \
495  _VoxelPooling<TReal, TFeat, \
496  Accumulator<TReal, TFeat, POS_FN, FEAT_FN>>( \
497  num_inp, inp_positions, in_channels, inp_features, voxel_size, \
498  output_allocator); \
499  }
500 
510 
511 #undef CALL_TEMPLATE
512 }
513 
529 template <class TReal, class TFeat>
530 void VoxelPoolingBackprop(TFeat* features_backprop,
531  size_t num_inp,
532  const TReal* const inp_positions,
533  int in_channels,
534  const TFeat* const inp_features,
535  size_t num_pooled,
536  const TReal* const pooled_positions,
537  const TFeat* const pooled_features_gradient,
538  TReal voxel_size,
539  AccumulationFn position_fn,
540  AccumulationFn feature_fn) {
541 #define CALL_TEMPLATE(POS_FN, FEAT_FN) \
542  if (POS_FN == position_fn && FEAT_FN == feature_fn) { \
543  _VoxelPoolingBackprop< \
544  TReal, TFeat, \
545  AccumulatorBackprop<TReal, TFeat, POS_FN, FEAT_FN>, FEAT_FN>( \
546  features_backprop, num_inp, inp_positions, in_channels, \
547  inp_features, num_pooled, pooled_positions, \
548  pooled_features_gradient, voxel_size); \
549  }
550 
560 
561 #undef CALL_TEMPLATE
562 }
563 
564 } // namespace impl
565 } // namespace ml
566 } // namespace open3d
#define CALL_TEMPLATE(POS_FN, FEAT_FN)
Definition: VoxelPooling.h:104
Eigen::Array< TFeat, Eigen::Dynamic, 1 > Features() const
Definition: VoxelPooling.h:180
void AddPoint(const Eigen::MatrixBase< Derived > &pos, const Eigen::MatrixBase< Derived2 > &voxel_center, const Eigen::ArrayBase< Derived3 > &feat, const size_t idx)
Definition: VoxelPooling.h:116
Eigen::Array< TReal, 3, 1 > Position() const
Definition: VoxelPooling.h:171
Eigen::Array< size_t, Eigen::Dynamic, 1 > Index() const
Definition: VoxelPooling.h:191
int Count() const
Definition: VoxelPooling.h:189
AccumulatorBackprop()
Definition: VoxelPooling.h:106
Definition: VoxelPooling.h:27
int Count() const
Definition: VoxelPooling.h:91
void AddPoint(const Eigen::MatrixBase< Derived > &pos, const Eigen::MatrixBase< Derived2 > &voxel_center, const Eigen::ArrayBase< Derived3 > &feat)
Definition: VoxelPooling.h:39
Eigen::Array< TReal, 3, 1 > Position() const
Definition: VoxelPooling.h:73
Accumulator()
Definition: VoxelPooling.h:29
Eigen::Array< TFeat, Eigen::Dynamic, 1 > Features() const
Definition: VoxelPooling.h:82
int count
Definition: FilePCD.cpp:42
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 image_handle timestamp_usec white_balance image_handle k4a_device_configuration_t config device_handle char size_t serial_number_size bool int32_t int32_t int32_t int32_t k4a_color_control_mode_t default_mode value const const k4a_calibration_t calibration char size_t
Definition: K4aPlugin.cpp:719
AccumulationFn
Definition: VoxelPooling.h:21
@ CENTER
Definition: VoxelPooling.h:21
@ NEAREST_NEIGHBOR
Definition: VoxelPooling.h:21
@ MAX
Definition: VoxelPooling.h:21
@ AVERAGE
Definition: VoxelPooling.h:21
void VoxelPoolingBackprop(TFeat *features_backprop, size_t num_inp, const TReal *const inp_positions, int in_channels, const TFeat *const inp_features, size_t num_pooled, const TReal *const pooled_positions, const TFeat *const pooled_features_gradient, TReal voxel_size, AccumulationFn position_fn, AccumulationFn feature_fn)
Definition: VoxelPooling.h:530
bool CheckVoxelSize(std::string &err, const size_t num_positions, const T *const positions, const T voxel_size)
Function for debugging. Checks if the voxel size is too small.
Definition: VoxelPooling.h:205
void VoxelPooling(size_t num_inp, const TReal *const inp_positions, int in_channels, const TFeat *inp_features, TReal voxel_size, OUTPUT_ALLOCATOR &output_allocator, AccumulationFn position_fn, AccumulationFn feature_fn)
Definition: VoxelPooling.h:485
void _VoxelPoolingBackprop(TFeat *features_backprop, size_t num_inp, const TReal *const inp_positions, int in_channels, const TFeat *const inp_features, size_t num_pooled, const TReal *const pooled_positions, const TFeat *const pooled_features_gradient, TReal voxel_size)
Definition: VoxelPooling.h:325
void _VoxelPooling(size_t num_inp, const TReal *const inp_positions, int in_channels, const TFeat *inp_features, TReal voxel_size, OUTPUT_ALLOCATOR &output_allocator)
Definition: VoxelPooling.h:262
HOST_DEVICE utility::MiniVec< int, 3 > ComputeVoxelIndex(const TVecf &pos, const typename TVecf::Scalar_t &inv_voxel_size)
Definition: NeighborSearchCommon.h:42
Definition: PinholeCameraIntrinsic.cpp:16
Definition: Device.h:107
Definition: Helper.h:71