From bfb3421692207a18c2bf30397c25b01958855179 Mon Sep 17 00:00:00 2001 From: crosstyan Date: Mon, 9 Feb 2026 07:36:27 +0000 Subject: [PATCH] fix: implement per-camera ground plane correction --- py_workspace/aruco/ground_plane.py | 80 +++++++++++++++---------- py_workspace/tests/test_ground_plane.py | 69 ++++++++++++++++++++- 2 files changed, 116 insertions(+), 33 deletions(-) diff --git a/py_workspace/aruco/ground_plane.py b/py_workspace/aruco/ground_plane.py index 32a274e..7dfdddb 100644 --- a/py_workspace/aruco/ground_plane.py +++ b/py_workspace/aruco/ground_plane.py @@ -51,7 +51,10 @@ class GroundPlaneMetrics: correction_applied: bool = False num_cameras_total: int = 0 num_cameras_valid: int = 0 - correction_transform: Mat44 = field(default_factory=lambda: np.eye(4)) + # Per-camera corrections + camera_corrections: Dict[str, Mat44] = field(default_factory=dict) + skipped_cameras: List[str] = field(default_factory=list) + # Summary stats (optional, maybe average of corrections?) rotation_deg: float = 0.0 translation_m: float = 0.0 camera_planes: Dict[str, FloorPlane] = field(default_factory=dict) @@ -354,7 +357,7 @@ def refine_ground_from_depth( metrics.message = f"Found {len(valid_planes)} valid planes, required {config.min_valid_cameras}" return extrinsics, metrics - # 3. Compute consensus + # 3. Compute consensus (for reporting/metrics only) try: consensus = compute_consensus_plane(valid_planes) metrics.consensus_plane = consensus @@ -362,42 +365,55 @@ def refine_ground_from_depth( metrics.message = f"Consensus computation failed: {e}" return extrinsics, metrics - # 4. Compute correction - correction = compute_floor_correction( - consensus, - target_floor_y=config.target_y, - max_rotation_deg=config.max_rotation_deg, - max_translation_m=config.max_translation_m, - ) + # 4. Compute and apply per-camera correction + new_extrinsics = extrinsics.copy() - metrics.correction_transform = correction.transform - - if not correction.valid: - metrics.message = f"Correction invalid: {correction.reason}" - return extrinsics, metrics - - # 5. Apply correction - # T_corr is the transform that moves the world frame. - # New world points P' = T_corr * P - # We want new extrinsics T'_world_cam such that P' = T'_world_cam * P_cam - # T'_world_cam * P_cam = T_corr * (T_world_cam * P_cam) - # So T'_world_cam = T_corr * T_world_cam - - new_extrinsics = {} - T_corr = correction.transform + # Track max rotation/translation for summary metrics + max_rot = 0.0 + max_trans = 0.0 + corrections_count = 0 for serial, T_old in extrinsics.items(): + # If we didn't find a plane for this camera, skip it + if serial not in metrics.camera_planes: + metrics.skipped_cameras.append(serial) + continue + + plane = metrics.camera_planes[serial] + + correction = compute_floor_correction( + plane, + target_floor_y=config.target_y, + max_rotation_deg=config.max_rotation_deg, + max_translation_m=config.max_translation_m, + ) + + if not correction.valid: + metrics.skipped_cameras.append(serial) + continue + + T_corr = correction.transform + metrics.camera_corrections[serial] = T_corr + + # Apply correction: T_new = T_corr @ T_old new_extrinsics[serial] = T_corr @ T_old - # Calculate metrics - # Rotation angle of T_corr - trace = np.trace(T_corr[:3, :3]) - cos_angle = np.clip((trace - 1) / 2, -1.0, 1.0) - metrics.rotation_deg = float(np.rad2deg(np.arccos(cos_angle))) - metrics.translation_m = float(np.linalg.norm(T_corr[:3, 3])) + # Update summary metrics + trace = np.trace(T_corr[:3, :3]) + cos_angle = np.clip((trace - 1) / 2, -1.0, 1.0) + rot_deg = float(np.rad2deg(np.arccos(cos_angle))) + trans_m = float(np.linalg.norm(T_corr[:3, 3])) + max_rot = max(max_rot, rot_deg) + max_trans = max(max_trans, trans_m) + corrections_count += 1 + + metrics.rotation_deg = max_rot + metrics.translation_m = max_trans metrics.success = True - metrics.correction_applied = True - metrics.message = "Success" + metrics.correction_applied = corrections_count > 0 + metrics.message = ( + f"Corrected {corrections_count} cameras, skipped {len(metrics.skipped_cameras)}" + ) return new_extrinsics, metrics diff --git a/py_workspace/tests/test_ground_plane.py b/py_workspace/tests/test_ground_plane.py index 0f30a49..6bb7ece 100644 --- a/py_workspace/tests/test_ground_plane.py +++ b/py_workspace/tests/test_ground_plane.py @@ -466,7 +466,7 @@ def test_refine_ground_from_depth_success(): # We started with floor at y=-1.0. Target is y=0.0. # So we expect translation of +1.0 in Y. # T_corr should have ty approx 1.0. - T_corr = metrics.correction_transform + T_corr = metrics.camera_corrections["cam1"] assert abs(T_corr[1, 3] - 1.0) < 0.1 # Allow some slack for RANSAC noise # Check new extrinsics @@ -475,3 +475,70 @@ def test_refine_ground_from_depth_success(): # New T origin y should be -3 + 1 = -2. T_new = new_extrinsics["cam1"] assert abs(T_new[1, 3] - (-2.0)) < 0.1 + + # Verify per-camera corrections + assert "cam1" in metrics.camera_corrections + assert "cam2" in metrics.camera_corrections + assert len(metrics.camera_corrections) == 2 + + +def test_refine_ground_from_depth_partial_success(): + # 2 cameras, but only 1 finds a plane + # Should still succeed if min_valid_cameras=1 (or if we relax it) + # But config default is 2. + # Let's set min_valid_cameras=1 + config = GroundPlaneConfig( + min_valid_cameras=1, + min_inliers=10, + target_y=0.0, + ) + + width, height = 20, 20 + K = np.eye(3) + K[0, 2] = 10 + K[1, 2] = 10 + K[0, 0] = 20 + K[1, 1] = 20 + + # Cam 1: Valid plane (same as success test) + Rx_neg90 = np.array([[1, 0, 0], [0, 0, 1], [0, -1, 0]]) + t = np.array([0, -3, 0]) + T_world_cam = np.eye(4) + T_world_cam[:3, :3] = Rx_neg90 + T_world_cam[:3, 3] = t + depth_map_valid = np.full((height, width), 2.0, dtype=np.float32) + + # Need to ensure we have enough points for RANSAC + config.stride = 1 + + # Cam 2: Invalid depth (zeros) + depth_map_invalid = np.zeros((height, width), dtype=np.float32) + + camera_data = { + "cam1": {"depth": depth_map_valid, "K": K}, + "cam2": {"depth": depth_map_invalid, "K": K}, + } + extrinsics = { + "cam1": T_world_cam, + "cam2": T_world_cam, + } + + new_extrinsics, metrics = refine_ground_from_depth(camera_data, extrinsics, config) + + assert metrics.success + assert metrics.num_cameras_valid == 1 + assert metrics.correction_applied + + # Cam 1 should be corrected + assert "cam1" in metrics.camera_corrections + assert "cam1" not in metrics.skipped_cameras + + # Cam 2 should be skipped + assert "cam2" not in metrics.camera_corrections + assert "cam2" in metrics.skipped_cameras + + # Cam 2 extrinsics should be unchanged + np.testing.assert_array_equal(new_extrinsics["cam2"], extrinsics["cam2"]) + + # Cam 1 extrinsics should be changed + assert not np.array_equal(new_extrinsics["cam1"], extrinsics["cam1"])