test: Add comprehensive tests for global ICP init and fallback

py_workspace/tests/test_icp_registration.py

@@ -487,3 +487,280 @@ def test_refine_with_icp_single_camera():
 
     assert metrics.num_cameras_optimized == 1
     assert metrics.success
+
+
+def test_compute_fpfh_features():
+    pcd = create_box_pcd()
+    voxel_size = 0.05
+    pcd_down = pcd.voxel_down_sample(voxel_size)
+
+    fpfh = compute_fpfh_features(pcd_down, voxel_size)
+
+    assert fpfh.dimension() == 33
+    assert fpfh.num() == len(pcd_down.points)
+
+
+def test_global_registration_known_transform():
+    source = create_box_pcd(size=1.0, num_points=1000)
+
+    # Create target with significant transform (rotation + translation)
+    T_true = np.eye(4)
+    # 30 degree rotation around Y
+    T_true[:3, :3] = Rotation.from_euler("y", 30, degrees=True).as_matrix()
+    T_true[:3, 3] = [0.5, 0.0, 0.2]
+
+    target = o3d.geometry.PointCloud()
+    target.points = o3d.utility.Vector3dVector(
+        (np.asarray(source.points) @ T_true[:3, :3].T) + T_true[:3, 3]
+    )
+
+    voxel_size = 0.05
+    source_down = source.voxel_down_sample(voxel_size)
+    target_down = target.voxel_down_sample(voxel_size)
+
+    source_fpfh = compute_fpfh_features(source_down, voxel_size)
+    target_fpfh = compute_fpfh_features(target_down, voxel_size)
+
+    result = global_registration(
+        source_down, target_down, source_fpfh, target_fpfh, voxel_size
+    )
+
+    assert result.fitness > 0.1
+    # RANSAC is stochastic, but with clean data it should be reasonably close
+    # We check if it found a transform close to truth
+    T_est = result.transformation
+
+    # Check rotation difference
+    R_diff = T_est[:3, :3] @ T_true[:3, :3].T
+    rot_diff = Rotation.from_matrix(R_diff).as_euler("xyz", degrees=True)
+    assert np.linalg.norm(rot_diff) < 5.0  # Within 5 degrees
+
+    # Check translation difference
+    trans_diff = np.linalg.norm(T_est[:3, 3] - T_true[:3, 3])
+    assert trans_diff < 0.1  # Within 10cm
+
+
+def test_refine_with_icp_global_init_success():
+    import aruco.icp_registration
+    import aruco.ground_plane
+
+    # Create two overlapping point clouds with large offset
+    box_points = create_box_pcd(size=1.0, num_points=2000).points
+    box_points = np.asarray(box_points)
+
+    # Mock unproject to return these points
+    def mock_unproject(depth, K, stride=1, **kwargs):
+        if depth[0, 0] == 1.0:  # cam1
+            return box_points
+        else:  # cam2
+            # Apply large transform to simulate misaligned camera
+            # Rotate 90 deg and translate
+            R = Rotation.from_euler("y", 90, degrees=True).as_matrix()
+            return (box_points @ R.T) + [2.0, 0, 0]
+
+    orig_unproject = aruco.ground_plane.unproject_depth_to_points
+    aruco.ground_plane.unproject_depth_to_points = mock_unproject
+
+    try:
+        camera_data = {
+            "cam1": {"depth": np.ones((10, 10)), "K": np.eye(3)},
+            "cam2": {"depth": np.zeros((10, 10)), "K": np.eye(3)},
+        }
+
+        # Initial extrinsics are identity (very wrong for cam2)
+        extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4)}
+
+        floor_planes = {
+            "cam1": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+            "cam2": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+        }
+
+        # Enable global init
+        config = ICPConfig(
+            global_init=True,
+            min_overlap_area=0.0,  # Disable overlap check for this test since initial overlap is zero
+            min_fitness=0.1,
+            voxel_size=0.05,
+            max_rotation_deg=180.0,  # Allow large rotation for this test
+            max_translation_m=5.0,  # Allow large translation
+        )
+
+        new_extrinsics, metrics = refine_with_icp(
+            camera_data, extrinsics, floor_planes, config
+        )
+
+        assert metrics.success
+        assert metrics.num_cameras_optimized == 2
+
+        # Check if cam2 moved significantly from identity
+        T_cam2 = new_extrinsics["cam2"]
+        assert np.linalg.norm(T_cam2[:3, 3]) > 1.0
+
+    finally:
+        aruco.ground_plane.unproject_depth_to_points = orig_unproject
+
+
+def test_refine_with_icp_global_init_disabled():
+    """Verify that global registration is skipped when disabled."""
+    import aruco.icp_registration
+    import aruco.ground_plane
+
+    # Mock global_registration to raise error if called
+    orig_global_reg = aruco.icp_registration.global_registration
+
+    def mock_global_reg(*args, **kwargs):
+        raise RuntimeError("Global registration should not be called")
+
+    aruco.icp_registration.global_registration = mock_global_reg
+
+    # Mock unproject
+    box_points = create_box_pcd(size=0.5).points
+    box_points = np.asarray(box_points)
+
+    def mock_unproject(depth, K, stride=1, **kwargs):
+        if depth[0, 0] == 1.0:
+            return box_points
+        else:
+            return box_points  # Perfect overlap
+
+    orig_unproject = aruco.ground_plane.unproject_depth_to_points
+    aruco.ground_plane.unproject_depth_to_points = mock_unproject
+
+    try:
+        camera_data = {
+            "cam1": {"depth": np.ones((10, 10)), "K": np.eye(3)},
+            "cam2": {"depth": np.zeros((10, 10)), "K": np.eye(3)},
+        }
+        extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4)}
+        floor_planes = {
+            "cam1": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+            "cam2": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+        }
+
+        # Default config (global_init=False)
+        config = ICPConfig(min_overlap_area=0.01)
+
+        # Should not raise RuntimeError
+        refine_with_icp(camera_data, extrinsics, floor_planes, config)
+
+    finally:
+        aruco.ground_plane.unproject_depth_to_points = orig_unproject
+        aruco.icp_registration.global_registration = orig_global_reg
+
+
+def test_refine_with_icp_global_init_fallback_low_fitness():
+    """Verify fallback to original init when global registration has low fitness."""
+    import aruco.icp_registration
+    import aruco.ground_plane
+
+    # Mock global_registration to return low fitness
+    orig_global_reg = aruco.icp_registration.global_registration
+
+    class MockResult:
+        def __init__(self):
+            self.fitness = 0.05  # Below 0.1 threshold
+            self.transformation = np.eye(4)
+            self.transformation[0, 3] = 100.0  # Crazy transform
+
+    def mock_global_reg(*args, **kwargs):
+        return MockResult()
+
+    aruco.icp_registration.global_registration = mock_global_reg
+
+    # Mock unproject
+    box_points = create_box_pcd(size=0.5).points
+    box_points = np.asarray(box_points)
+
+    def mock_unproject(depth, K, stride=1, **kwargs):
+        if depth[0, 0] == 1.0:
+            return box_points
+        else:
+            return box_points  # Perfect overlap
+
+    orig_unproject = aruco.ground_plane.unproject_depth_to_points
+    aruco.ground_plane.unproject_depth_to_points = mock_unproject
+
+    try:
+        camera_data = {
+            "cam1": {"depth": np.ones((10, 10)), "K": np.eye(3)},
+            "cam2": {"depth": np.zeros((10, 10)), "K": np.eye(3)},
+        }
+        extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4)}
+        floor_planes = {
+            "cam1": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+            "cam2": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+        }
+
+        config = ICPConfig(global_init=True, min_overlap_area=0.01)
+
+        new_extrinsics, metrics = refine_with_icp(
+            camera_data, extrinsics, floor_planes, config
+        )
+
+        # Should have ignored the crazy transform from global reg
+        assert np.allclose(new_extrinsics["cam2"], np.eye(4), atol=0.1)
+
+    finally:
+        aruco.ground_plane.unproject_depth_to_points = orig_unproject
+        aruco.icp_registration.global_registration = orig_global_reg
+
+
+def test_refine_with_icp_global_init_fallback_bounds_check():
+    """Verify fallback when global registration exceeds safety bounds."""
+    import aruco.icp_registration
+    import aruco.ground_plane
+
+    # Mock global_registration to return unsafe transform
+    orig_global_reg = aruco.icp_registration.global_registration
+
+    class MockResult:
+        def __init__(self):
+            self.fitness = 1.0  # High fitness
+            self.transformation = np.eye(4)
+            self.transformation[0, 3] = 10.0  # 10m translation (unsafe)
+
+    def mock_global_reg(*args, **kwargs):
+        return MockResult()
+
+    aruco.icp_registration.global_registration = mock_global_reg
+
+    # Mock unproject
+    box_points = create_box_pcd(size=0.5).points
+    box_points = np.asarray(box_points)
+
+    def mock_unproject(depth, K, stride=1, **kwargs):
+        if depth[0, 0] == 1.0:
+            return box_points
+        else:
+            return box_points  # Perfect overlap
+
+    orig_unproject = aruco.ground_plane.unproject_depth_to_points
+    aruco.ground_plane.unproject_depth_to_points = mock_unproject
+
+    try:
+        camera_data = {
+            "cam1": {"depth": np.ones((10, 10)), "K": np.eye(3)},
+            "cam2": {"depth": np.zeros((10, 10)), "K": np.eye(3)},
+        }
+        extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4)}
+        floor_planes = {
+            "cam1": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+            "cam2": FloorPlane(normal=np.array([0, 1, 0]), d=0.0),
+        }
+
+        config = ICPConfig(
+            global_init=True,
+            min_overlap_area=0.01,
+            max_translation_m=0.5,  # Strict bound
+        )
+
+        new_extrinsics, metrics = refine_with_icp(
+            camera_data, extrinsics, floor_planes, config
+        )
+
+        # Should have ignored the unsafe transform
+        assert np.allclose(new_extrinsics["cam2"], np.eye(4), atol=0.1)
+
+    finally:
+        aruco.ground_plane.unproject_depth_to_points = orig_unproject
+        aruco.icp_registration.global_registration = orig_global_reg