Open3D (C++ API)  0.17.0
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
16
17namespace open3d {
18namespace ml {
19namespace impl {
20
22
23template <class TReal,
24 class TFeat,
25 AccumulationFn POS_FN,
26 AccumulationFn FEAT_FN>
28public:
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
93private:
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
100template <class TReal,
101 class TFeat,
102 AccumulationFn POS_FN,
103 AccumulationFn FEAT_FN>
105public:
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
195private:
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
204template <class T>
205bool 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
248template <class TDerived>
249Eigen::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.
261template <class TReal, class TFeat, class ACCUMULATOR, class OUTPUT_ALLOCATOR>
262void _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.
324template <class TReal, class TFeat, class ACCUMULATOR, AccumulationFn FEAT_FN>
325void _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
484template <class TReal, class TFeat, class OUTPUT_ALLOCATOR>
485void 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
529template <class TReal, class TFeat>
530void 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< TReal, 3, 1 > Position() const
Definition: VoxelPooling.h:171
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< size_t, Eigen::Dynamic, 1 > Index() const
Definition: VoxelPooling.h:191
int Count() const
Definition: VoxelPooling.h:189
AccumulatorBackprop()
Definition: VoxelPooling.h:106
Eigen::Array< TFeat, Eigen::Dynamic, 1 > Features() const
Definition: VoxelPooling.h:180
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
Eigen::Array< TFeat, Eigen::Dynamic, 1 > Features() const
Definition: VoxelPooling.h:82
Accumulator()
Definition: VoxelPooling.h:29
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