return transformation matrix from registration.icp, make max_iter optional

1 files changed, 26 insertions, 6 deletions

diff --git a/pcl/registration.pyx b/pcl/registration.pyx
index 386f7a3..6b8c124 100644
--- a/pcl/registration.pyx
+++ b/pcl/registration.pyx
@@ -4,13 +4,23 @@
 
 from libcpp cimport bool
 
+cimport numpy as np
+import numpy as np
+
 cimport _pcl
 cimport pcl_defs as cpp
 
+np.import_array()
+
+
 cdef extern from "pcl/registration/icp.h" namespace "pcl":
     cdef cppclass IterativeClosestPoint[Source, Target]:
+        cppclass Matrix4:
+            float *data()
+
         IterativeClosestPoint() except +
         void align(cpp.PointCloud[Source] &) except +
+        Matrix4 getFinalTransformation() except +
         double getFitnessScore() except +
         bool hasConverged() except +
         void setInputSource(cpp.PointCloudPtr_t) except +
@@ -19,7 +29,7 @@ cdef extern from "pcl/registration/icp.h" namespace "pcl":
 
 
 def icp(_pcl.PointCloud source, _pcl.PointCloud target,
-        unsigned max_iter):
+        max_iter=None):
     """Align source to target using iterative closests point (ICP).
 
     Parameters
@@ -28,8 +38,9 @@ def icp(_pcl.PointCloud source, _pcl.PointCloud target,
         Source point cloud.
     target : PointCloud
         Target point cloud.
-    max_iter : integer
-        Maximum number of iterations.
+    max_iter : integer, optional
+        Maximum number of iterations. If not given, uses the default number
+        hardwired into PCL.
 
     Returns
     -------
@@ -53,11 +64,20 @@ def icp(_pcl.PointCloud source, _pcl.PointCloud target,
     icp.setInputSource(source.thisptr.makeShared())
     icp.setInputTarget(target.thisptr.makeShared())
 
-    icp.setMaximumIterations(max_iter)
+    if max_iter is not None:
+        icp.setMaximumIterations(max_iter)
 
     cdef _pcl.PointCloud result = _pcl.PointCloud()
 
     icp.align(result.thisptr[0])
 
-    # TODO return transformation as promised.
-    return icp.hasConverged(), result, icp.getFitnessScore()
+    # Convert transformation from Eigen to NumPy format.
+    cdef const float *transf_data = icp.getFinalTransformation().data()
+    cdef np.ndarray[dtype=np.float32_t, ndim=1, mode='c'] transf_flat
+    transf = np.zeros((4, 4), dtype=np.float32)
+    transf_flat = transf.ravel()
+
+    for i in range(16):
+        transf_flat[i] = transf_data[i]
+
+    return icp.hasConverged(), transf, result, icp.getFitnessScore()