File size: 1,813 Bytes
604e535 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """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)
|