flat_container.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright 2024 Kenji Koide
// SPDX-License-Identifier: MIT
#pragma once
 
#include <queue>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <small_gicp/ann/traits.hpp>
#include <small_gicp/points/traits.hpp>
#include <small_gicp/ann/knn_result.hpp>
 
namespace small_gicp {
 
template <bool HasNormals = false, bool HasCovs = false>
struct FlatContainer {
public:
  struct Setting {
    double min_sq_dist_in_cell = 0.1 * 0.1;  
    size_t max_num_points_in_cell = 10;      
  };
 
  FlatContainer() { points.reserve(5); }
 
  size_t size() const { return points.size(); }
 
  template <typename PointCloud>
  void add(const Setting& setting, const Eigen::Vector4d& transformed_pt, const PointCloud& points, size_t i, const Eigen::Isometry3d& T) {
    if (
      this->points.size() >= setting.max_num_points_in_cell ||                                                                                                  //
      std::any_of(this->points.begin(), this->points.end(), [&](const auto& pt) { return (pt - transformed_pt).squaredNorm() < setting.min_sq_dist_in_cell; })  //
    ) {
      return;
    }
 
    this->points.emplace_back(transformed_pt);
    if constexpr (HasNormals) {
      this->normals.emplace_back(T.matrix() * traits::normal(points, i));
    }
    if constexpr (HasCovs) {
      this->covs.emplace_back(T.matrix() * traits::cov(points, i) * T.matrix().transpose());
    }
  }
 
  void finalize() {}
 
  size_t nearest_neighbor_search(const Eigen::Vector4d& pt, size_t* k_index, double* k_sq_dist) const {
    if (points.empty()) {
      return 0;
    }
 
    KnnResult<1> result(k_index, k_sq_dist);
    knn_search(pt, result);
    return result.num_found();
  }
 
  size_t knn_search(const Eigen::Vector4d& pt, int k, size_t* k_indices, double* k_sq_dists) const {
    if (points.empty()) {
      return 0;
    }
 
    KnnResult<-1> result(k_indices, k_sq_dists, k);
    knn_search(pt, result);
    return result.num_found();
  }
 
  template <typename Result>
  void knn_search(const Eigen::Vector4d& pt, Result& result) const {
    if (points.empty()) {
      return;
    }
 
    for (size_t i = 0; i < points.size(); i++) {
      const double sq_dist = (points[i] - pt).squaredNorm();
      result.push(i, sq_dist);
    }
  }
 
public:
  struct Empty {};
 
  std::vector<Eigen::Vector4d> points;                                          
  std::conditional_t<HasNormals, std::vector<Eigen::Vector4d>, Empty> normals;  
  std::conditional_t<HasCovs, std::vector<Eigen::Matrix4d>, Empty> covs;        
};
 
using FlatContainerPoints = FlatContainer<false, false>;
using FlatContainerNormal = FlatContainer<true, false>;
using FlatContainerCov = FlatContainer<false, true>;
using FlatContainerNormalCov = FlatContainer<true, true>;
 
namespace traits {
 
template <bool HasNormals, bool HasCovs>
struct Traits<FlatContainer<HasNormals, HasCovs>> {
  static size_t size(const FlatContainer<HasNormals, HasCovs>& container) { return container.size(); }
  static bool has_points(const FlatContainer<HasNormals, HasCovs>& container) { return container.size(); }
  static bool has_normals(const FlatContainer<HasNormals, HasCovs>& container) { return HasNormals && container.size(); }
  static bool has_covs(const FlatContainer<HasNormals, HasCovs>& container) { return HasCovs && container.size(); }
 
  static const Eigen::Vector4d& point(const FlatContainer<HasNormals, HasCovs>& container, size_t i) { return container.points[i]; }
  static const Eigen::Vector4d& normal(const FlatContainer<HasNormals, HasCovs>& container, size_t i) { return container.normals[i]; }
  static const Eigen::Matrix4d& cov(const FlatContainer<HasNormals, HasCovs>& container, size_t i) { return container.covs[i]; }
 
  static size_t nearest_neighbor_search(const FlatContainer<HasNormals, HasCovs>& container, const Eigen::Vector4d& pt, size_t* k_index, double* k_sq_dist) {
    return container.nearest_neighbor_search(pt, k_index, k_sq_dist);
  }
 
  static size_t knn_search(const FlatContainer<HasNormals, HasCovs>& container, const Eigen::Vector4d& pt, size_t k, size_t* k_index, double* k_sq_dist) {
    return container.knn_search(pt, k, k_index, k_sq_dist);
  }
 
  template <typename Result>
  static void knn_search(const FlatContainer<HasNormals, HasCovs>& container, const Eigen::Vector4d& pt, Result& result) {
    container.knn_search(pt, result);
  }
};
 
}  // namespace traits
 
}  // namespace small_gicp