63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
|
"""Shared types for swappable vision targets."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class Detection:
|
||
|
|
"""One detected object in image pixel coordinates."""
|
||
|
|
|
||
|
|
label: str
|
||
|
|
cx: float
|
||
|
|
cy: float
|
||
|
|
radius: float
|
||
|
|
score: float # 0..1 confidence / quality
|
||
|
|
width: float | None = None # minAreaRect width (rect detections)
|
||
|
|
height: float | None = None
|
||
|
|
angle: float | None = None # minAreaRect angle in degrees
|
||
|
|
|
||
|
|
|
||
|
|
class TargetDetector:
|
||
|
|
"""Interface every target module should implement."""
|
||
|
|
|
||
|
|
name: str = "target"
|
||
|
|
|
||
|
|
def detect(self, frame_bgr: np.ndarray) -> list[Detection]:
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def annotate(self, frame_bgr: np.ndarray, detections: list[Detection]) -> np.ndarray:
|
||
|
|
"""Draw detections on a copy of the frame and return it."""
|
||
|
|
import cv2
|
||
|
|
|
||
|
|
out = frame_bgr.copy()
|
||
|
|
for d in detections:
|
||
|
|
center = (int(round(d.cx)), int(round(d.cy)))
|
||
|
|
if d.width is not None and d.height is not None and d.angle is not None:
|
||
|
|
rect = (
|
||
|
|
(float(d.cx), float(d.cy)),
|
||
|
|
(float(d.width), float(d.height)),
|
||
|
|
float(d.angle),
|
||
|
|
)
|
||
|
|
box = cv2.boxPoints(rect).astype(int)
|
||
|
|
cv2.drawContours(out, [box], 0, (0, 255, 0), 2)
|
||
|
|
cv2.drawMarker(out, center, (0, 255, 255), cv2.MARKER_CROSS, 12, 2)
|
||
|
|
else:
|
||
|
|
radius = max(1, int(round(d.radius)))
|
||
|
|
cv2.circle(out, center, radius, (0, 255, 0), 2)
|
||
|
|
cv2.drawMarker(out, center, (0, 255, 255), cv2.MARKER_CROSS, 12, 2)
|
||
|
|
cv2.putText(
|
||
|
|
out,
|
||
|
|
f"{d.label} {d.score:.2f}",
|
||
|
|
(center[0] + 8, center[1] - 8),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
||
|
|
0.45,
|
||
|
|
(0, 255, 0),
|
||
|
|
1,
|
||
|
|
cv2.LINE_AA,
|
||
|
|
)
|
||
|
|
return out
|