docs: update README and finalize ground plane refinement integration

2026-02-09 08:30:50 +00:00
parent 9d9e95de81
commit ae3c9eba76
4 changed files with 221 additions and 0 deletions
@@ -0,0 +1,160 @@
+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,
+        patch("calibrate_extrinsics.save_depth_data") as mock_save_depth,
+    ):
+        # 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, mock_save_depth
+
+
+def test_save_depth_data_integration(mock_dependencies):
+    """
+    Test that save_depth_data is called with correct data when save_depth_path is provided.
+    """
+    mock_verify, _, _, mock_save_depth = mock_dependencies
+
+    serial = "123456"
+    serial_int = int(serial)
+    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}
+
+    # Create frames
+    # Frame 1
+    f1 = MagicMock()
+    f1.depth_map = np.ones((10, 20)) * 2.0  # H=10, W=20
+    f1.confidence_map = np.zeros((10, 20))
+
+    # Frame 2
+    f2 = MagicMock()
+    f2.depth_map = np.ones((10, 20)) * 2.2
+    f2.confidence_map = np.zeros((10, 20))
+
+    vfs = [
+        {
+            "frame": f1,
+            "ids": np.array([[1]]),
+            "corners": np.zeros((1, 4, 2)),
+            "score": 100,
+            "frame_index": 10,
+        },
+        {
+            "frame": f2,
+            "ids": np.array([[1]]),
+            "corners": np.zeros((1, 4, 2)),
+            "score": 90,
+            "frame_index": 20,
+        },
+    ]
+
+    verification_frames = {serial_int: vfs}
+    marker_geometry = {1: np.zeros((4, 3))}
+    camera_matrices = {serial_int: np.eye(3)}
+
+    save_path = "output_depth.h5"
+
+    # Run with save_depth_path
+    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,
+        save_depth_path=save_path,
+    )
+
+    # Verify save_depth_data was called
+    mock_save_depth.assert_called_once()
+
+    # Check arguments
+    call_args = mock_save_depth.call_args
+    assert call_args[0][0] == save_path
+
+    camera_data = call_args[0][1]
+    assert serial in camera_data
+
+    cam_data = camera_data[serial]
+    assert "intrinsics" in cam_data
+    assert "resolution" in cam_data
+    assert cam_data["resolution"] == (20, 10)  # (W, H)
+
+    assert "pooled_depth" in cam_data
+    assert "pooled_confidence" in cam_data
+    assert "pool_metadata" in cam_data
+    assert "raw_frames" in cam_data
+
+    assert len(cam_data["raw_frames"]) == 2
+    assert cam_data["raw_frames"][0]["frame_index"] == 10
+    assert cam_data["raw_frames"][1]["frame_index"] == 20
+
+    # Check that raw frames contain depth maps
+    assert "depth_map" in cam_data["raw_frames"][0]
+    np.testing.assert_array_equal(cam_data["raw_frames"][0]["depth_map"], f1.depth_map)
+
+
+def test_save_depth_skipped_when_no_path(mock_dependencies):
+    """
+    Test that save_depth_data is NOT called when save_depth_path is None.
+    """
+    _, _, _, mock_save_depth = mock_dependencies
+
+    serial = "123456"
+    serial_int = int(serial)
+    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}
+
+    f1 = MagicMock()
+    f1.depth_map = np.ones((10, 10))
+    f1.confidence_map = np.zeros((10, 10))
+
+    vfs = [
+        {
+            "frame": f1,
+            "ids": np.array([[1]]),
+            "corners": np.zeros((1, 4, 2)),
+            "score": 100,
+            "frame_index": 10,
+        }
+    ]
+
+    apply_depth_verify_refine_postprocess(
+        results=results,
+        verification_frames={serial_int: vfs},
+        marker_geometry={1: np.zeros((4, 3))},
+        camera_matrices={serial_int: np.eye(3)},
+        verify_depth=True,
+        refine_depth=False,
+        use_confidence_weights=False,
+        depth_confidence_threshold=50,
+        save_depth_path=None,
+    )
+
+    mock_save_depth.assert_not_called()