feat(calibration): robust depth refinement pipeline with diagnostics and benchmarking

py_workspace/tests/test_depth_refine.py

@@ -37,6 +37,14 @@ def test_refine_extrinsics_with_depth_no_change():
     # np.testing.assert_allclose(T_initial, T_refined, atol=1e-5)
     # assert stats["success"] is True
     assert stats["final_cost"] <= stats["initial_cost"] + 1e-10
+    assert "termination_status" in stats
+    assert "nfev" in stats
+    assert "optimality" in stats
+    assert "n_active_bounds" in stats
+    assert "n_depth_valid" in stats
+    assert "n_points_total" in stats
+    assert "loss_function" in stats
+    assert "f_scale" in stats
 
 
 def test_refine_extrinsics_with_depth_with_offset():
@@ -95,48 +103,50 @@ def test_refine_extrinsics_respects_bounds():
 
 def test_robust_loss_handles_outliers():
     K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
-    
+
     # True pose: camera moved 0.1m forward
     T_true = np.eye(4)
     T_true[2, 3] = 0.1
-    
+
     # Initial pose: identity
     T_initial = np.eye(4)
-    
+
     # Create synthetic depth map
     # Marker at (0,0,2.1) in world -> (0,0,2.0) in camera (since cam moved 0.1 forward)
     depth_map = np.full((720, 1280), 2.0, dtype=np.float32)
-    
+
     # Add outliers: 30% of pixels are garbage (e.g. 0.5m or 5.0m)
     # We'll simulate this by having multiple markers, some with bad depth
     marker_corners_world = {}
-    
+
     # 7 good markers (depth 2.0)
     # 3 bad markers (depth 5.0 - huge outlier)
-    
+
     # We need to ensure these project to unique pixels.
     # K = 1000 focal.
     # x = 0.1 * i. Z = 2.1 (world).
     # u = 1000 * x / Z + 640
-    
+
     marker_corners_world[0] = []
-    
+
     for i in range(10):
         u = int(50 * i + 640)
         v = 360
-        
+
         world_pt = np.array([0.1 * i, 0, 2.1])
         marker_corners_world[0].append(world_pt)
-        
+
         # Paint a wide strip to cover T_initial to T_true movement
         # u_initial = 47.6 * i + 640. u_true = 50 * i + 640.
         # Diff is ~2.4 * i. Max diff (i=9) is ~22 pixels.
         # So +/- 30 pixels should cover it.
-        
+
         if i < 7:
-            depth_map[v-5:v+6, u-30:u+31] = 2.0  # Good measurement
+            depth_map[v - 5 : v + 6, u - 30 : u + 31] = 2.0  # Good measurement
         else:
-            depth_map[v-5:v+6, u-30:u+31] = 5.0  # Outlier measurement (3m error)
+            depth_map[v - 5 : v + 6, u - 30 : u + 31] = (
+                5.0  # Outlier measurement (3m error)
+            )
 
     marker_corners_world[0] = np.array(marker_corners_world[0])
 
@@ -148,15 +158,17 @@ def test_robust_loss_handles_outliers():
         K,
         max_translation_m=0.2,
         max_rotation_deg=5.0,
-        regularization_weight=0.0, # Disable reg to see if data term wins
+        regularization_weight=0.0,  # Disable reg to see if data term wins
         loss="soft_l1",
-        f_scale=0.1
+        f_scale=0.1,
     )
-    
+
     # With robust loss, it should ignore the 3m errors and converge to the 0.1m shift
     # The 0.1m shift explains the 7 inliers perfectly.
     # T_refined[2, 3] should be close to 0.1
-    assert abs(T_refined[2, 3] - 0.1) < 0.02  # Allow small error due to outliers pulling slightly
+    assert (
+        abs(T_refined[2, 3] - 0.1) < 0.02
+    )  # Allow small error due to outliers pulling slightly
     assert stats["success"] is True
 
     # Run with linear loss (MSE) - should fail or be pulled significantly
@@ -168,14 +180,61 @@ def test_robust_loss_handles_outliers():
         max_translation_m=0.2,
         max_rotation_deg=5.0,
         regularization_weight=0.0,
-        loss="linear"
+        loss="linear",
     )
-    
+
     # MSE will try to average 0.0 error (7 points) and 3.0 error (3 points)
     # Mean error target ~ 0.9m
     # So it will likely pull the camera way back to reduce the 3m errors
     # The result should be WORSE than the robust one
     error_robust = abs(T_refined[2, 3] - 0.1)
     error_mse = abs(T_refined_mse[2, 3] - 0.1)
-    
+
     assert error_robust < error_mse
+
+
+def test_refine_with_confidence_weights():
+    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
+    T_initial = np.eye(4)
+
+    # 2 points: one with good depth, one with bad depth but low confidence
+    # Point 1: World (0,0,2.1), Depth 2.0 (True shift 0.1)
+    # Point 2: World (0.5,0,2.1), Depth 5.0 (Outlier)
+    marker_corners_world = {1: np.array([[0, 0, 2.1], [0.5, 0, 2.1]])}
+    depth_map = np.full((720, 1280), 2.0, dtype=np.float32)
+    # Paint outlier depth
+    depth_map[360, int(1000 * 0.5 / 2.1 + 640)] = 5.0
+
+    # Confidence map: Point 1 is confident (1), Point 2 is NOT confident (90)
+    confidence_map = np.full((720, 1280), 1.0, dtype=np.float32)
+    confidence_map[360, int(1000 * 0.5 / 2.1 + 640)] = 90.0
+
+    # 1. Without weights: Outlier should pull the result significantly
+    T_no_weights, stats_no_weights = refine_extrinsics_with_depth(
+        T_initial,
+        marker_corners_world,
+        depth_map,
+        K,
+        regularization_weight=0.0,
+        confidence_map=None,
+        loss="linear",  # Use linear to make weighting effect more obvious
+    )
+
+    # 2. With weights: Outlier should be suppressed
+    T_weighted, stats_weighted = refine_extrinsics_with_depth(
+        T_initial,
+        marker_corners_world,
+        depth_map,
+        K,
+        regularization_weight=0.0,
+        confidence_map=confidence_map,
+        confidence_thresh=100.0,
+        loss="linear",
+    )
+
+    error_no_weights = abs(T_no_weights[2, 3] - 0.1)
+    error_weighted = abs(T_weighted[2, 3] - 0.1)
+
+    # Weighted error should be much smaller because the 5.0 depth was suppressed
+    assert error_weighted < error_no_weights
+    assert error_weighted < 0.06