crosstyan
43e736faac
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
145 lines
4.2 KiB
Python
145 lines
4.2 KiB
Python
import cv2
|
|
import numpy as np
|
|
import pyzed.sl as sl
|
|
from typing import Optional, Tuple, Dict
|
|
from .pose_math import compute_reprojection_error
|
|
|
|
|
|
def create_detector(
|
|
dictionary_id: int = cv2.aruco.DICT_4X4_50,
|
|
params: Optional[cv2.aruco.DetectorParameters] = None,
|
|
) -> cv2.aruco.ArucoDetector:
|
|
"""
|
|
Creates an ArUco detector with the specified dictionary and parameters.
|
|
|
|
Args:
|
|
dictionary_id: The ArUco dictionary ID.
|
|
params: Optional detector parameters. If None, default parameters are used.
|
|
|
|
Returns:
|
|
cv2.aruco.ArucoDetector: The configured ArUco detector.
|
|
"""
|
|
dictionary = cv2.aruco.getPredefinedDictionary(dictionary_id)
|
|
if params is None:
|
|
params = cv2.aruco.DetectorParameters()
|
|
return cv2.aruco.ArucoDetector(dictionary, params)
|
|
|
|
|
|
def detect_markers(
|
|
image: np.ndarray, detector: cv2.aruco.ArucoDetector
|
|
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
|
|
"""
|
|
Detects ArUco markers in an image.
|
|
|
|
Args:
|
|
image: Input image (BGR or grayscale).
|
|
detector: The ArUco detector to use.
|
|
|
|
Returns:
|
|
Tuple[np.ndarray, Optional[np.ndarray]]:
|
|
- corners: Detected corners as float32 shape (N, 4, 2).
|
|
- ids: Detected marker IDs as int shape (N,) or None if no markers detected.
|
|
"""
|
|
if image is None:
|
|
raise ValueError("Input image is None")
|
|
|
|
if len(image.shape) == 3:
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
else:
|
|
gray = image
|
|
|
|
corners, ids, _ = detector.detectMarkers(gray)
|
|
|
|
if ids is not None:
|
|
# corners is a list of (1, 4, 2) arrays, reshape to (N, 4, 2)
|
|
corners_array = np.array(corners, dtype=np.float32).reshape(-1, 4, 2)
|
|
ids_array = ids.flatten().astype(np.int32)
|
|
return corners_array, ids_array
|
|
else:
|
|
return np.empty((0, 4, 2), dtype=np.float32), None
|
|
|
|
|
|
def build_camera_matrix_from_zed(
|
|
cam: sl.Camera, resizer: Optional[sl.Resolution] = None
|
|
) -> np.ndarray:
|
|
"""
|
|
Builds a 3x3 camera intrinsic matrix from a ZED camera's calibration parameters.
|
|
|
|
Args:
|
|
cam: The ZED camera object.
|
|
resizer: Optional resolution to get parameters for.
|
|
|
|
Returns:
|
|
np.ndarray: 3x3 float64 camera matrix K.
|
|
"""
|
|
cam_params = cam.get_camera_information(
|
|
resizer
|
|
).camera_configuration.calibration_parameters.left_cam
|
|
fx = cam_params.fx
|
|
fy = cam_params.fy
|
|
cx = cam_params.cx
|
|
cy = cam_params.cy
|
|
|
|
K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
|
|
|
|
return K
|
|
|
|
|
|
def estimate_pose_from_detections(
|
|
corners: np.ndarray,
|
|
ids: np.ndarray,
|
|
marker_geometry: Dict[int, np.ndarray],
|
|
K: np.ndarray,
|
|
dist: Optional[np.ndarray] = None,
|
|
min_markers: int = 4,
|
|
pnp_flag: int = cv2.SOLVEPNP_SQPNP,
|
|
) -> Optional[Tuple[np.ndarray, np.ndarray, float, int]]:
|
|
"""
|
|
Estimates the camera pose from detected ArUco markers.
|
|
|
|
Args:
|
|
corners: Detected corners (N, 4, 2).
|
|
ids: Detected marker IDs (N,).
|
|
marker_geometry: Dictionary mapping marker ID to its 3D corners (4, 3).
|
|
K: 3x3 camera matrix.
|
|
dist: Optional distortion coefficients.
|
|
min_markers: Minimum number of markers required for pose estimation.
|
|
pnp_flag: OpenCV solvePnP flag.
|
|
|
|
Returns:
|
|
Optional[Tuple[np.ndarray, np.ndarray, float, int]]:
|
|
(rvec, tvec, reproj_err, n_markers) or None if estimation fails or not enough markers.
|
|
"""
|
|
if ids is None or len(ids) == 0:
|
|
return None
|
|
|
|
obj_pts = []
|
|
img_pts = []
|
|
n_markers = 0
|
|
|
|
for i, marker_id in enumerate(ids):
|
|
if marker_id in marker_geometry:
|
|
obj_pts.append(marker_geometry[marker_id])
|
|
img_pts.append(corners[i])
|
|
n_markers += 1
|
|
|
|
if n_markers < min_markers:
|
|
return None
|
|
|
|
# Stack points into (M*4, 3) and (M*4, 2)
|
|
obj_pts_array = np.vstack(obj_pts).astype(np.float32)
|
|
img_pts_array = np.vstack(img_pts).astype(np.float32)
|
|
|
|
success, rvec, tvec = cv2.solvePnP(
|
|
obj_pts_array, img_pts_array, K, dist, flags=pnp_flag
|
|
)
|
|
|
|
if not success:
|
|
return None
|
|
|
|
reproj_err = compute_reprojection_error(
|
|
obj_pts_array, img_pts_array, rvec, tvec, K, dist
|
|
)
|
|
|
|
return rvec, tvec, reproj_err, n_markers
|