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

py_workspace/aruco/depth_refine.py

@@ -1,8 +1,12 @@
 import numpy as np
-from typing import Dict, Tuple, Any
+from typing import Dict, Tuple, Any, Optional
 from scipy.optimize import least_squares
 from .pose_math import rvec_tvec_to_matrix, matrix_to_rvec_tvec
-from .depth_verify import compute_depth_residual
+from .depth_verify import (
+    compute_depth_residual,
+    get_confidence_weight,
+    project_point_to_pixel,
+)
 
 
 def extrinsics_to_params(T: np.ndarray) -> np.ndarray:
@@ -24,6 +28,8 @@ def depth_residuals(
     initial_params: np.ndarray,
     reg_rot: float = 0.1,
     reg_trans: float = 1.0,
+    confidence_map: Optional[np.ndarray] = None,
+    confidence_thresh: float = 100.0,
 ) -> np.ndarray:
     T = params_to_extrinsics(params)
     residuals = []
@@ -32,15 +38,25 @@ def depth_residuals(
         for corner in corners:
             residual = compute_depth_residual(corner, T, depth_map, K, window_size=5)
             if residual is not None:
+                if confidence_map is not None:
+                    u, v = project_point_to_pixel(
+                        (np.linalg.inv(T) @ np.append(corner, 1.0))[:3], K
+                    )
+                    if u is not None and v is not None:
+                        h, w = confidence_map.shape[:2]
+                        if 0 <= u < w and 0 <= v < h:
+                            conf = confidence_map[v, u]
+                            weight = get_confidence_weight(conf, confidence_thresh)
+                            residual *= np.sqrt(weight)
                 residuals.append(residual)
 
     # Regularization as pseudo-residuals
     param_diff = params - initial_params
-    
+
     # Rotation regularization (first 3 params)
     if reg_rot > 0:
         residuals.extend(param_diff[:3] * reg_rot)
-        
+
     # Translation regularization (last 3 params)
     if reg_trans > 0:
         residuals.extend(param_diff[3:] * reg_trans)
@@ -60,6 +76,8 @@ def refine_extrinsics_with_depth(
     f_scale: float = 0.1,
     reg_rot: float | None = None,
     reg_trans: float | None = None,
+    confidence_map: Optional[np.ndarray] = None,
+    confidence_thresh: float = 100.0,
 ) -> Tuple[np.ndarray, dict[str, Any]]:
     initial_params = extrinsics_to_params(T_initial)
 
@@ -72,14 +90,29 @@ def refine_extrinsics_with_depth(
         reg_trans = regularization_weight * 10.0
 
     # Check for valid depth points first
-    data_residual_count = 0
+    n_points_total = 0
+    n_depth_valid = 0
+    n_confidence_rejected = 0
+
     for marker_id, corners in marker_corners_world.items():
         for corner in corners:
+            n_points_total += 1
             res = compute_depth_residual(corner, T_initial, depth_map, K, window_size=5)
             if res is not None:
-                data_residual_count += 1
-    
-    if data_residual_count == 0:
+                n_depth_valid += 1
+                if confidence_map is not None:
+                    u, v = project_point_to_pixel(
+                        (np.linalg.inv(T_initial) @ np.append(corner, 1.0))[:3], K
+                    )
+                    if u is not None and v is not None:
+                        h, w = confidence_map.shape[:2]
+                        if 0 <= u < w and 0 <= v < h:
+                            conf = confidence_map[v, u]
+                            weight = get_confidence_weight(conf, confidence_thresh)
+                            if weight <= 0:
+                                n_confidence_rejected += 1
+
+    if n_depth_valid == 0:
         return T_initial, {
             "success": False,
             "reason": "no_valid_depth_points",
@@ -89,22 +122,30 @@ def refine_extrinsics_with_depth(
             "delta_rotation_deg": 0.0,
             "delta_translation_norm_m": 0.0,
             "termination_message": "No valid depth points found at marker corners",
+            "termination_status": -1,
             "nfev": 0,
+            "njev": 0,
             "optimality": 0.0,
+            "n_active_bounds": 0,
             "active_mask": np.zeros(6, dtype=int),
-            "cost": 0.0
+            "cost": 0.0,
+            "n_points_total": n_points_total,
+            "n_depth_valid": n_depth_valid,
+            "n_confidence_rejected": n_confidence_rejected,
+            "loss_function": loss,
+            "f_scale": f_scale,
         }
 
     max_rotation_rad = np.deg2rad(max_rotation_deg)
 
     lower_bounds = initial_params.copy()
     upper_bounds = initial_params.copy()
-    
+
     lower_bounds[:3] -= max_rotation_rad
     upper_bounds[:3] += max_rotation_rad
     lower_bounds[3:] -= max_translation_m
     upper_bounds[3:] += max_translation_m
-    
+
     bounds = (lower_bounds, upper_bounds)
 
     result = least_squares(
@@ -117,6 +158,8 @@ def refine_extrinsics_with_depth(
             initial_params,
             reg_rot,
             reg_trans,
+            confidence_map,
+            confidence_thresh,
         ),
         method="trf",
         loss=loss,
@@ -142,6 +185,8 @@ def refine_extrinsics_with_depth(
         initial_params,
         reg_rot,
         reg_trans,
+        confidence_map,
+        confidence_thresh,
     )
     initial_cost = 0.5 * np.sum(initial_residuals**2)
 
@@ -153,10 +198,20 @@ def refine_extrinsics_with_depth(
         "delta_rotation_deg": float(delta_rotation_deg),
         "delta_translation_norm_m": float(delta_translation),
         "termination_message": result.message,
-        "nfev": result.nfev,
+        "termination_status": int(result.status),
+        "nfev": int(result.nfev),
+        "njev": int(getattr(result, "njev", 0)),
         "optimality": float(result.optimality),
-        "active_mask": result.active_mask,
+        "n_active_bounds": int(np.sum(result.active_mask != 0)),
+        "active_mask": result.active_mask.tolist()
+        if hasattr(result.active_mask, "tolist")
+        else result.active_mask,
         "cost": float(result.cost),
+        "n_points_total": n_points_total,
+        "n_depth_valid": n_depth_valid,
+        "n_confidence_rejected": n_confidence_rejected,
+        "loss_function": loss,
+        "f_scale": f_scale,
     }
 
     return T_refined, stats