feat: implement ICP registration for ground plane refinement and add tests

2026-02-10 03:04:43 +00:00
parent 2d42e2cdfa
commit 206c6e58ee
7 changed files with 1068 additions and 15 deletions
@@ -0,0 +1,312 @@
+import numpy as np
+import pytest
+import open3d as o3d
+from scipy.spatial.transform import Rotation
+from aruco.icp_registration import (
+    ICPConfig,
+    ICPResult,
+    ICPMetrics,
+    extract_near_floor_band,
+    compute_overlap_xz,
+    apply_gravity_constraint,
+    pairwise_icp,
+    build_pose_graph,
+    refine_with_icp,
+)
+from aruco.ground_plane import FloorPlane
+
+
+def test_extract_near_floor_band_basic():
+    points = np.array(
+        [[0, -0.1, 0], [0, 0.1, 0], [0, 0.2, 0], [0, 0.4, 0]], dtype=np.float64
+    )
+
+    floor_y = 0.0
+    band_height = 0.3
+    floor_normal = np.array([0, 1, 0], dtype=np.float64)
+
+    result = extract_near_floor_band(points, floor_y, band_height, floor_normal)
+    assert len(result) == 2
+    assert np.all(result[:, 1] >= 0)
+    assert np.all(result[:, 1] <= 0.3)
+
+
+def test_extract_near_floor_band_empty():
+    points = np.zeros((0, 3))
+    result = extract_near_floor_band(points, 0.0, 0.3, np.array([0, 1, 0]))
+    assert len(result) == 0
+
+
+def test_extract_near_floor_band_all_in():
+    points = np.random.uniform(0, 0.2, (100, 3))
+    points[:, 1] = np.random.uniform(0.05, 0.25, 100)
+    result = extract_near_floor_band(points, 0.0, 0.3, np.array([0, 1, 0]))
+    assert len(result) == 100
+
+
+def test_compute_overlap_xz_full():
+    points_a = np.array([[0, 0, 0], [1, 0, 1]])
+    points_b = np.array([[0, 0, 0], [1, 0, 1]])
+    area = compute_overlap_xz(points_a, points_b)
+    assert abs(area - 1.0) < 1e-6
+
+
+def test_compute_overlap_xz_no():
+    points_a = np.array([[0, 0, 0], [1, 0, 1]])
+    points_b = np.array([[2, 0, 2], [3, 0, 3]])
+    area = compute_overlap_xz(points_a, points_b)
+    assert area == 0.0
+
+
+def test_compute_overlap_xz_partial():
+    points_a = np.array([[0, 0, 0], [1, 0, 1]])
+    points_b = np.array([[0.5, 0, 0.5], [1.5, 0, 1.5]])
+    area = compute_overlap_xz(points_a, points_b)
+    assert abs(area - 0.25) < 1e-6
+
+
+def test_compute_overlap_xz_with_margin():
+    points_a = np.array([[0, 0, 0], [1, 0, 1]])
+    points_b = np.array([[1.2, 0, 0], [2.2, 0, 1]])
+    area_no_margin = compute_overlap_xz(points_a, points_b)
+    area_with_margin = compute_overlap_xz(points_a, points_b, margin=0.5)
+    assert area_no_margin == 0.0
+    assert area_with_margin > 0.0
+
+
+def test_apply_gravity_constraint_identity():
+    T = np.eye(4)
+    result = apply_gravity_constraint(T, T)
+    np.testing.assert_allclose(result, T, atol=1e-6)
+
+
+def test_apply_gravity_constraint_preserves_yaw():
+    T_orig = np.eye(4)
+    R_icp = Rotation.from_euler("xyz", [0, 10, 0], degrees=True).as_matrix()
+    T_icp = np.eye(4)
+    T_icp[:3, :3] = R_icp
+
+    result = apply_gravity_constraint(T_icp, T_orig, penalty_weight=10.0)
+
+    res_euler = Rotation.from_matrix(result[:3, :3]).as_euler("xyz", degrees=True)
+    assert abs(res_euler[1] - 10.0) < 1e-6
+    assert abs(res_euler[0]) < 1e-6
+    assert abs(res_euler[2]) < 1e-6
+
+
+def test_apply_gravity_constraint_regularizes_pitch_roll():
+    T_orig = np.eye(4)
+    R_icp = Rotation.from_euler("xyz", [10, 0, 10], degrees=True).as_matrix()
+    T_icp = np.eye(4)
+    T_icp[:3, :3] = R_icp
+
+    result = apply_gravity_constraint(T_icp, T_orig, penalty_weight=9.0)
+
+    res_euler = Rotation.from_matrix(result[:3, :3]).as_euler("xyz", degrees=True)
+    assert abs(res_euler[0] - 1.0) < 1e-2
+    assert abs(res_euler[2] - 1.0) < 1e-2
+
+
+def create_box_pcd(size=1.0, num_points=500, seed=42):
+    rng = np.random.default_rng(seed)
+    points = rng.uniform(0, size, (num_points, 3))
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    return pcd
+
+
+def test_pairwise_icp_identity():
+    pcd = create_box_pcd()
+    config = ICPConfig(min_fitness=0.1)
+    result = pairwise_icp(pcd, pcd, config, np.eye(4))
+
+    assert result.converged
+    assert result.fitness > 0.9
+    np.testing.assert_allclose(result.transformation, np.eye(4), atol=1e-3)
+
+
+def test_pairwise_icp_known_transform():
+    source = create_box_pcd()
+    T_true = np.eye(4)
+    T_true[:3, :3] = Rotation.from_euler("y", 5, degrees=True).as_matrix()
+    T_true[:3, 3] = [0.05, 0, 0.02]
+
+    target = o3d.geometry.PointCloud()
+    target.points = o3d.utility.Vector3dVector(
+        (np.asarray(source.points) @ T_true[:3, :3].T) + T_true[:3, 3]
+    )
+
+    config = ICPConfig(min_fitness=0.5, voxel_size=0.01)
+    result = pairwise_icp(source, target, config, np.eye(4))
+
+    assert result.converged
+    np.testing.assert_allclose(result.transformation, T_true, atol=1e-2)
+
+
+def test_pairwise_icp_known_transform_gicp():
+    source = create_box_pcd()
+    T_true = np.eye(4)
+    T_true[:3, :3] = Rotation.from_euler("y", 5, degrees=True).as_matrix()
+    T_true[:3, 3] = [0.05, 0, 0.02]
+
+    target = o3d.geometry.PointCloud()
+    target.points = o3d.utility.Vector3dVector(
+        (np.asarray(source.points) @ T_true[:3, :3].T) + T_true[:3, 3]
+    )
+
+    # GICP usually needs normals, which pairwise_icp estimates internally
+    config = ICPConfig(min_fitness=0.5, voxel_size=0.01, method="gicp")
+    result = pairwise_icp(source, target, config, np.eye(4))
+
+    assert result.converged
+    np.testing.assert_allclose(result.transformation, T_true, atol=1e-2)
+
+
+def test_pairwise_icp_insufficient_points():
+    source = o3d.geometry.PointCloud()
+    source.points = o3d.utility.Vector3dVector(np.random.rand(5, 3))
+    target = o3d.geometry.PointCloud()
+    target.points = o3d.utility.Vector3dVector(np.random.rand(5, 3))
+
+    config = ICPConfig(min_fitness=0.9)
+    result = pairwise_icp(source, target, config, np.eye(4))
+    assert not result.converged
+
+
+def test_build_pose_graph_basic():
+    serials = ["cam1", "cam2"]
+    extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4)}
+
+    res = ICPResult(
+        transformation=np.eye(4),
+        fitness=1.0,
+        inlier_rmse=0.0,
+        information_matrix=np.eye(6),
+        converged=True,
+    )
+    pair_results = {("cam1", "cam2"): res}
+
+    graph = build_pose_graph(serials, extrinsics, pair_results, "cam1")
+    assert len(graph.nodes) == 2
+    assert len(graph.edges) == 1
+
+
+def test_build_pose_graph_disconnected():
+    serials = ["cam1", "cam2", "cam3"]
+    extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4), "cam3": np.eye(4)}
+
+    res = ICPResult(
+        transformation=np.eye(4),
+        fitness=1.0,
+        inlier_rmse=0.0,
+        information_matrix=np.eye(6),
+        converged=True,
+    )
+    pair_results = {("cam1", "cam2"): res}
+
+    graph = build_pose_graph(serials, extrinsics, pair_results, "cam1")
+    assert len(graph.nodes) == 2
+    assert len(graph.edges) == 1
+
+
+def test_refine_with_icp_synthetic_offset():
+    import aruco.icp_registration
+    import aruco.ground_plane
+
+    box_points = create_box_pcd(size=0.5).points
+    box_points = np.asarray(box_points)
+    box_points[:, 1] -= 1.0
+    box_points[:, 2] += 2.0
+
+    def mock_unproject(depth, K, stride=1, **kwargs):
+        if depth[0, 0] == 1.0:
+            return box_points
+        else:
+            return box_points - np.array([1.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)},
+        }
+        T_w1 = np.eye(4)
+        T_w2_est = np.eye(4)
+        T_w2_est[0, 3] = 1.05
+
+        extrinsics = {"cam1": T_w1, "cam2": T_w2_est}
+
+        floor_planes = {
+            "cam1": FloorPlane(normal=np.array([0, 1, 0]), d=1.0),
+            "cam2": FloorPlane(normal=np.array([0, 1, 0]), d=1.0),
+        }
+
+        config = ICPConfig(
+            min_overlap_area=0.01,
+            min_fitness=0.1,
+            voxel_size=0.05,
+            max_iterations=[20, 10, 5],
+            max_translation_m=3.0,
+        )
+
+        new_extrinsics, metrics = refine_with_icp(
+            camera_data, extrinsics, floor_planes, config
+        )
+
+        assert metrics.success
+        assert metrics.num_cameras_optimized == 2
+        assert abs(new_extrinsics["cam2"][0, 3] - T_w2_est[0, 3]) > 0.01
+
+    finally:
+        aruco.ground_plane.unproject_depth_to_points = orig_unproject
+
+
+def test_refine_with_icp_no_overlap():
+    import aruco.icp_registration
+    import aruco.ground_plane
+
+    def mock_unproject(depth, K, stride=1, **kwargs):
+        if depth[0, 0] == 1.0:
+            return np.random.rand(200, 3) + [0, -1, 0]
+        else:
+            return np.random.rand(200, 3) + [10, -1, 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)},
+        }
+        extrinsics = {"cam1": np.eye(4), "cam2": np.eye(4)}
+        floor_planes = {
+            "cam1": FloorPlane(normal=np.array([0, 1, 0]), d=1.0),
+            "cam2": FloorPlane(normal=np.array([0, 1, 0]), d=1.0),
+        }
+
+        config = ICPConfig(min_overlap_area=1.0)
+        new_extrinsics, metrics = refine_with_icp(
+            camera_data, extrinsics, floor_planes, config
+        )
+
+        assert not metrics.success
+        assert "No converged ICP pairs" in metrics.message
+
+    finally:
+        aruco.ground_plane.unproject_depth_to_points = orig_unproject
+
+
+def test_refine_with_icp_single_camera():
+    camera_data = {"cam1": {"depth": np.ones((10, 10)), "K": np.eye(3)}}
+    extrinsics = {"cam1": np.eye(4)}
+    floor_planes = {"cam1": FloorPlane(normal=np.array([0, 1, 0]), d=1.0)}
+
+    config = ICPConfig()
+    new_extrinsics, metrics = refine_with_icp(
+        camera_data, extrinsics, floor_planes, config
+    )
+
+    assert not metrics.success