small_gicp/src/python/preprocess.cpp

304 lines
11 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/points/eigen.hpp>
#include <small_gicp/points/point_cloud.hpp>
#include <small_gicp/ann/kdtree_omp.hpp>
#include <small_gicp/util/downsampling.hpp>
#include <small_gicp/util/downsampling_omp.hpp>
#include <small_gicp/util/normal_estimation.hpp>
#include <small_gicp/util/normal_estimation_omp.hpp>
namespace py = pybind11;
using namespace small_gicp;
void define_preprocess(py::module& m) {
// voxelgrid_sampling
m.def(
"voxelgrid_sampling",
[](const PointCloud& points, double resolution, int num_threads) {
if (num_threads == 1) {
return voxelgrid_sampling(points, resolution);
}
return voxelgrid_sampling_omp(points, resolution, num_threads);
},
py::arg("points"),
py::arg("downsampling_resolution"),
py::arg("num_threads") = 1,
R"pbdoc(
Voxelgrid downsampling.
Notes
-----
When multi-threading is enabled, this function has minor run-by-run non-deterministic behavior due to parallel data collection that results
in a deviation of the number of points in the downsampling results (up to 10% increase from the single-thread version).
Discretized voxel coords must be in 21bit range [-1048576, 1048575].
For example, if the downsampling resolution is 0.01 m, point coordinates must be in [-10485.76, 10485.75] m.
Points outside the valid range will be ignored.
Parameters
----------
points : :class:`PointCloud`
Input point cloud.
resolution : float
Voxel size.
num_threads : int, optional
Number of threads. (default: 1)
Returns
-------
:class:`PointCloud`
Downsampled point cloud.
)pbdoc");
// voxelgrid_sampling (numpy)
m.def(
"voxelgrid_sampling",
[](const Eigen::MatrixXd& points, double resolution, int num_threads) -> std::shared_ptr<PointCloud> {
if (points.cols() != 3 && points.cols() != 4) {
std::cerr << "points must be Nx3 or Nx4" << std::endl;
return nullptr;
}
if (num_threads == 1) {
return voxelgrid_sampling<Eigen::MatrixXd, PointCloud>(points, resolution);
} else {
return voxelgrid_sampling_omp<Eigen::MatrixXd, PointCloud>(points, resolution, num_threads);
}
},
py::arg("points"),
py::arg("downsampling_resolution"),
py::arg("num_threads") = 1,
R"pbdoc(
Voxelgrid downsampling.
Notes
-----
When multi-threading is enabled, this function has minor run-by-run non-deterministic behavior due to parallel data collection that results
in a deviation of the number of points in the downsampling results (up to 10% increase from the single-thread version).
Discretized voxel coords must be in 21bit range [-1048576, 1048575].
For example, if the downsampling resolution is 0.01 m, point coordinates must be in [-10485.76, 10485.75] m.
Points outside the valid range will be ignored.
Parameters
----------
points : [np.float64]
Input point cloud. Nx3 or Nx4.
resolution : float
Voxel size.
num_threads : int, optional
Number of threads. (default: 1)
Returns
-------
:class:`PointCloud`
Downsampled point cloud.
)pbdoc");
// estimate_normals
m.def(
"estimate_normals",
[](PointCloud::Ptr points, std::shared_ptr<KdTree<PointCloud>> tree, int num_neighbors, int num_threads) -> std::shared_ptr<PointCloud> {
if (tree == nullptr) {
tree = std::make_shared<KdTree<PointCloud>>(points, KdTreeBuilderOMP(num_threads));
}
if (num_threads == 1) {
estimate_normals(*points, *tree, num_neighbors);
} else {
estimate_normals_omp(*points, *tree, num_neighbors, num_threads);
}
return points;
},
py::arg("points"),
py::arg("tree") = nullptr,
py::arg("num_neighbors") = 20,
py::arg("num_threads") = 1,
R"pbdoc(
Estimate point normals.
If a sufficient number of neighbor points (5 points) is not found, a zero vector is set to the point.
Parameters
----------
points : :class:`PointCloud`
Input point cloud. Normals will be estimated in-place. (in/out)
tree : :class:`KdTree`, optional
Nearest neighbor search. If None, create a new KdTree (default: None)
num_neighbors : int, optional
Number of neighbors. (default: 20)
num_threads : int, optional
Number of threads. (default: 1)
)pbdoc");
// estimate_covariances
m.def(
"estimate_covariances",
[](PointCloud::Ptr points, std::shared_ptr<KdTree<PointCloud>> tree, int num_neighbors, int num_threads) -> std::shared_ptr<PointCloud> {
if (tree == nullptr) {
tree = std::make_shared<KdTree<PointCloud>>(points, KdTreeBuilderOMP(num_threads));
}
if (num_threads == 1) {
estimate_covariances(*points, *tree, num_neighbors);
} else {
estimate_covariances_omp(*points, *tree, num_neighbors, num_threads);
}
return points;
},
py::arg("points"),
py::arg("tree") = nullptr,
py::arg("num_neighbors") = 20,
py::arg("num_threads") = 1,
R"pbdoc(
Estimate point covariances.
If a sufficient number of neighbor points (5 points) is not found, an identity matrix is set to the point.
Parameters
----------
points : :class:`PointCloud`
Input point cloud. Covariances will be estimated in-place. (in/out)
tree : :class:`KdTree`, optional
Nearest neighbor search. If None, create a new KdTree (default: None)
num_neighbors : int, optional
Number of neighbors. (default: 20)
num_threads : int, optional
Number of threads. (default: 1)
)pbdoc");
// estimate_normals_covariances
m.def(
"estimate_normals_covariances",
[](PointCloud::Ptr points, std::shared_ptr<KdTree<PointCloud>> tree, int num_neighbors, int num_threads) {
if (tree == nullptr) {
tree = std::make_shared<KdTree<PointCloud>>(points, KdTreeBuilderOMP(num_threads));
}
if (num_threads == 1) {
estimate_normals_covariances(*points, *tree, num_neighbors);
} else {
estimate_normals_covariances_omp(*points, *tree, num_neighbors, num_threads);
}
},
py::arg("points"),
py::arg("tree") = nullptr,
py::arg("num_neighbors") = 20,
py::arg("num_threads") = 1,
R"pbdoc(
Estimate point normals and covariances.
If a sufficient number of neighbor points (5 points) is not found, a zero vector and an identity matrix is set to the point.
Parameters
----------
points : :class:`PointCloud`
Input point cloud. Normals and covariances will be estimated in-place. (in/out)
tree : :class:`KdTree`, optional
Nearest neighbor search. If None, create a new KdTree (default: None)
num_neighbors : int, optional
Number of neighbors. (default: 20)
num_threads : int, optional
Number of threads. (default: 1)
)pbdoc");
// preprocess_points (numpy)
m.def(
"preprocess_points",
[](const Eigen::MatrixXd& points_numpy, double downsampling_resolution, int num_neighbors, int num_threads) -> std::pair<PointCloud::Ptr, KdTree<PointCloud>::Ptr> {
if (points_numpy.cols() != 3 && points_numpy.cols() != 4) {
std::cerr << "points_numpy must be Nx3 or Nx4" << std::endl;
return {nullptr, nullptr};
}
auto points = std::make_shared<PointCloud>();
points->resize(points_numpy.rows());
for (size_t i = 0; i < points_numpy.rows(); i++) {
if (points_numpy.cols() == 3) {
points->point(i) << points_numpy.row(i).transpose(), 1.0;
} else {
points->point(i) << points_numpy.row(i).transpose();
}
}
auto downsampled = voxelgrid_sampling_omp(*points, downsampling_resolution, num_threads);
auto kdtree = std::make_shared<KdTree<PointCloud>>(downsampled, KdTreeBuilderOMP(num_threads));
estimate_normals_covariances_omp(*downsampled, *kdtree, num_neighbors, num_threads);
return {downsampled, kdtree};
},
py::arg("points"),
py::arg("downsampling_resolution") = 0.25,
py::arg("num_neighbors") = 10,
py::arg("num_threads") = 1,
R"pbdoc(
Preprocess point cloud (Downsampling and normal/covariance estimation).
See also: :func:`voxelgrid_sampling`, :func:`estimate_normals_covariances`.
Parameters
----------
points : [np.float64]
Input point cloud. Nx3 or Nx4.
downsampling_resolution : float, optional
Resolution for downsampling the point clouds. Input points out of 21bit range after discretization are ignored (See also: :func:`voxelgrid_sampling`). (default: 0.25)
num_neighbors : int, optional
Number of neighbors used for attribute estimation. (default: 10)
num_threads : int, optional
Number of threads. (default: 1)
Returns
-------
:class:`PointCloud`
Downsampled point cloud.
:class:`KdTree`
KdTree for the downsampled point cloud.
)pbdoc");
// preprocess_points
m.def(
"preprocess_points",
[](const PointCloud& points, double downsampling_resolution, int num_neighbors, int num_threads) -> std::pair<PointCloud::Ptr, KdTree<PointCloud>::Ptr> {
if (points.empty()) {
std::cerr << "warning: points is empty" << std::endl;
return {nullptr, nullptr};
}
auto downsampled = voxelgrid_sampling_omp(points, downsampling_resolution, num_threads);
auto kdtree = std::make_shared<KdTree<PointCloud>>(downsampled, KdTreeBuilderOMP(num_threads));
estimate_normals_covariances_omp(*downsampled, *kdtree, num_neighbors, num_threads);
return {downsampled, kdtree};
},
py::arg("points"),
py::arg("downsampling_resolution") = 0.25,
py::arg("num_neighbors") = 10,
py::arg("num_threads") = 1,
R"pbdoc(
Preprocess point cloud (Downsampling and normal/covariance estimation).
See also: :func:`voxelgrid_sampling`, :func:`estimate_normals_covariances`.
Parameters
----------
points : :class:`PointCloud`
Input point cloud.
downsampling_resolution : float, optional
Resolution for downsampling the point clouds. Input points out of 21bit range after discretization are ignored (See also: :func:`voxelgrid_sampling`). (default: 0.25)
num_neighbors : int, optional
Number of neighbors used for attribute estimation. (default: 10)
num_threads : int, optional
Number of threads. (default: 1)
Returns
-------
:class:`PointCloud`
Downsampled point cloud.
:class:`KdTree`
KdTree for the downsampled point cloud.
)pbdoc");
}