Merge pull request #19 from UCL/18-limit-go-icp

18 limit go icp rotation/translation

sksurgerysurfacematch/algorithms/goicp_registration.py

@@ -6,7 +6,7 @@
 
 import logging
 import numpy as np
-from sksurgerygoicppython import  GoICP, POINT3D
+from sksurgerygoicppython import  GoICP, POINT3D, ROTNODE, TRANSNODE
 import sksurgerysurfacematch.interfaces.rigid_registration as rr
 
 LOGGER = logging.getLogger(__name__)
@@ -69,22 +69,61 @@ def demean_and_normalise(points_a: np.ndarray,
     return a_normalised, b_normalised, scale_matrix, \
          translate_a_matrix, translate_b_matrix
 
+def set_rotnode(limits_degrees) -> ROTNODE:
+    """ Setup a ROTNODE with upper/lower rotation limits"""
+
+    lower_degrees = limits_degrees[0]
+    upper_degrees = limits_degrees[1]
+
+    l_rads = lower_degrees * 3.14 / 180
+    u_rads = upper_degrees * 3.14 / 180
+
+    r_node = ROTNODE()
+
+    r_node.a = l_rads
+    r_node.b = l_rads
+    r_node.c = l_rads
+    r_node.w = u_rads - l_rads
+
+    return r_node
+
+def set_transnode(trans_limits) -> TRANSNODE:
+    """ Setup a TRANSNODE with upper/lower limits"""
+
+    t_node = TRANSNODE()
+
+    t_node.x = trans_limits[0]
+    t_node.y = trans_limits[0]
+    t_node.z = trans_limits[0]
+
+    t_node.w = trans_limits[1] - trans_limits[0]
+
+    return t_node
+
 class RigidRegistration(rr.RigidRegistration):
     """
     Class that uses GoICP implementation to register fixed/moving clouds.
     At the moment, we are just relying on all default parameters.
     :param dt_size: Nodes per dimension of distance transform
     :param dt_factor: GoICP distance transform factor
+    TODO: rest of params
     """
-
+    #pylint:disable=dangerous-default-value
     def __init__(self,
                  dt_size: int = 200,
                  dt_factor: float = 2.0,
                  normalise: bool = True,
-                 num_moving_points: int = 1000):
+                 num_moving_points: int = 1000,
+                 rotation_limits=[-45, 45],
+                 trans_limits=[-0.5, 0.5]):
+
+        r_node = set_rotnode(rotation_limits)
+        t_node = set_transnode(trans_limits)
 
         self.goicp = GoICP()
         self.goicp.setDTSizeAndFactor(dt_size, dt_factor)
+        self.goicp.setInitNodeRot(r_node)
+        self.goicp.setInitNodeTrans(t_node)
 
         self.normalise = normalise
         self.num_moving_points = num_moving_points

sksurgerysurfacematch/algorithms/reconstructor_with_rectified_images.py

@@ -38,6 +38,10 @@ def __init__(self,
         self.lower_disparity_multiplier = lower_disparity_multiplier
         self.upper_disparity_multiplier = upper_disparity_multiplier
 
+        self.left_rectified = None
+        self.right_rectified = None
+        self.left_mask = None
+
     # pylint:disable=too-many-arguments
     def reconstruct(self,
                     left_image: np.ndarray,
@@ -81,7 +85,6 @@ def reconstruct(self,
                               (width, height),
                               left_to_right_rmat,
                               left_to_right_tvec,
-                              alpha=0  # Only keep points visible in both
                               )
 
         undistort_rectify_map_l_x, undistort_rectify_map_l_y = \
@@ -96,25 +99,27 @@ def reconstruct(self,
                                         r_2, p_2,
                                         (width, height), cv2.CV_32FC1)
 
-        left_rectified = cv2.remap(left_image, undistort_rectify_map_l_x,
-                                   undistort_rectify_map_l_y, cv2.INTER_LINEAR)
+        self.left_rectified = \
+            cv2.remap(left_image, undistort_rectify_map_l_x,
+                      undistort_rectify_map_l_y, cv2.INTER_LINEAR)
 
-        right_rectified = cv2.remap(right_image, undistort_rectify_map_r_x,
-                                    undistort_rectify_map_r_y, cv2.INTER_LINEAR)
+        self.right_rectified = \
+            cv2.remap(right_image, undistort_rectify_map_r_x,
+                      undistort_rectify_map_r_y, cv2.INTER_LINEAR)
 
         # Need to remap the mask if we have one
+        self.left_mask = left_mask
         if left_mask is not None:
 
-            left_mask = cv2.remap(left_mask, undistort_rectify_map_l_x,
-                                  undistort_rectify_map_l_y, cv2.INTER_NEAREST)
-
-        self.disparity = self._compute_disparity(left_rectified,
-                                                 right_rectified)
+            self.left_mask = \
+                cv2.remap(left_mask, undistort_rectify_map_l_x,
+                          undistort_rectify_map_l_y, cv2.INTER_NEAREST)
 
-        print(f'Disparity: {self.disparity.shape}')
+        self.disparity = self._compute_disparity(self.left_rectified,
+                                                 self.right_rectified)
 
         self.points = cv2.reprojectImageTo3D(self.disparity, q_mat)
-        self.rgb_image = cv2.cvtColor(left_rectified, cv2.COLOR_BGR2RGB)
+        self.rgb_image = cv2.cvtColor(self.left_rectified, cv2.COLOR_BGR2RGB)
 
         # Calls method below to extract data.
         return self.extract(left_mask)

tests/test_icp_goicp.py

@@ -29,7 +29,7 @@ def test_goicp_reg():
     fixed = np.loadtxt('tests/data/icp/rabbit_full.xyz')
     moving = np.loadtxt('tests/data/icp/rabbit_partial.xyz')
 
-    goicp_reg = goicp.RigidRegistration()
+    goicp_reg = goicp.RigidRegistration(rotation_limits=[-180, 180])
 
     # Data already normalsied
     residual, moving_to_fixed = goicp_reg.register(moving, fixed)
@@ -68,7 +68,7 @@ def test_goicp_known_transform():
                                         [0, 0, 0, 1]])
 
     moving = transform_points(fixed, fixed_to_moving)
-    goicp_reg = goicp.RigidRegistration()
+    goicp_reg = goicp.RigidRegistration(rotation_limits=[-180, 180])
     residual, moving_to_fixed = goicp_reg.register(moving, fixed)
 
 

tests/test_surface_recon.py

@@ -101,7 +101,7 @@ def test_stoyanov_and_sgbm():
     print("SGBM, cloud=" + str(points.shape))
 
     assert points.shape[1] == 6
-    assert 104000 >= points.shape[0] >= 103000 # Slightly different answers between windows and mac/linux
+    assert 104000 >= points.shape[0] >= 102000 # Slightly different answers between windows and mac/linux
 
     points = reconstructor.reconstruct(left_undistorted,
                                        left_intrinsics,
@@ -116,6 +116,6 @@ def test_stoyanov_and_sgbm():
     pl.write_pointcloud(points[:, 0:3], points[:, 3:6], 'tests/output/sgbm_masked.ply')
 
     assert points.shape[1] == 6
-    assert 8000 >= points.shape[0] >= 7800
+    assert 8000 >= points.shape[0] >= 6000
 
     print("SGBM masked, cloud=" + str(points.shape))