integrated_icp_factor_impl.hpp

// SPDX-License-Identifier: MIT
// Copyright (c) 2021  Kenji Koide (k.koide@aist.go.jp)
 
#include <gtsam_points/factors/integrated_icp_factor.hpp>
 
#include <gtsam/geometry/Pose3.h>
#include <gtsam/linear/HessianFactor.h>
#include <gtsam_points/config.hpp>
#include <gtsam_points/ann/kdtree2.hpp>
#include <gtsam_points/types/frame_traits.hpp>
#include <gtsam_points/util/parallelism.hpp>
#include <gtsam_points/factors/impl/scan_matching_reduction.hpp>
 
#ifdef GTSAM_POINTS_USE_TBB
#include <tbb/parallel_for.h>
#endif
 
namespace gtsam_points {
 
template <typename TargetFrame, typename SourceFrame>
IntegratedICPFactor_<TargetFrame, SourceFrame>::IntegratedICPFactor_(
  gtsam::Key target_key,
  gtsam::Key source_key,
  const std::shared_ptr<const TargetFrame>& target,
  const std::shared_ptr<const SourceFrame>& source,
  const std::shared_ptr<const NearestNeighborSearch>& target_tree,
  bool use_point_to_plane)
: gtsam_points::IntegratedMatchingCostFactor(target_key, source_key),
  num_threads(1),
  max_correspondence_distance_sq(1.0),
  use_point_to_plane(use_point_to_plane),
  correspondence_update_tolerance_rot(0.0),
  correspondence_update_tolerance_trans(0.0),
  target(target),
  source(source) {
  //
  if (!frame::has_points(*target) || (use_point_to_plane && !frame::has_normals(*target))) {
    std::cerr << "error: target frame doesn't have required attributes for icp" << std::endl;
    abort();
  }
 
  if (!frame::has_points(*source)) {
    std::cerr << "error: source frame doesn't have required attributes for icp" << std::endl;
    abort();
  }
 
  if (target_tree) {
    this->target_tree = target_tree;
  } else {
    this->target_tree.reset(new KdTree2<TargetFrame>(target));
  }
}
 
template <typename TargetFrame, typename SourceFrame>
IntegratedICPFactor_<TargetFrame, SourceFrame>::IntegratedICPFactor_(
  gtsam::Key target_key,
  gtsam::Key source_key,
  const std::shared_ptr<const TargetFrame>& target,
  const std::shared_ptr<const SourceFrame>& source,
  bool use_point_to_plane)
: gtsam_points::IntegratedICPFactor_<TargetFrame, SourceFrame>(target_key, source_key, target, source, nullptr, use_point_to_plane) {}
 
template <typename TargetFrame, typename SourceFrame>
IntegratedICPFactor_<TargetFrame, SourceFrame>::IntegratedICPFactor_(
  const gtsam::Pose3& fixed_target_pose,
  gtsam::Key source_key,
  const std::shared_ptr<const TargetFrame>& target,
  const std::shared_ptr<const SourceFrame>& source,
  const std::shared_ptr<const NearestNeighborSearch>& target_tree,
  bool use_point_to_plane)
: gtsam_points::IntegratedMatchingCostFactor(fixed_target_pose, source_key),
  num_threads(1),
  max_correspondence_distance_sq(1.0),
  use_point_to_plane(use_point_to_plane),
  correspondence_update_tolerance_rot(0.0),
  correspondence_update_tolerance_trans(0.0),
  target(target),
  source(source) {
  //
  if (!frame::has_points(*target) || (use_point_to_plane && !frame::has_normals(*target))) {
    std::cerr << "error: target frame doesn't have required attributes for icp" << std::endl;
    abort();
  }
 
  if (!frame::has_points(*source)) {
    std::cerr << "error: source frame doesn't have required attributes for icp" << std::endl;
    abort();
  }
 
  if (target_tree) {
    this->target_tree = target_tree;
  } else {
    this->target_tree.reset(new KdTree2<TargetFrame>(target));
  }
}
 
template <typename TargetFrame, typename SourceFrame>
IntegratedICPFactor_<TargetFrame, SourceFrame>::IntegratedICPFactor_(
  const gtsam::Pose3& fixed_target_pose,
  gtsam::Key source_key,
  const std::shared_ptr<const TargetFrame>& target,
  const std::shared_ptr<const SourceFrame>& source,
  bool use_point_to_plane)
: gtsam_points::IntegratedICPFactor_<TargetFrame, SourceFrame>(fixed_target_pose, source_key, target, source, nullptr, use_point_to_plane) {}
 
template <typename TargetFrame, typename SourceFrame>
IntegratedICPFactor_<TargetFrame, SourceFrame>::~IntegratedICPFactor_() {}
 
template <typename TargetFrame, typename SourceFrame>
void IntegratedICPFactor_<TargetFrame, SourceFrame>::print(const std::string& s, const gtsam::KeyFormatter& keyFormatter) const {
  std::cout << s << "IntegratedICPFactor";
  if (is_binary) {
    std::cout << "(" << keyFormatter(this->keys()[0]) << ", " << keyFormatter(this->keys()[1]) << ")" << std::endl;
  } else {
    std::cout << "(fixed, " << keyFormatter(this->keys()[0]) << ")" << std::endl;
  }
 
  std::cout << "|target|=" << frame::size(*target) << "pts, |source|=" << frame::size(*source) << "pts" << std::endl;
  std::cout << "num_threads=" << num_threads << ", max_corr_dist=" << std::sqrt(max_correspondence_distance_sq) << std::endl;
}
 
template <typename TargetFrame, typename SourceFrame>
size_t IntegratedICPFactor_<TargetFrame, SourceFrame>::memory_usage() const {
  return sizeof(*this) + sizeof(long) * correspondences.capacity();
}
 
template <typename TargetFrame, typename SourceFrame>
void IntegratedICPFactor_<TargetFrame, SourceFrame>::update_correspondences(const Eigen::Isometry3d& delta) const {
  bool do_update = true;
  if (correspondences.size() == frame::size(*source) && (correspondence_update_tolerance_trans > 0.0 || correspondence_update_tolerance_rot > 0.0)) {
    Eigen::Isometry3d diff = delta.inverse() * last_correspondence_point;
    double diff_rot = Eigen::AngleAxisd(diff.linear()).angle();
    double diff_trans = diff.translation().norm();
    if (diff_rot < correspondence_update_tolerance_rot && diff_trans < correspondence_update_tolerance_trans) {
      do_update = false;
    }
  }
 
  if (!do_update) {
    return;
  }
 
  last_correspondence_point = delta;
  correspondences.resize(frame::size(*source));
 
  const auto perpoint_task = [&](int i) {
    Eigen::Vector4d pt = delta * frame::point(*source, i);
 
    size_t k_index = -1;
    double k_sq_dist = -1;
    size_t num_found = target_tree->knn_search(pt.data(), 1, &k_index, &k_sq_dist, max_correspondence_distance_sq);
 
    if (num_found == 0 || k_sq_dist > max_correspondence_distance_sq) {
      correspondences[i] = -1;
    } else {
      correspondences[i] = k_index;
    }
  };
 
  if (is_omp_default() || num_threads == 1) {
#pragma omp parallel for num_threads(num_threads) schedule(guided, 8)
    for (int i = 0; i < frame::size(*source); i++) {
      perpoint_task(i);
    }
  } else {
#ifdef GTSAM_POINTS_USE_TBB
    tbb::parallel_for(tbb::blocked_range<int>(0, frame::size(*source), 8), [&](const tbb::blocked_range<int>& range) {
      for (int i = range.begin(); i < range.end(); i++) {
        perpoint_task(i);
      }
    });
#else
    std::cerr << "error: TBB is not available" << std::endl;
    abort();
#endif
  }
}
 
template <typename TargetFrame, typename SourceFrame>
double IntegratedICPFactor_<TargetFrame, SourceFrame>::evaluate(
  const Eigen::Isometry3d& delta,
  Eigen::Matrix<double, 6, 6>* H_target,
  Eigen::Matrix<double, 6, 6>* H_source,
  Eigen::Matrix<double, 6, 6>* H_target_source,
  Eigen::Matrix<double, 6, 1>* b_target,
  Eigen::Matrix<double, 6, 1>* b_source) const {
  //
  if (correspondences.size() != frame::size(*source)) {
    update_correspondences(delta);
  }
 
  const auto perpoint_task = [&](
                               int i,
                               Eigen::Matrix<double, 6, 6>* H_target,
                               Eigen::Matrix<double, 6, 6>* H_source,
                               Eigen::Matrix<double, 6, 6>* H_target_source,
                               Eigen::Matrix<double, 6, 1>* b_target,
                               Eigen::Matrix<double, 6, 1>* b_source) {
    long target_index = correspondences[i];
    if (target_index < 0) {
      return 0.0;
    }
 
    const auto& mean_A = frame::point(*source, i);
    const auto& mean_B = frame::point(*target, target_index);
 
    const Eigen::Vector4d transed_mean_A = delta * mean_A;
    Eigen::Vector4d residual = mean_B - transed_mean_A;
 
    if (use_point_to_plane) {
      const auto& normal_B = frame::normal(*target, target_index);
      residual = normal_B.array() * residual.array();
    }
 
    const double error = residual.transpose() * residual;
    if (H_target == nullptr) {
      return error;
    }
 
    Eigen::Matrix<double, 4, 6> J_target = Eigen::Matrix<double, 4, 6>::Zero();
    J_target.block<3, 3>(0, 0) = -gtsam::SO3::Hat(transed_mean_A.head<3>());
    J_target.block<3, 3>(0, 3) = Eigen::Matrix3d::Identity();
 
    Eigen::Matrix<double, 4, 6> J_source = Eigen::Matrix<double, 4, 6>::Zero();
    J_source.block<3, 3>(0, 0) = delta.linear() * gtsam::SO3::Hat(mean_A.template head<3>());
    J_source.block<3, 3>(0, 3) = -delta.linear();
 
    if (use_point_to_plane) {
      const auto& normal_B = frame::normal(*target, target_index);
      J_target = normal_B.asDiagonal() * J_target;
      J_source = normal_B.asDiagonal() * J_source;
    }
 
    *H_target += J_target.transpose() * J_target;
    *H_source += J_source.transpose() * J_source;
    *H_target_source += J_target.transpose() * J_source;
    *b_target += J_target.transpose() * residual;
    *b_source += J_source.transpose() * residual;
 
    return error;
  };
 
  if (is_omp_default() || num_threads == 1) {
    return scan_matching_reduce_omp(perpoint_task, frame::size(*source), num_threads, H_target, H_source, H_target_source, b_target, b_source);
  } else {
    return scan_matching_reduce_tbb(perpoint_task, frame::size(*source), H_target, H_source, H_target_source, b_target, b_source);
  }
}
 
}  // namespace gtsam_points