567 lines
21 KiB
C
567 lines
21 KiB
C
|
// ----------------------------------------------------------------------------
|
||
|
|
// - Open3D: www.open3d.org -
|
||
|
|
// ----------------------------------------------------------------------------
|
||
|
|
// Copyright (c) 2018-2023 www.open3d.org
|
||
|
|
// SPDX-License-Identifier: MIT
|
||
|
|
// ----------------------------------------------------------------------------
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
#include <tbb/task_group.h>
|
||
|
|
|
||
|
|
#include <Eigen/Core>
|
||
|
|
#include <unordered_map>
|
||
|
|
|
||
|
|
#include "open3d/utility/Helper.h"
|
||
|
|
|
||
|
|
namespace open3d {
|
||
|
|
namespace ml {
|
||
|
|
namespace impl {
|
||
|
|
|
||
|
|
enum AccumulationFn { AVERAGE = 0, NEAREST_NEIGHBOR, MAX, CENTER };
|
||
|
|
|
||
|
|
template <class TReal,
|
||
|
|
class TFeat,
|
||
|
|
AccumulationFn POS_FN,
|
||
|
|
AccumulationFn FEAT_FN>
|
||
|
|
class Accumulator {
|
||
|
|
public:
|
||
|
|
Accumulator()
|
||
|
|
: count_(0),
|
||
|
|
min_sqr_dist_to_center_(std::numeric_limits<TReal>::max()),
|
||
|
|
position_(0, 0, 0) {
|
||
|
|
static_assert(POS_FN != MAX, "MAX is not allowed for point positions");
|
||
|
|
static_assert(FEAT_FN != CENTER,
|
||
|
|
"CENTER is not allowed for feature vectors");
|
||
|
|
}
|
||
|
|
|
||
|
|
template <class Derived, class Derived2, class Derived3>
|
||
|
|
inline void AddPoint(const Eigen::MatrixBase<Derived>& pos,
|
||
|
|
const Eigen::MatrixBase<Derived2>& voxel_center,
|
||
|
|
const Eigen::ArrayBase<Derived3>& feat) {
|
||
|
|
bool new_nearest_neighbor = false;
|
||
|
|
TReal sqr_d = 0;
|
||
|
|
if (POS_FN == NEAREST_NEIGHBOR || FEAT_FN == NEAREST_NEIGHBOR) {
|
||
|
|
sqr_d = (voxel_center - pos).squaredNorm();
|
||
|
|
if (sqr_d < min_sqr_dist_to_center_) {
|
||
|
|
new_nearest_neighbor = true;
|
||
|
|
min_sqr_dist_to_center_ = sqr_d;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (POS_FN == AVERAGE) {
|
||
|
|
position_ += pos.array();
|
||
|
|
} else if (POS_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
|
||
|
|
position_ = pos;
|
||
|
|
} else if (POS_FN == CENTER) {
|
||
|
|
if (count_ == 0) position_ = voxel_center;
|
||
|
|
}
|
||
|
|
|
||
|
|
if (count_ == 0) {
|
||
|
|
features_.resizeLike(feat);
|
||
|
|
features_.setZero();
|
||
|
|
}
|
||
|
|
if (FEAT_FN == AVERAGE) {
|
||
|
|
features_ += feat;
|
||
|
|
} else if (FEAT_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
|
||
|
|
features_ = feat;
|
||
|
|
} else if (FEAT_FN == MAX) {
|
||
|
|
features_ = features_.max(feat);
|
||
|
|
}
|
||
|
|
++count_;
|
||
|
|
}
|
||
|
|
|
||
|
|
inline Eigen::Array<TReal, 3, 1> Position() const {
|
||
|
|
if (POS_FN == AVERAGE) {
|
||
|
|
return position_ / count_;
|
||
|
|
} else // if( POS_FN == NEAREST_NEIGHBOR || POS_FN == CENTER )
|
||
|
|
{
|
||
|
|
return position_;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline Eigen::Array<TFeat, Eigen::Dynamic, 1> Features() const {
|
||
|
|
if (FEAT_FN == AVERAGE) {
|
||
|
|
return features_ / count_;
|
||
|
|
} else // if( FEAT_FN == NEAREST_NEIGHBOR || FEAT_FN == MAX )
|
||
|
|
{
|
||
|
|
return features_;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline int Count() const { return count_; }
|
||
|
|
|
||
|
|
private:
|
||
|
|
int count_;
|
||
|
|
TReal min_sqr_dist_to_center_;
|
||
|
|
Eigen::Array<TReal, 3, 1> position_;
|
||
|
|
Eigen::Array<TFeat, Eigen::Dynamic, 1> features_;
|
||
|
|
}; // namespace
|
||
|
|
|
||
|
|
template <class TReal,
|
||
|
|
class TFeat,
|
||
|
|
AccumulationFn POS_FN,
|
||
|
|
AccumulationFn FEAT_FN>
|
||
|
|
class AccumulatorBackprop {
|
||
|
|
public:
|
||
|
|
AccumulatorBackprop()
|
||
|
|
: count_(0),
|
||
|
|
min_sqr_dist_to_center_(std::numeric_limits<TReal>::max()),
|
||
|
|
position_(0, 0, 0) {
|
||
|
|
static_assert(POS_FN != MAX, "MAX is not allowed for point positions");
|
||
|
|
static_assert(FEAT_FN != CENTER,
|
||
|
|
"CENTER is not allowed for feature vectors");
|
||
|
|
}
|
||
|
|
|
||
|
|
template <class Derived, class Derived2, class Derived3>
|
||
|
|
inline void AddPoint(const Eigen::MatrixBase<Derived>& pos,
|
||
|
|
const Eigen::MatrixBase<Derived2>& voxel_center,
|
||
|
|
const Eigen::ArrayBase<Derived3>& feat,
|
||
|
|
const size_t idx) {
|
||
|
|
bool new_nearest_neighbor = false;
|
||
|
|
TReal sqr_d = 0;
|
||
|
|
if (POS_FN == NEAREST_NEIGHBOR || FEAT_FN == NEAREST_NEIGHBOR) {
|
||
|
|
sqr_d = (voxel_center - pos).squaredNorm();
|
||
|
|
if (sqr_d < min_sqr_dist_to_center_) {
|
||
|
|
new_nearest_neighbor = true;
|
||
|
|
min_sqr_dist_to_center_ = sqr_d;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if (POS_FN == AVERAGE) {
|
||
|
|
position_ += pos.array();
|
||
|
|
} else if (POS_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
|
||
|
|
position_ = pos;
|
||
|
|
} else if (POS_FN == CENTER) {
|
||
|
|
if (count_ == 0) position_ = voxel_center;
|
||
|
|
}
|
||
|
|
|
||
|
|
if (count_ == 0) {
|
||
|
|
features_.resizeLike(feat);
|
||
|
|
features_.setZero();
|
||
|
|
if (FEAT_FN == NEAREST_NEIGHBOR) {
|
||
|
|
features_ = feat;
|
||
|
|
index_.resize(1);
|
||
|
|
index_(0) = idx;
|
||
|
|
++count_;
|
||
|
|
return;
|
||
|
|
} else if (FEAT_FN == MAX) {
|
||
|
|
features_ = feat;
|
||
|
|
index_.resizeLike(feat);
|
||
|
|
index_ = idx;
|
||
|
|
++count_;
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if (FEAT_FN == AVERAGE) {
|
||
|
|
features_ += feat;
|
||
|
|
} else if (FEAT_FN == NEAREST_NEIGHBOR && new_nearest_neighbor) {
|
||
|
|
features_ = feat;
|
||
|
|
index_(0) = idx;
|
||
|
|
} else if (FEAT_FN == MAX) {
|
||
|
|
for (int i = 0; i < features_.rows(); ++i) {
|
||
|
|
if (feat(i) > features_(i)) {
|
||
|
|
features_(i) = feat(i);
|
||
|
|
index_(i) = idx;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
++count_;
|
||
|
|
}
|
||
|
|
|
||
|
|
inline Eigen::Array<TReal, 3, 1> Position() const {
|
||
|
|
if (POS_FN == AVERAGE) {
|
||
|
|
return position_ / count_;
|
||
|
|
} else // if( POS_FN == NEAREST_NEIGHBOR || POS_FN == CENTER )
|
||
|
|
{
|
||
|
|
return position_;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline Eigen::Array<TFeat, Eigen::Dynamic, 1> Features() const {
|
||
|
|
if (FEAT_FN == AVERAGE) {
|
||
|
|
return features_ / count_;
|
||
|
|
} else // if( FEAT_FN == NEAREST_NEIGHBOR || FEAT_FN == MAX )
|
||
|
|
{
|
||
|
|
return features_;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
inline int Count() const { return count_; }
|
||
|
|
|
||
|
|
inline Eigen::Array<size_t, Eigen::Dynamic, 1> Index() const {
|
||
|
|
return index_;
|
||
|
|
}
|
||
|
|
|
||
|
|
private:
|
||
|
|
int count_;
|
||
|
|
TReal min_sqr_dist_to_center_;
|
||
|
|
Eigen::Array<TReal, 3, 1> position_;
|
||
|
|
Eigen::Array<TFeat, Eigen::Dynamic, 1> features_;
|
||
|
|
Eigen::Array<size_t, Eigen::Dynamic, 1> index_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// Function for debugging. Checks if the voxel size is too small.
|
||
|
|
template <class T>
|
||
|
|
bool CheckVoxelSize(std::string& err,
|
||
|
|
const size_t num_positions,
|
||
|
|
const T* const positions,
|
||
|
|
const T voxel_size) {
|
||
|
|
typedef Eigen::Array<double, 3, 1> Vec3_t;
|
||
|
|
if (num_positions == 0) {
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
Vec3_t bb_min, bb_max;
|
||
|
|
bb_min << positions[0], positions[1], positions[2];
|
||
|
|
bb_max = bb_min;
|
||
|
|
|
||
|
|
Vec3_t voxel_size3(voxel_size, voxel_size, voxel_size);
|
||
|
|
|
||
|
|
for (size_t i = 1; i < num_positions; ++i) {
|
||
|
|
Vec3_t pos(positions[i * 3 + 0], positions[i * 3 + 1],
|
||
|
|
positions[i * 3 + 2]);
|
||
|
|
bb_min = bb_min.min(pos);
|
||
|
|
bb_max = bb_max.max(pos);
|
||
|
|
}
|
||
|
|
|
||
|
|
// the min and max bounding box shall be a multiple of the voxel size
|
||
|
|
bb_min /= voxel_size3;
|
||
|
|
bb_min = bb_min.floor() * voxel_size3;
|
||
|
|
bb_max /= voxel_size3;
|
||
|
|
bb_max = bb_max.ceil() * voxel_size3;
|
||
|
|
|
||
|
|
if (voxel_size * double(std::numeric_limits<int>::max()) <
|
||
|
|
bb_max.maxCoeff() ||
|
||
|
|
voxel_size * double(std::numeric_limits<int>::min()) >
|
||
|
|
bb_min.maxCoeff()) {
|
||
|
|
err = "voxel_size is too small\n";
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Computes an integer voxel index for a 3D position.
|
||
|
|
///
|
||
|
|
/// \param pos A 3D position.
|
||
|
|
/// \param inv_voxel_size The reciprocal of the voxel size
|
||
|
|
///
|
||
|
|
template <class TDerived>
|
||
|
|
Eigen::Vector3i ComputeVoxelIndex(
|
||
|
|
const Eigen::ArrayBase<TDerived>& pos,
|
||
|
|
const typename TDerived::Scalar& inv_voxel_size) {
|
||
|
|
typedef typename TDerived::Scalar Scalar_t;
|
||
|
|
Eigen::Array<Scalar_t, 3, 1> ref_coord = pos * inv_voxel_size;
|
||
|
|
|
||
|
|
Eigen::Vector3i voxel_index;
|
||
|
|
voxel_index = ref_coord.floor().template cast<int>();
|
||
|
|
return voxel_index;
|
||
|
|
}
|
||
|
|
|
||
|
|
// implementation for VoxelPooling with template parameter for the accumulator.
|
||
|
|
template <class TReal, class TFeat, class ACCUMULATOR, class OUTPUT_ALLOCATOR>
|
||
|
|
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) {
|
||
|
|
if (num_inp == 0) {
|
||
|
|
TReal* out_pos_ptr;
|
||
|
|
TFeat* out_feat_ptr;
|
||
|
|
output_allocator.AllocPooledPositions(&out_pos_ptr, 0);
|
||
|
|
output_allocator.AllocPooledFeatures(&out_feat_ptr, 0, in_channels);
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
typedef Eigen::Array<TReal, 3, 1> Vec3_t;
|
||
|
|
typedef Eigen::Array<TFeat, Eigen::Dynamic, 1> FeatureVec_t;
|
||
|
|
|
||
|
|
std::unordered_map<Eigen::Vector3i, ACCUMULATOR,
|
||
|
|
open3d::utility::hash_eigen<Eigen::Vector3i>>
|
||
|
|
voxelindex_to_accpoint;
|
||
|
|
|
||
|
|
Vec3_t voxel_center;
|
||
|
|
Eigen::Vector3i voxel_index;
|
||
|
|
TReal inv_voxel_size = 1 / voxel_size;
|
||
|
|
TReal half_voxel_size = 0.5 * voxel_size;
|
||
|
|
for (size_t i = 0; i < num_inp; ++i) {
|
||
|
|
Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
|
||
|
|
|
||
|
|
voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
|
||
|
|
|
||
|
|
voxel_center << voxel_index(0) * voxel_size + half_voxel_size,
|
||
|
|
voxel_index(1) * voxel_size + half_voxel_size,
|
||
|
|
voxel_index(2) * voxel_size + half_voxel_size;
|
||
|
|
|
||
|
|
Eigen::Map<const FeatureVec_t> feat(inp_features + in_channels * i,
|
||
|
|
in_channels);
|
||
|
|
voxelindex_to_accpoint[voxel_index].AddPoint(
|
||
|
|
pos.matrix(), voxel_center.matrix(), feat);
|
||
|
|
}
|
||
|
|
|
||
|
|
const size_t num_out = voxelindex_to_accpoint.size();
|
||
|
|
|
||
|
|
TReal* out_pos_ptr;
|
||
|
|
TFeat* out_feat_ptr;
|
||
|
|
output_allocator.AllocPooledPositions(&out_pos_ptr, num_out);
|
||
|
|
output_allocator.AllocPooledFeatures(&out_feat_ptr, num_out, in_channels);
|
||
|
|
|
||
|
|
size_t i = 0;
|
||
|
|
for (const auto point : voxelindex_to_accpoint) {
|
||
|
|
Vec3_t pos = point.second.Position();
|
||
|
|
Eigen::Map<Vec3_t> out_pos(out_pos_ptr + i * 3);
|
||
|
|
out_pos = pos;
|
||
|
|
|
||
|
|
Eigen::Map<FeatureVec_t> out_feat(out_feat_ptr + i * in_channels,
|
||
|
|
in_channels);
|
||
|
|
out_feat = point.second.Features();
|
||
|
|
++i;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// implementation for VoxelPoolingBackprop with template parameter for the
|
||
|
|
// accumulator.
|
||
|
|
template <class TReal, class TFeat, class ACCUMULATOR, AccumulationFn FEAT_FN>
|
||
|
|
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) {
|
||
|
|
if (num_inp == 0) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
memset(features_backprop, 0, sizeof(TFeat) * num_inp * in_channels);
|
||
|
|
|
||
|
|
typedef Eigen::Array<TReal, 3, 1> Vec3_t;
|
||
|
|
typedef Eigen::Array<TFeat, Eigen::Dynamic, 1> FeatureVec_t;
|
||
|
|
|
||
|
|
Vec3_t voxel_size3(voxel_size, voxel_size, voxel_size);
|
||
|
|
|
||
|
|
// create the two hash maps in parallel
|
||
|
|
tbb::task_group task_group;
|
||
|
|
|
||
|
|
std::unordered_map<Eigen::Vector3i, ACCUMULATOR,
|
||
|
|
open3d::utility::hash_eigen<Eigen::Vector3i>>
|
||
|
|
voxelindex_to_accpoint;
|
||
|
|
|
||
|
|
task_group.run([&] {
|
||
|
|
Vec3_t voxel_center;
|
||
|
|
Eigen::Vector3i voxel_index;
|
||
|
|
TReal inv_voxel_size = 1 / voxel_size;
|
||
|
|
TReal half_voxel_size = 0.5 * voxel_size;
|
||
|
|
for (size_t i = 0; i < num_inp; ++i) {
|
||
|
|
Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
|
||
|
|
|
||
|
|
voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
|
||
|
|
|
||
|
|
voxel_center << voxel_index(0) * voxel_size + half_voxel_size,
|
||
|
|
voxel_index(1) * voxel_size + half_voxel_size,
|
||
|
|
voxel_index(2) * voxel_size + half_voxel_size;
|
||
|
|
|
||
|
|
Eigen::Map<const FeatureVec_t> feat(inp_features + in_channels * i,
|
||
|
|
in_channels);
|
||
|
|
voxelindex_to_accpoint[voxel_index].AddPoint(
|
||
|
|
pos.matrix(), voxel_center.matrix(), feat, i);
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
std::unordered_map<Eigen::Vector3i, size_t,
|
||
|
|
open3d::utility::hash_eigen<Eigen::Vector3i>>
|
||
|
|
voxelindex_to_gradindex;
|
||
|
|
|
||
|
|
task_group.run([&] {
|
||
|
|
Eigen::Vector3i voxel_index;
|
||
|
|
TReal inv_voxel_size = 1 / voxel_size;
|
||
|
|
for (size_t i = 0; i < num_pooled; ++i) {
|
||
|
|
Eigen::Map<const Vec3_t> pos(pooled_positions + i * 3);
|
||
|
|
|
||
|
|
voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
|
||
|
|
|
||
|
|
voxelindex_to_gradindex[voxel_index] = i;
|
||
|
|
}
|
||
|
|
});
|
||
|
|
|
||
|
|
task_group.wait();
|
||
|
|
|
||
|
|
if (FEAT_FN == AVERAGE) {
|
||
|
|
Eigen::Vector3i voxel_index;
|
||
|
|
TReal inv_voxel_size = 1 / voxel_size;
|
||
|
|
for (size_t i = 0; i < num_inp; ++i) {
|
||
|
|
Eigen::Map<const Vec3_t> pos(inp_positions + i * 3);
|
||
|
|
|
||
|
|
voxel_index = ComputeVoxelIndex(pos, inv_voxel_size);
|
||
|
|
|
||
|
|
Eigen::Map<FeatureVec_t> feat_bp(
|
||
|
|
features_backprop + in_channels * i, in_channels);
|
||
|
|
|
||
|
|
size_t grad_idx = voxelindex_to_gradindex[voxel_index];
|
||
|
|
int count = voxelindex_to_accpoint[voxel_index].Count();
|
||
|
|
Eigen::Map<const FeatureVec_t> grad(
|
||
|
|
pooled_features_gradient + in_channels * grad_idx,
|
||
|
|
in_channels);
|
||
|
|
feat_bp = grad / count;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if (FEAT_FN == NEAREST_NEIGHBOR) {
|
||
|
|
for (const auto point : voxelindex_to_accpoint) {
|
||
|
|
size_t idx = point.second.Index()(0);
|
||
|
|
Eigen::Map<FeatureVec_t> feat_bp(
|
||
|
|
features_backprop + in_channels * idx, in_channels);
|
||
|
|
|
||
|
|
size_t grad_idx = voxelindex_to_gradindex[point.first];
|
||
|
|
Eigen::Map<const FeatureVec_t> grad(
|
||
|
|
pooled_features_gradient + in_channels * grad_idx,
|
||
|
|
in_channels);
|
||
|
|
feat_bp = grad;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if (FEAT_FN == MAX) {
|
||
|
|
for (const auto point : voxelindex_to_accpoint) {
|
||
|
|
size_t grad_idx = voxelindex_to_gradindex[point.first];
|
||
|
|
Eigen::Map<const FeatureVec_t> grad(
|
||
|
|
pooled_features_gradient + in_channels * grad_idx,
|
||
|
|
in_channels);
|
||
|
|
for (int i = 0; i < in_channels; ++i) {
|
||
|
|
size_t idx = point.second.Index()(i);
|
||
|
|
Eigen::Map<FeatureVec_t> feat_bp(
|
||
|
|
features_backprop + in_channels * idx, in_channels);
|
||
|
|
feat_bp(i) = grad(i);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Pooling operation for point clouds. Aggregates points that are inside the
|
||
|
|
/// same voxel.
|
||
|
|
///
|
||
|
|
/// \tparam TReal Scalar type for point positions.
|
||
|
|
///
|
||
|
|
/// \tparam TFeat Scalar type for the features.
|
||
|
|
///
|
||
|
|
/// \tparam OUTPUT_ALLOCATOR Type of the output_allocator. See
|
||
|
|
/// \p output_allocator for more information.
|
||
|
|
///
|
||
|
|
/// \param num_inp The number of input points.
|
||
|
|
///
|
||
|
|
/// \param inp_positions Array with 3D point positions.
|
||
|
|
///
|
||
|
|
/// \param in_channels The number of input feature channels.
|
||
|
|
///
|
||
|
|
/// \param inp_features The array with the point features. The shape is
|
||
|
|
/// [num_inp, in_channels].
|
||
|
|
///
|
||
|
|
/// \param voxel_size The voxel size (voxel edge length) used for pooling.
|
||
|
|
///
|
||
|
|
/// \param output_allocator An object that implements functions for
|
||
|
|
/// allocating the output arrays. The object must implement functions
|
||
|
|
/// AllocPooledPositions(TReal** ptr, size_t num) and
|
||
|
|
/// AllocPooledFeatures(TFeat** ptr, size_t num, channels), with 'num'
|
||
|
|
/// as the number of output points and 'channels' as \p in_channels .
|
||
|
|
/// Both functions should allocate memory and return a pointer
|
||
|
|
/// to that memory in ptr.
|
||
|
|
/// Both functions must accept the argument num==0. In this case ptr
|
||
|
|
/// does not need to be set. Note that AllocPooledPositions must
|
||
|
|
/// allocate memory for num*3*sizeof(TReal) bytes.
|
||
|
|
///
|
||
|
|
/// \param position_fn Defines how the new point positions will be computed.
|
||
|
|
/// AVERAGE computes the center of gravity for the points within one
|
||
|
|
/// voxel.
|
||
|
|
/// NEAREST_NEIGHBOR selects the point closest to the voxel center.
|
||
|
|
/// CENTER uses the voxel center for the position of the generated point.
|
||
|
|
///
|
||
|
|
/// \param feature_fn Defines how the new features will be computed.
|
||
|
|
/// AVERAGE computes the average feature vector.
|
||
|
|
/// NEAREST_NEIGHBOR selects the feature vector of the point closest to
|
||
|
|
/// the voxel center.
|
||
|
|
/// MAX uses the maximum feature among all points within the voxel.
|
||
|
|
///
|
||
|
|
template <class TReal, class TFeat, class OUTPUT_ALLOCATOR>
|
||
|
|
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) {
|
||
|
|
#define CALL_TEMPLATE(POS_FN, FEAT_FN) \
|
||
|
|
if (POS_FN == position_fn && FEAT_FN == feature_fn) { \
|
||
|
|
_VoxelPooling<TReal, TFeat, \
|
||
|
|
Accumulator<TReal, TFeat, POS_FN, FEAT_FN>>( \
|
||
|
|
num_inp, inp_positions, in_channels, inp_features, voxel_size, \
|
||
|
|
output_allocator); \
|
||
|
|
}
|
||
|
|
|
||
|
|
CALL_TEMPLATE(AVERAGE, AVERAGE)
|
||
|
|
CALL_TEMPLATE(AVERAGE, NEAREST_NEIGHBOR)
|
||
|
|
CALL_TEMPLATE(AVERAGE, MAX)
|
||
|
|
CALL_TEMPLATE(NEAREST_NEIGHBOR, AVERAGE)
|
||
|
|
CALL_TEMPLATE(NEAREST_NEIGHBOR, NEAREST_NEIGHBOR)
|
||
|
|
CALL_TEMPLATE(NEAREST_NEIGHBOR, MAX)
|
||
|
|
CALL_TEMPLATE(CENTER, AVERAGE)
|
||
|
|
CALL_TEMPLATE(CENTER, NEAREST_NEIGHBOR)
|
||
|
|
CALL_TEMPLATE(CENTER, MAX)
|
||
|
|
|
||
|
|
#undef CALL_TEMPLATE
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Backpropagation to the features for VoxelPooling.
|
||
|
|
///
|
||
|
|
/// \param features_backprop The output array with the gradients for the
|
||
|
|
/// features.
|
||
|
|
///
|
||
|
|
/// \param num_pooled The number of points after pooling with \p
|
||
|
|
/// VoxelPooling.
|
||
|
|
///
|
||
|
|
/// \param pooled_positions Array with the 3D positions after pooling.
|
||
|
|
///
|
||
|
|
/// \param pooled_features_gradient Array with the point features after
|
||
|
|
/// pooling.
|
||
|
|
///
|
||
|
|
/// See \p VoxelPooling for the description of the remaining parameters.
|
||
|
|
///
|
||
|
|
template <class TReal, class TFeat>
|
||
|
|
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) {
|
||
|
|
#define CALL_TEMPLATE(POS_FN, FEAT_FN) \
|
||
|
|
if (POS_FN == position_fn && FEAT_FN == feature_fn) { \
|
||
|
|
_VoxelPoolingBackprop< \
|
||
|
|
TReal, TFeat, \
|
||
|
|
AccumulatorBackprop<TReal, TFeat, POS_FN, FEAT_FN>, FEAT_FN>( \
|
||
|
|
features_backprop, num_inp, inp_positions, in_channels, \
|
||
|
|
inp_features, num_pooled, pooled_positions, \
|
||
|
|
pooled_features_gradient, voxel_size); \
|
||
|
|
}
|
||
|
|
|
||
|
|
CALL_TEMPLATE(AVERAGE, AVERAGE)
|
||
|
|
CALL_TEMPLATE(AVERAGE, NEAREST_NEIGHBOR)
|
||
|
|
CALL_TEMPLATE(AVERAGE, MAX)
|
||
|
|
CALL_TEMPLATE(NEAREST_NEIGHBOR, AVERAGE)
|
||
|
|
CALL_TEMPLATE(NEAREST_NEIGHBOR, NEAREST_NEIGHBOR)
|
||
|
|
CALL_TEMPLATE(NEAREST_NEIGHBOR, MAX)
|
||
|
|
CALL_TEMPLATE(CENTER, AVERAGE)
|
||
|
|
CALL_TEMPLATE(CENTER, NEAREST_NEIGHBOR)
|
||
|
|
CALL_TEMPLATE(CENTER, MAX)
|
||
|
|
|
||
|
|
#undef CALL_TEMPLATE
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace impl
|
||
|
|
} // namespace ml
|
||
|
|
} // namespace open3d
|