feat(calibrate): integrate multi-frame depth pooling with --depth-pool-size flag

py_workspace/tests/test_depth_pool_integration.py

@@ -0,0 +1,253 @@
+import pytest
+import numpy as np
+from unittest.mock import MagicMock, patch
+import sys
+from pathlib import Path
+
+# Add py_workspace to path
+sys.path.append(str(Path(__file__).parent.parent))
+
+from calibrate_extrinsics import apply_depth_verify_refine_postprocess
+
+
+@pytest.fixture
+def mock_dependencies():
+    with (
+        patch("calibrate_extrinsics.verify_extrinsics_with_depth") as mock_verify,
+        patch("calibrate_extrinsics.refine_extrinsics_with_depth") as mock_refine,
+        patch("calibrate_extrinsics.click.echo") as mock_echo,
+    ):
+        # Setup mock return values
+        mock_verify_res = MagicMock()
+        mock_verify_res.rmse = 0.05
+        mock_verify_res.mean_abs = 0.04
+        mock_verify_res.median = 0.03
+        mock_verify_res.depth_normalized_rmse = 0.02
+        mock_verify_res.n_valid = 100
+        mock_verify_res.n_total = 120
+        mock_verify_res.residuals = []
+        mock_verify.return_value = mock_verify_res
+
+        mock_refine.return_value = (np.eye(4), {"success": True})
+
+        yield mock_verify, mock_refine, mock_echo
+
+
+def test_pool_size_1_equivalence(mock_dependencies):
+    """
+    Regression test: Ensure pool_size=1 behaves exactly like the old single-frame path.
+    """
+    mock_verify, _, _ = mock_dependencies
+
+    serial = "123456"
+    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}
+
+    # Create a frame with specific depth values
+    depth_map = np.ones((10, 10)) * 2.0
+    conf_map = np.zeros((10, 10))
+
+    frame_mock = MagicMock()
+    frame_mock.depth_map = depth_map
+    frame_mock.confidence_map = conf_map
+
+    vf = {
+        "frame": frame_mock,
+        "ids": np.array([[1]]),
+        "corners": np.zeros((1, 4, 2)),
+        "score": 100.0,
+    }
+
+    # Structure for new implementation: list of frames
+    verification_frames = {serial: [vf]}
+    marker_geometry = {1: np.zeros((4, 3))}
+    camera_matrices = {serial: np.eye(3)}
+
+    # Run with pool_size=1
+    apply_depth_verify_refine_postprocess(
+        results=results,
+        verification_frames=verification_frames,
+        marker_geometry=marker_geometry,
+        camera_matrices=camera_matrices,
+        verify_depth=True,
+        refine_depth=False,
+        use_confidence_weights=False,
+        depth_confidence_threshold=50,
+        depth_pool_size=1,
+    )
+
+    # Verify that verify_extrinsics_with_depth was called with the exact depth map from the frame
+    args, _ = mock_verify.call_args
+    passed_depth_map = args[2]
+
+    np.testing.assert_array_equal(passed_depth_map, depth_map)
+    assert passed_depth_map is depth_map
+
+
+def test_pool_size_5_integration(mock_dependencies):
+    """
+    Test that pool_size > 1 actually calls pooling and uses the result.
+    """
+    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 3 frames with different depth values
+    # Frame 1: 2.0m
+    # Frame 2: 2.2m
+    # Frame 3: 1.8m
+    # Median should be 2.0m
+
+    frames = []
+    for d in [2.0, 2.2, 1.8]:
+        f = MagicMock()
+        f.depth_map = np.ones((10, 10)) * d
+        f.confidence_map = np.zeros((10, 10))
+        frames.append(f)
+
+    vfs = []
+    for i, f in enumerate(frames):
+        vfs.append(
+            {
+                "frame": f,
+                "ids": np.array([[1]]),
+                "corners": np.zeros((1, 4, 2)),
+                "score": 100.0 - i,
+            }
+        )
+
+    verification_frames = {serial: vfs}
+    marker_geometry = {1: np.zeros((4, 3))}
+    camera_matrices = {serial: np.eye(3)}
+
+    # Run with pool_size=3
+    apply_depth_verify_refine_postprocess(
+        results=results,
+        verification_frames=verification_frames,
+        marker_geometry=marker_geometry,
+        camera_matrices=camera_matrices,
+        verify_depth=True,
+        refine_depth=False,
+        use_confidence_weights=False,
+        depth_confidence_threshold=50,
+        depth_pool_size=3,
+    )
+
+    # Check that "Using pooled depth" was logged
+    any_pooled = any(
+        "Using pooled depth" in str(call.args[0]) for call in mock_echo.call_args_list
+    )
+    assert any_pooled
+
+    # Check that the depth map passed to verify is the median (2.0)
+    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
+    assert "depth_pool" in results[serial]
+    assert results[serial]["depth_pool"]["pooled"] is True
+    assert results[serial]["depth_pool"]["pool_size_actual"] == 3
+
+
+def test_pool_fallback_insufficient_valid(mock_dependencies):
+    """
+    Test fallback to single frame when pooled result has too few valid points.
+    """
+    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"}}
+
+    # Frame 1: Good depth
+    f1 = MagicMock()
+    f1.depth_map = np.ones((10, 10)) * 2.0
+    f1.confidence_map = np.zeros((10, 10))
+
+    # Frame 2: NaN depth (simulating misalignment or noise)
+    f2 = MagicMock()
+    f2.depth_map = np.full((10, 10), np.nan)
+    f2.confidence_map = np.zeros((10, 10))
+
+    # Frame 3: NaN depth
+    f3 = MagicMock()
+    f3.depth_map = np.full((10, 10), np.nan)
+    f3.confidence_map = np.zeros((10, 10))
+
+    # With median pooling, if >50% are NaN, result is NaN (standard median behavior with NaNs usually propagates or ignores)
+    # Our pool_depth_maps uses nanmedian, which ignores NaNs.
+    # But if we have [2.0, NaN, NaN], median of [2.0] is 2.0.
+    # Wait, let's make it so they are valid but inconsistent to cause variance?
+    # Or just force the pooled result to be bad by making them all different and sparse?
+
+    # Let's use the fact that we can patch pool_depth_maps in the test!
+    with patch("calibrate_extrinsics.pool_depth_maps") as mock_pool:
+        # Return empty/invalid map
+        mock_pool.return_value = (
+            np.zeros((10, 10)),
+            None,
+        )  # Zeros are invalid depth (<=0)
+
+        # Frame 1: Valid on left half
+        d1 = np.full((10, 10), np.nan)
+        d1[:, :5] = 2.0
+        f1.depth_map = d1
+        f1.confidence_map = np.zeros((10, 10))
+
+        # Frame 2: Valid on right half
+        d2 = np.full((10, 10), np.nan)
+        d2[:, 5:] = 2.0
+        f2.depth_map = d2
+        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)}
+
+        apply_depth_verify_refine_postprocess(
+            results=results,
+            verification_frames=verification_frames,
+            marker_geometry=marker_geometry,
+            camera_matrices=camera_matrices,
+            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(
+            "Falling back to best single frame" in str(call.args[0])
+            for call in mock_echo.call_args_list
+        )
+        assert any_fallback
+
+        # Verify we used the best frame (f1)
+        args, _ = mock_verify.call_args
+        passed_depth_map = args[2]
+        assert passed_depth_map is d1
+
+        # Verify metadata
+        assert results[serial]["depth_pool"]["pooled"] is False
+        assert (
+            results[serial]["depth_pool"]["fallback_reason"]
+            == "insufficient_valid_points"
+        )