From 8dbf892ce8607418aac03231447627c25f9b0fcf Mon Sep 17 00:00:00 2001 From: crosstyan Date: Sat, 7 Feb 2026 09:07:23 +0000 Subject: [PATCH] feat: add RMSE-based fallback for depth pooling --- py_workspace/calibrate_extrinsics.py | 72 +++++++++++++-- .../tests/test_depth_pool_integration.py | 92 ++++++++++++++++++- 2 files changed, 153 insertions(+), 11 deletions(-) diff --git a/py_workspace/calibrate_extrinsics.py b/py_workspace/calibrate_extrinsics.py index c8b2fc1..0f4e1f4 100644 --- a/py_workspace/calibrate_extrinsics.py +++ b/py_workspace/calibrate_extrinsics.py @@ -199,9 +199,9 @@ def apply_depth_verify_refine_postprocess( # If pooled result is much worse (e.g. < 50% of valid points of single frame), fallback # This can happen if frames are misaligned or pooling logic fails - if n_valid_pooled < (n_valid_best * 0.5): + if n_valid_pooled < n_valid_best: click.echo( - f"Camera {serial}: Pooled depth has too few valid points ({n_valid_pooled} vs {n_valid_best}). " + f"Camera {serial}: Pooled depth has fewer valid points ({n_valid_pooled} vs {n_valid_best}). " "Falling back to best single frame." ) final_depth = best_depth @@ -213,16 +213,68 @@ def apply_depth_verify_refine_postprocess( "fallback_reason": "insufficient_valid_points", } else: - final_depth = pooled_depth - final_conf = pooled_conf - pool_metadata = { - "pool_size_requested": depth_pool_size, - "pool_size_actual": len(depth_maps), - "pooled": True, + # A/B Test: Compare RMSE of pooled vs best single frame + # We need to compute RMSE for both using the current T_mean + pose_str = results[str(serial)]["pose"] + T_mean = np.fromstring(pose_str, sep=" ").reshape(4, 4) + cam_matrix = camera_matrices[serial] + + marker_corners_world = { + int(mid): marker_geometry[int(mid)] + for mid in ids.flatten() + if int(mid) in marker_geometry } - click.echo( - f"Camera {serial}: Using pooled depth from {len(depth_maps)} frames." + + # Verify pooled + verify_pooled = verify_extrinsics_with_depth( + T_mean, + marker_corners_world, + pooled_depth, + cam_matrix, + confidence_map=pooled_conf, + confidence_thresh=depth_confidence_threshold, ) + + # Verify best single frame + verify_best = verify_extrinsics_with_depth( + T_mean, + marker_corners_world, + best_depth, + cam_matrix, + confidence_map=best_conf, + confidence_thresh=depth_confidence_threshold, + ) + + # If pooled RMSE is worse than best single frame, fallback + # We use a small epsilon to avoid flipping on noise, but generally strict + if verify_pooled.rmse > verify_best.rmse: + click.echo( + f"Camera {serial}: Pooled depth RMSE ({verify_pooled.rmse:.4f}m) worse than single frame ({verify_best.rmse:.4f}m). " + "Falling back to best single frame." + ) + final_depth = best_depth + final_conf = best_conf + pool_metadata = { + "pool_size_requested": depth_pool_size, + "pool_size_actual": len(depth_maps), + "pooled": False, + "fallback_reason": "worse_verify_rmse", + "pooled_rmse": verify_pooled.rmse, + "single_rmse": verify_best.rmse, + } + else: + final_depth = pooled_depth + final_conf = pooled_conf + pool_metadata = { + "pool_size_requested": depth_pool_size, + "pool_size_actual": len(depth_maps), + "pooled": True, + "pooled_rmse": verify_pooled.rmse, + "single_rmse": verify_best.rmse, + } + click.echo( + f"Camera {serial}: Using pooled depth from {len(depth_maps)} frames (RMSE {verify_pooled.rmse:.4f}m vs {verify_best.rmse:.4f}m)." + ) except Exception as e: click.echo( f"Camera {serial}: Pooling failed with error: {e}. Falling back to single frame.", diff --git a/py_workspace/tests/test_depth_pool_integration.py b/py_workspace/tests/test_depth_pool_integration.py index a25c48a..61915dc 100644 --- a/py_workspace/tests/test_depth_pool_integration.py +++ b/py_workspace/tests/test_depth_pool_integration.py @@ -140,13 +140,15 @@ def test_pool_size_5_integration(mock_dependencies): assert any_pooled # Check that the depth map passed to verify is the median (2.0) + # Note: verify is called twice now (once for check, once for final) + # We check the last call args, _ = mock_verify.call_args passed_depth_map = args[2] expected_median = np.ones((10, 10)) * 2.0 np.testing.assert_allclose(passed_depth_map, expected_median) - # Verify metadata was added + # Verify metadata assert "depth_pool" in results[serial] assert results[serial]["depth_pool"]["pooled"] is True # pyright: ignore assert results[serial]["depth_pool"]["pool_size_actual"] == 3 # pyright: ignore @@ -251,3 +253,91 @@ def test_pool_fallback_insufficient_valid(mock_dependencies): results[serial]["depth_pool"]["fallback_reason"] # pyright: ignore == "insufficient_valid_points" ) + + +def test_pool_fallback_worse_rmse(mock_dependencies): + """ + Test fallback to single frame when pooled result has worse RMSE. + """ + mock_verify, _, mock_echo = mock_dependencies + + serial = "123456" + results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}} + + # Create frames + f1 = MagicMock() + f1.depth_map = np.ones((10, 10)) * 2.0 + f1.confidence_map = np.zeros((10, 10)) + + f2 = MagicMock() + f2.depth_map = np.ones((10, 10)) * 2.2 + f2.confidence_map = np.zeros((10, 10)) + + vfs = [ + { + "frame": f1, + "ids": np.array([[1]]), + "corners": np.zeros((1, 4, 2)), + "score": 100, + }, + { + "frame": f2, + "ids": np.array([[1]]), + "corners": np.zeros((1, 4, 2)), + "score": 90, + }, + ] + + verification_frames = {serial: vfs} + marker_geometry = {1: np.zeros((4, 3))} + camera_matrices = {serial: np.eye(3)} + + # Mock verify_extrinsics_with_depth to return worse RMSE for pooled + # First call is pooled, second is best single frame + res_pooled = MagicMock() + res_pooled.rmse = 0.10 + res_pooled.n_valid = 100 + res_pooled.n_total = 100 + + res_best = MagicMock() + res_best.rmse = 0.05 + res_best.n_valid = 100 + res_best.n_total = 100 + + # Third call is the final verification with the chosen depth (best frame) + res_final = MagicMock() + res_final.rmse = 0.05 + res_final.n_valid = 100 + res_final.n_total = 100 + res_final.residuals = [] + + mock_verify.side_effect = [res_pooled, res_best, res_final] + + # Run with pool_size=2 + apply_depth_verify_refine_postprocess( + results=results, # type: ignore + verification_frames=verification_frames, # pyright: ignore + marker_geometry=marker_geometry, + camera_matrices=camera_matrices, # pyright: ignore + verify_depth=True, + refine_depth=False, + use_confidence_weights=False, + depth_confidence_threshold=50, + depth_pool_size=2, + ) + + # Check for fallback message + any_fallback = any( + "worse than single frame" in str(call.args[0]) + for call in mock_echo.call_args_list + ) + assert any_fallback + + # Verify metadata + assert results[serial]["depth_pool"]["pooled"] is False # pyright: ignore + assert ( + results[serial]["depth_pool"]["fallback_reason"] # pyright: ignore + == "worse_verify_rmse" + ) + assert results[serial]["depth_pool"]["pooled_rmse"] == 0.10 # pyright: ignore + assert results[serial]["depth_pool"]["single_rmse"] == 0.05 # pyright: ignore