| """Pose extraction for traditional controllers from clean images.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
|
|
| def estimate_pose_from_clean_image( |
| image: np.ndarray, |
| workspace: tuple[float, float, float, float] = (0.0, 10.0, 0.0, 10.0), |
| pad: int = 4, |
| visual_scale: float = 2.5, |
| ) -> np.ndarray: |
| img = image.astype(np.float32) |
| mask = (img[..., 2] > 110.0) & (img[..., 0] < 100.0) & (img[..., 1] < 140.0) |
| ys, xs = np.nonzero(mask) |
| centroid_x = xs.mean() |
| centroid_y = ys.mean() |
| coords = np.stack([xs - centroid_x, ys - centroid_y], axis=1) |
| cov = coords.T @ coords / coords.shape[0] |
| vals, vecs = np.linalg.eigh(cov) |
| axis = vecs[:, int(np.argmax(vals))] |
| theta_px = np.arctan2(-axis[1], axis[0]) |
| marker = (img[..., 0] > 180.0) & (img[..., 1] > 140.0) & (img[..., 2] < 130.0) |
| my, mx = np.nonzero(marker) |
| marker_x = marker_y = None |
| if len(mx): |
| marker_x = float(mx.mean()) |
| marker_y = float(my.mean()) |
| theta_px = np.arctan2(-(marker_y - centroid_y), marker_x - centroid_x) |
| image_size = image.shape[0] |
| xmin, xmax, ymin, ymax = workspace |
| x = xmin + (centroid_x - pad) / (image_size - 2 * pad) * (xmax - xmin) |
| y = ymin + (image_size - pad - centroid_y) / (image_size - 2 * pad) * (ymax - ymin) |
| if marker_x is not None and marker_y is not None: |
| marker_world_x = xmin + (marker_x - pad) / (image_size - 2 * pad) * (xmax - xmin) |
| marker_world_y = ymin + (image_size - pad - marker_y) / (image_size - 2 * pad) * (ymax - ymin) |
| marker_offset = 0.22 * float(visual_scale) |
| x = marker_world_x - marker_offset * np.cos(theta_px) |
| y = marker_world_y - marker_offset * np.sin(theta_px) |
| return np.array([x, y, np.cos(theta_px), np.sin(theta_px)], dtype=np.float32) |
|
|
