161 lines
4.6 KiB
Python
161 lines
4.6 KiB
Python
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()
|