diff --git a/pclpy/__init__.py b/pclpy/__init__.py index f74e19d..d633ae2 100644 --- a/pclpy/__init__.py +++ b/pclpy/__init__.py @@ -15,4 +15,5 @@ moving_least_squares, radius_outlier_removal, octree_voxel_centroid, + fit, ) diff --git a/pclpy/api.py b/pclpy/api.py index 1d44801..7ba12b4 100644 --- a/pclpy/api.py +++ b/pclpy/api.py @@ -1,5 +1,7 @@ import math +import numpy as np + from . import pcl from .view.vtk import Viewer @@ -198,6 +200,69 @@ def octree_voxel_centroid(cloud, resolution, epsilon=None): return centroids +@register_point_cloud_function +def fit(cloud, model, distance, method=pcl.sample_consensus.SAC_RANSAC, indices=None, optimize=True): + """ + Fit a model to a cloud using a sample consensus method + :param cloud: input point cloud + :param model: str (ex.: 'line', 'sphere', ...) or an instance of pcl.sample_consensus.SacModel + :param distance: distance threshold + :param method: SAC method to use + :param indices: optional indices of the input cloud to use + :param optimize: passed to setOptimizeCoefficients + :return: (inliers: pcl.PointIndices, coefficients: pcl.ModelCoefficients) + """ + models = [ + "circle2d", + "circle3d", + "cone", + "cylinder", # needs normals + "line", # needs normals + "normal_parallel_plane", # needs normals + "normal_plane", # needs normals + "normal_sphere", # needs normals + "parallel_line", + "parallel_lines", + "parallel_plane", + "perpendicular_plane", + "plane", + "registration", + "registration_2d", + "sphere", + "stick", + "torus", # needs normals + ] + if isinstance(model, str): + model = model.lower() + for model_name in models: + if model_name == model: + model = getattr(pcl.sample_consensus, "SACMODEL_" + model_name.upper()) + break + + if not isinstance(model, pcl.sample_consensus.SacModel): # pcl.sample_consensus.SACMODEL_* + message = ("Unrecognized model: %s. Must be either a string " + "or an enum from pcl.sample_consensus.SACMODEL_*") + raise ValueError(message) + + pc_type = utils.get_point_cloud_type(cloud) + seg = getattr(pcl.segmentation.SACSegmentation, pc_type)() + pcl.segmentation + seg.setOptimizeCoefficients(optimize) + seg.setModelType(model) + seg.setMethodType(method) + seg.setDistanceThreshold(distance) + seg.setInputCloud(cloud) + + if indices is not None: + if isinstance(indices, np.ndarray): + indices = pcl.vectors.Int(indices) + seg.setIndices(indices) + coefficients = pcl.ModelCoefficients() + inliers = pcl.PointIndices() + seg.segment(inliers, coefficients) + return inliers, coefficients + + @register_point_cloud_function def show(cloud, *other_clouds, **kwargs): """ diff --git a/pclpy/tests/test_sample_consensus.py b/pclpy/tests/test_sample_consensus.py new file mode 100644 index 0000000..8407388 --- /dev/null +++ b/pclpy/tests/test_sample_consensus.py @@ -0,0 +1,29 @@ +import pytest +import os +import numpy as np +from math import cos, sin, pi + +from pclpy import pcl +import pclpy + + +def test_data(*args): + return os.path.join("test_data", *args) + + +def test_fit_line(): + line = np.array([(1, 2, 3), (2, 4, 6), (3, 7, 9), (5, 10, 15)]) + pc = pcl.PointCloud.PointXYZ(line) + inliers, coefficients = pclpy.fit(pc, "line", distance=0.1) + assert len(inliers.indices) == 3 + assert np.allclose(coefficients.values, pcl.vectors.Float([2.66667, 5.33333, 8, 0.267261, 0.534522, 0.801784])) + + +def test_fit_cylinder(): + points = np.array([(cos(a), sin(a), z) for z in np.linspace(0, 5, num=10) for a in np.linspace(0, 2 * pi, num=20)]) + points = np.vstack((np.array([10, 10, 10]), points)) + pc = pcl.PointCloud.PointXYZ(points) + inliers, coefficients = pclpy.fit(pc, "circle2d", distance=0.01, indices=np.arange(100)) + assert 0 not in inliers.indices + assert len(inliers.indices) == 99 + assert np.allclose(coefficients.values, [0., 0., 1.], atol=0.00001)