feat: add RMSE-based fallback for depth pooling

2026-02-07 09:07:23 +00:00
parent 8f6aee7f22
commit 8dbf892ce8
2 changed files with 153 additions and 11 deletions
@@ -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.",
@@ -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