diff --git a/sksurgerysurfacematch/algorithms/goicp_registration.py b/sksurgerysurfacematch/algorithms/goicp_registration.py index c4b3f46..ab7b425 100644 --- a/sksurgerysurfacematch/algorithms/goicp_registration.py +++ b/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 diff --git a/sksurgerysurfacematch/algorithms/reconstructor_with_rectified_images.py b/sksurgerysurfacematch/algorithms/reconstructor_with_rectified_images.py index 1c93ebe..55d1a21 100644 --- a/sksurgerysurfacematch/algorithms/reconstructor_with_rectified_images.py +++ b/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) diff --git a/tests/test_icp_goicp.py b/tests/test_icp_goicp.py index c49233c..bf6609d 100644 --- a/tests/test_icp_goicp.py +++ b/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) diff --git a/tests/test_surface_recon.py b/tests/test_surface_recon.py index c214e23..8ef5cff 100644 --- a/tests/test_surface_recon.py +++ b/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))