Add ICP support for 3d point clouds

SLAM/iterative_closest_point/iterative_closest_point.py

@@ -1,10 +1,11 @@
 """
 Iterative Closest Point (ICP) SLAM example
-author: Atsushi Sakai (@Atsushi_twi), Göktuğ Karakaşlı
+author: Atsushi Sakai (@Atsushi_twi), Göktuğ Karakaşlı, Shamil Gemuev
 """
 
 import math
 
+from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -19,8 +20,8 @@ def icp_matching(previous_points, current_points):
     """
     Iterative Closest Point matching
     - input
-    previous_points: 2D points in the previous frame
-    current_points: 2D points in the current frame
+    previous_points: 2D or 3D points in the previous frame
+    current_points: 2D or 3D points in the current frame
     - output
     R: Rotation matrix
     T: Translation vector
@@ -31,19 +32,16 @@ def icp_matching(previous_points, current_points):
     preError = np.inf
     count = 0
 
+    if show_animation:
+        fig = plt.figure()
+        if previous_points.shape[0] == 3:
+            fig.add_subplot(111, projection='3d')
+
     while dError >= EPS:
         count += 1
 
         if show_animation:  # pragma: no cover
-            plt.cla()
-            # for stopping simulation with the esc key.
-            plt.gcf().canvas.mpl_connect(
-                'key_release_event',
-                lambda event: [exit(0) if event.key == 'escape' else None])
-            plt.plot(previous_points[0, :], previous_points[1, :], ".r")
-            plt.plot(current_points[0, :], current_points[1, :], ".b")
-            plt.plot(0.0, 0.0, "xr")
-            plt.axis("equal")
+            plot_points(previous_points, current_points, fig)
             plt.pause(0.1)
 
         indexes, error = nearest_neighbor_association(previous_points, current_points)
@@ -68,24 +66,20 @@ def icp_matching(previous_points, current_points):
             print("Not Converge...", error, dError, count)
             break
 
-    R = np.array(H[0:2, 0:2])
-    T = np.array(H[0:2, 2])
+    R = np.array(H[0:-1, 0:-1])
+    T = np.array(H[0:-1, -1])
 
     return R, T
 
 
 def update_homogeneous_matrix(Hin, R, T):
 
-    H = np.zeros((3, 3))
-
-    H[0, 0] = R[0, 0]
-    H[1, 0] = R[1, 0]
-    H[0, 1] = R[0, 1]
-    H[1, 1] = R[1, 1]
-    H[2, 2] = 1.0
+    r_size = R.shape[0]
+    H = np.zeros((r_size + 1, r_size + 1))
 
-    H[0, 2] = T[0]
-    H[1, 2] = T[1]
+    H[0:r_size, 0:r_size] = R
+    H[0:r_size, r_size] = T
+    H[r_size, r_size] = 1.0
 
     if Hin is None:
         return H
@@ -124,6 +118,31 @@ def svd_motion_estimation(previous_points, current_points):
     return R, t
 
 
+def plot_points(previous_points, current_points, figure):
+    # for stopping simulation with the esc key.
+    plt.gcf().canvas.mpl_connect(
+        'key_release_event',
+        lambda event: [exit(0) if event.key == 'escape' else None])
+    if previous_points.shape[0] == 3:
+        # axes = figure.axes[0]
+        # axes.clear()
+        figure = plt.figure()
+        axes = figure.add_subplot(111, projection='3d')
+        axes.scatter(previous_points[0, :], previous_points[1, :],
+                     previous_points[2, :], c="r", marker=".")
+        axes.scatter(current_points[0, :], current_points[1, :],
+                     current_points[2, :], c="b", marker=".")
+        axes.scatter(0.0, 0.0, 0.0, c="r", marker="x")
+        # figure.canvas.draw()
+        # figure.canvas.flush_events()
+    else:
+        plt.cla()
+        plt.plot(previous_points[0, :], previous_points[1, :], ".r")
+        plt.plot(current_points[0, :], current_points[1, :], ".b")
+        plt.plot(0.0, 0.0, "xr")
+        plt.axis("equal")
+
+
 def main():
     print(__file__ + " start!!")
 
@@ -153,5 +172,37 @@ def main():
         print("T:", T)
 
 
+def main_3d_points():
+    print(__file__ + " start!!")
+
+    # simulation parameters for 3d point set
+    nPoint = 1000
+    fieldLength = 50.0
+    motion = [0.5, 2.0, -5, np.deg2rad(-10.0)]  # [x[m],y[m],z[m],roll[deg]]
+
+    nsim = 3  # number of simulation
+
+    for _ in range(nsim):
+
+        # previous points
+        px = (np.random.rand(nPoint) - 0.5) * fieldLength
+        py = (np.random.rand(nPoint) - 0.5) * fieldLength
+        pz = (np.random.rand(nPoint) - 0.5) * fieldLength
+        previous_points = np.vstack((px, py, pz))
+
+        # current points
+        cx = [math.cos(motion[3]) * x - math.sin(motion[3]) * z + motion[0]
+              for (x, z) in zip(px, pz)]
+        cy = [y + motion[1] for y in py]
+        cz = [math.sin(motion[3]) * x + math.cos(motion[3]) * z + motion[2]
+              for (x, z) in zip(px, pz)]
+        current_points = np.vstack((cx, cy, cz))
+
+        R, T = icp_matching(previous_points, current_points)
+        print("R:", R)
+        print("T:", T)
+
+
 if __name__ == '__main__':
     main()
+    main_3d_points()

tests/test_iterative_closest_point.py

@@ -7,5 +7,10 @@ def test_1():
     m.main()
 
 
+def test_2():
+    m.show_animation = False
+    m.main_3d_points()
+
+
 if __name__ == '__main__':
     conftest.run_this_test(__file__)