fit using ransac and tests

pclpy/__init__.py

@@ -15,4 +15,5 @@
                        moving_least_squares,
                        radius_outlier_removal,
                        octree_voxel_centroid,
+                       fit,
                        )

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):
     """

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)