small_gicp/src/python/factors.cpp

107 lines
4.1 KiB
C++

// SPDX-FileCopyrightText: Copyright 2024 Kenji Koide
// SPDX-License-Identifier: MIT
#include <iostream>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <pybind11/eigen.h>
#include <pybind11/functional.h>
#include <small_gicp/factors/icp_factor.hpp>
#include <small_gicp/factors/plane_icp_factor.hpp>
#include <small_gicp/factors/gicp_factor.hpp>
#include <small_gicp/points/point_cloud.hpp>
#include <small_gicp/ann/kdtree_omp.hpp>
#include <small_gicp/registration/rejector.hpp>
namespace py = pybind11;
using namespace small_gicp;
void define_factors(py::module& m) {
// DistanceRejector
py::class_<DistanceRejector>(m, "DistanceRejector", "Correspondence rejection based on the distance between points.")
.def(py::init<>())
.def("set_max_distance", [](DistanceRejector& rejector, double dist) { rejector.max_dist_sq = dist * dist; }, py::arg("dist"), "Set the maximum distance.");
// ICPFactor
py::class_<ICPFactor>(m, "ICPFactor", "ICP per-point factor")
.def(py::init<>())
.def(
"linearize",
[](
ICPFactor& factor,
const PointCloud& target,
const PointCloud& source,
const KdTree<PointCloud>& kdtree,
const Eigen::Matrix4d& T,
size_t source_index,
const DistanceRejector& rejector) -> std::tuple<bool, Eigen::Matrix<double, 6, 6>, Eigen::Matrix<double, 6, 1>, double> {
Eigen::Matrix<double, 6, 6> H = Eigen::Matrix<double, 6, 6>::Zero();
Eigen::Matrix<double, 6, 1> b = Eigen::Matrix<double, 6, 1>::Zero();
double e = 0.0;
bool succ = factor.linearize(target, source, kdtree, Eigen::Isometry3d(T), source_index, rejector, &H, &b, &e);
return std::make_tuple(succ, H, b, e);
},
py::arg("target"),
py::arg("source"),
py::arg("kdtree"),
py::arg("T"),
py::arg("source_index"),
py::arg("rejector"),
"Linearize the factor. Returns a tuple of success, Hessian, gradient, and error.");
// PointToPlaneICPFactor
py::class_<PointToPlaneICPFactor>(m, "PointToPlaneICPFactor", "Point-to-plane ICP per-point factor")
.def(py::init<>())
.def(
"linearize",
[](
PointToPlaneICPFactor& factor,
const PointCloud& target,
const PointCloud& source,
const KdTree<PointCloud>& kdtree,
const Eigen::Matrix4d& T,
size_t source_index,
const DistanceRejector& rejector) -> std::tuple<bool, Eigen::Matrix<double, 6, 6>, Eigen::Matrix<double, 6, 1>, double> {
Eigen::Matrix<double, 6, 6> H = Eigen::Matrix<double, 6, 6>::Zero();
Eigen::Matrix<double, 6, 1> b = Eigen::Matrix<double, 6, 1>::Zero();
double e = 0.0;
bool succ = factor.linearize(target, source, kdtree, Eigen::Isometry3d(T), source_index, rejector, &H, &b, &e);
return std::make_tuple(succ, H, b, e);
},
py::arg("target"),
py::arg("source"),
py::arg("kdtree"),
py::arg("T"),
py::arg("source_index"),
py::arg("rejector"),
"Linearize the factor. Returns a tuple of success, Hessian, gradient, and error.");
// GICPFactor
py::class_<GICPFactor>(m, "GICPFactor", "Generalized ICP per-point factor") //
.def(py::init<>())
.def(
"linearize",
[](
GICPFactor& factor,
const PointCloud& target,
const PointCloud& source,
const KdTree<PointCloud>& kdtree,
const Eigen::Matrix4d& T,
size_t source_index,
const DistanceRejector& rejector) -> std::tuple<bool, Eigen::Matrix<double, 6, 6>, Eigen::Matrix<double, 6, 1>, double> {
Eigen::Matrix<double, 6, 6> H = Eigen::Matrix<double, 6, 6>::Zero();
Eigen::Matrix<double, 6, 1> b = Eigen::Matrix<double, 6, 1>::Zero();
double e = 0.0;
bool succ = factor.linearize(target, source, kdtree, Eigen::Isometry3d(T), source_index, rejector, &H, &b, &e);
return std::make_tuple(succ, H, b, e);
},
py::arg("target"),
py::arg("source"),
py::arg("kdtree"),
py::arg("T"),
py::arg("source_index"),
py::arg("rejector"),
"Linearize the factor. Returns a tuple of success, Hessian, gradient, and error.");
}