163 lines
4.8 KiB
Python
163 lines
4.8 KiB
Python
import numpy as np
|
|
from typing import Dict, Tuple, Any
|
|
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
|
|
|
|
|
|
def extrinsics_to_params(T: np.ndarray) -> np.ndarray:
|
|
rvec, tvec = matrix_to_rvec_tvec(T)
|
|
return np.concatenate([rvec, tvec])
|
|
|
|
|
|
def params_to_extrinsics(params: np.ndarray) -> np.ndarray:
|
|
rvec = params[:3]
|
|
tvec = params[3:]
|
|
return rvec_tvec_to_matrix(rvec, tvec)
|
|
|
|
|
|
def depth_residuals(
|
|
params: np.ndarray,
|
|
marker_corners_world: Dict[int, np.ndarray],
|
|
depth_map: np.ndarray,
|
|
K: np.ndarray,
|
|
initial_params: np.ndarray,
|
|
reg_rot: float = 0.1,
|
|
reg_trans: float = 1.0,
|
|
) -> np.ndarray:
|
|
T = params_to_extrinsics(params)
|
|
residuals = []
|
|
|
|
for marker_id, corners in marker_corners_world.items():
|
|
for corner in corners:
|
|
residual = compute_depth_residual(corner, T, depth_map, K, window_size=5)
|
|
if residual is not None:
|
|
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)
|
|
|
|
return np.array(residuals)
|
|
|
|
|
|
def refine_extrinsics_with_depth(
|
|
T_initial: np.ndarray,
|
|
marker_corners_world: Dict[int, np.ndarray],
|
|
depth_map: np.ndarray,
|
|
K: np.ndarray,
|
|
max_translation_m: float = 0.05,
|
|
max_rotation_deg: float = 5.0,
|
|
regularization_weight: float = 0.1,
|
|
loss: str = "soft_l1",
|
|
f_scale: float = 0.1,
|
|
reg_rot: float | None = None,
|
|
reg_trans: float | None = None,
|
|
) -> Tuple[np.ndarray, dict[str, Any]]:
|
|
initial_params = extrinsics_to_params(T_initial)
|
|
|
|
# Handle legacy regularization_weight if specific weights aren't provided
|
|
# Default behavior: reg_rot = weight, reg_trans = weight * 10
|
|
# This matches the plan's default of (0.1, 1.0) when weight is 0.1
|
|
if reg_rot is None:
|
|
reg_rot = regularization_weight
|
|
if reg_trans is None:
|
|
reg_trans = regularization_weight * 10.0
|
|
|
|
# Check for valid depth points first
|
|
data_residual_count = 0
|
|
for marker_id, corners in marker_corners_world.items():
|
|
for corner in corners:
|
|
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:
|
|
return T_initial, {
|
|
"success": False,
|
|
"reason": "no_valid_depth_points",
|
|
"initial_cost": 0.0,
|
|
"final_cost": 0.0,
|
|
"iterations": 0,
|
|
"delta_rotation_deg": 0.0,
|
|
"delta_translation_norm_m": 0.0,
|
|
"termination_message": "No valid depth points found at marker corners",
|
|
"nfev": 0,
|
|
"optimality": 0.0,
|
|
"active_mask": np.zeros(6, dtype=int),
|
|
"cost": 0.0
|
|
}
|
|
|
|
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(
|
|
depth_residuals,
|
|
initial_params,
|
|
args=(
|
|
marker_corners_world,
|
|
depth_map,
|
|
K,
|
|
initial_params,
|
|
reg_rot,
|
|
reg_trans,
|
|
),
|
|
method="trf",
|
|
loss=loss,
|
|
f_scale=f_scale,
|
|
bounds=bounds,
|
|
x_scale="jac",
|
|
max_nfev=200,
|
|
)
|
|
|
|
T_refined = params_to_extrinsics(result.x)
|
|
|
|
delta_params = result.x - initial_params
|
|
delta_rotation_rad = np.linalg.norm(delta_params[:3])
|
|
delta_rotation_deg = np.rad2deg(delta_rotation_rad)
|
|
delta_translation = np.linalg.norm(delta_params[3:])
|
|
|
|
# Calculate initial cost for comparison
|
|
initial_residuals = depth_residuals(
|
|
initial_params,
|
|
marker_corners_world,
|
|
depth_map,
|
|
K,
|
|
initial_params,
|
|
reg_rot,
|
|
reg_trans,
|
|
)
|
|
initial_cost = 0.5 * np.sum(initial_residuals**2)
|
|
|
|
stats = {
|
|
"iterations": result.nfev,
|
|
"success": result.success,
|
|
"initial_cost": float(initial_cost),
|
|
"final_cost": float(result.cost),
|
|
"delta_rotation_deg": float(delta_rotation_deg),
|
|
"delta_translation_norm_m": float(delta_translation),
|
|
"termination_message": result.message,
|
|
"nfev": result.nfev,
|
|
"optimality": float(result.optimality),
|
|
"active_mask": result.active_mask,
|
|
"cost": float(result.cost),
|
|
}
|
|
|
|
return T_refined, stats
|