Add ArUco detection and pose estimation modules

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-opencode) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
2026-02-04 11:44:37 +00:00
parent 0c4e906dcf
commit 43e736faac
13 changed files with 1141 additions and 0 deletions
@@ -0,0 +1,93 @@
+import cv2
+import numpy as np
+
+
+def draw_detected_markers(
+    image: np.ndarray, corners: list[np.ndarray] | np.ndarray, ids: np.ndarray | None
+) -> np.ndarray:
+    """
+    Draws ArUco markers outlines and IDs on the image.
+
+    Args:
+        image: Input image (BGR or grayscale).
+        corners: List of marker corners or np.ndarray.
+        ids: Array of marker IDs.
+
+    Returns:
+        Image with markers drawn.
+    """
+    if image.ndim == 2:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+
+    if corners is not None and len(corners) > 0:
+        # Draw outlines
+        # corners can be list of (1, 4, 2) or np.ndarray of (N, 4, 2)
+        int_corners = [c.astype(np.int32).reshape((-1, 1, 2)) for c in corners]
+        cv2.polylines(image, int_corners, True, (0, 255, 0), 2)
+
+        if ids is not None:
+            # Flatten ids to handle (N,) or (N, 1)
+            flat_ids = ids.flatten()
+            for i, corner in enumerate(corners):
+                # Put ID near the first corner
+                # corner can be (1, 4, 2) or (4, 2)
+                c = corner.reshape((-1, 2))
+                pos = c[0].astype(int)
+                cv2.putText(
+                    image,
+                    str(flat_ids[i]),
+                    tuple(pos),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.5,
+                    (0, 0, 255),
+                    2,
+                )
+    return image
+
+
+def draw_pose_axes(
+    image: np.ndarray,
+    rvec: np.ndarray,
+    tvec: np.ndarray,
+    K: np.ndarray,
+    dist: np.ndarray | None = None,
+    length: float = 0.1,
+) -> np.ndarray:
+    """
+    Draws pose axes on the image.
+
+    Args:
+        image: Input image.
+        rvec: Rotation vector.
+        tvec: Translation vector.
+        K: Camera intrinsic matrix.
+        dist: Distortion coefficients.
+        length: Length of the axes.
+
+    Returns:
+        Image with axes drawn.
+    """
+    if image.ndim == 2:
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+
+    cv2.drawFrameAxes(image, K, dist, rvec, tvec, length)
+    return image
+
+
+def show_preview(
+    frames: dict[int, np.ndarray], *, wait_ms: int = 1, window_prefix: str = "View "
+) -> int:
+    """
+    Shows multiple camera frames in separate windows.
+
+    Args:
+        frames: Dictionary mapping camera serials to images.
+        wait_ms: Time to wait for a key press in milliseconds.
+        window_prefix: Prefix for the window names.
+
+    Returns:
+        The key code pressed, or -1 if no key was pressed.
+    """
+    for serial, frame in frames.items():
+        cv2.imshow(f"{window_prefix}{serial}", frame)
+    return cv2.waitKey(wait_ms)