VoxelPooling.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 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_;
};

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;
}

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);
            }
        }
    }
}

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
}

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