409 lines
14 KiB
Python
409 lines
14 KiB
Python
|
"""HSV colour mask + swappable ball finders (contour / hough / blob / TM)."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
import threading
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .base import Detection, TargetDetector
|
||
|
|
from .teachable_machine import TeachableMachineModel
|
||
|
|
|
||
|
|
METHODS = ("contour", "hough", "blob", "rect", "teachable")
|
||
|
|
METHOD_LABELS = {
|
||
|
|
"contour": "Contour",
|
||
|
|
"hough": "Hough circles",
|
||
|
|
"blob": "Blob",
|
||
|
|
"rect": "Rectangle",
|
||
|
|
"teachable": "Teachable Machine",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class OrangeBallDetector(TargetDetector):
|
||
|
|
name = "orange_ball"
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
# Centre colour in OpenCV HSV (H: 0-179)
|
||
|
|
hsv_center: tuple[int, int, int] = (15, 180, 200),
|
||
|
|
h_tol: int = 10,
|
||
|
|
s_tol: int = 80,
|
||
|
|
v_tol: int = 80,
|
||
|
|
min_area: float = 80.0,
|
||
|
|
min_circularity: float = 0.55,
|
||
|
|
min_rectangularity: float = 0.65,
|
||
|
|
max_detections: int = 3,
|
||
|
|
method: str = "contour",
|
||
|
|
teachable_dir: str | Path | None = None,
|
||
|
|
) -> None:
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
self.hsv_center = np.array(hsv_center, dtype=np.int32)
|
||
|
|
self.h_tol = int(h_tol)
|
||
|
|
self.s_tol = int(s_tol)
|
||
|
|
self.v_tol = int(v_tol)
|
||
|
|
self.min_area = float(min_area)
|
||
|
|
self.min_circularity = float(min_circularity)
|
||
|
|
self.min_rectangularity = float(min_rectangularity)
|
||
|
|
self.max_detections = max_detections
|
||
|
|
self.method = self._normalize_method(method)
|
||
|
|
self._blob_detector = self._make_blob_detector(
|
||
|
|
self.min_area, self.min_circularity
|
||
|
|
)
|
||
|
|
self.teachable = TeachableMachineModel(teachable_dir)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _normalize_method(method: str) -> str:
|
||
|
|
key = method.strip().lower()
|
||
|
|
if key not in METHODS:
|
||
|
|
raise ValueError(f"unknown method '{method}'. known: {', '.join(METHODS)}")
|
||
|
|
return key
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _make_blob_detector(min_area: float, min_circularity: float):
|
||
|
|
params = cv2.SimpleBlobDetector_Params()
|
||
|
|
params.filterByColor = True
|
||
|
|
params.blobColor = 255
|
||
|
|
params.filterByArea = True
|
||
|
|
params.minArea = float(max(1.0, min_area))
|
||
|
|
params.maxArea = 1e6
|
||
|
|
params.filterByCircularity = True
|
||
|
|
params.minCircularity = float(max(0.01, min(1.0, min_circularity)))
|
||
|
|
params.filterByInertia = True
|
||
|
|
params.minInertiaRatio = 0.4
|
||
|
|
params.filterByConvexity = True
|
||
|
|
params.minConvexity = 0.6
|
||
|
|
params.minThreshold = 50
|
||
|
|
params.maxThreshold = 255
|
||
|
|
params.thresholdStep = 50
|
||
|
|
return cv2.SimpleBlobDetector_create(params)
|
||
|
|
|
||
|
|
def get_hsv_center(self) -> tuple[int, int, int]:
|
||
|
|
with self._lock:
|
||
|
|
h, s, v = (int(x) for x in self.hsv_center)
|
||
|
|
return h, s, v
|
||
|
|
|
||
|
|
def get_bgr_center(self) -> tuple[int, int, int]:
|
||
|
|
h, s, v = self.get_hsv_center()
|
||
|
|
pixel = np.uint8([[[h, s, v]]])
|
||
|
|
bgr = cv2.cvtColor(pixel, cv2.COLOR_HSV2BGR)[0, 0]
|
||
|
|
return int(bgr[0]), int(bgr[1]), int(bgr[2])
|
||
|
|
|
||
|
|
def get_tolerance(self) -> tuple[int, int, int]:
|
||
|
|
with self._lock:
|
||
|
|
return self.h_tol, self.s_tol, self.v_tol
|
||
|
|
|
||
|
|
def get_method(self) -> str:
|
||
|
|
with self._lock:
|
||
|
|
return self.method
|
||
|
|
|
||
|
|
def set_method(self, method: str) -> None:
|
||
|
|
key = self._normalize_method(method)
|
||
|
|
with self._lock:
|
||
|
|
self.method = key
|
||
|
|
if key == "teachable":
|
||
|
|
self.teachable.ensure_loaded()
|
||
|
|
|
||
|
|
def set_teachable_dir(self, model_dir: str | Path) -> None:
|
||
|
|
self.teachable.set_model_dir(model_dir)
|
||
|
|
if self.get_method() == "teachable":
|
||
|
|
self.teachable.ensure_loaded()
|
||
|
|
|
||
|
|
def teachable_status(self) -> str:
|
||
|
|
if self.teachable.ready:
|
||
|
|
labels = ", ".join(self.teachable.labels) or "?"
|
||
|
|
return f"ready ({labels})"
|
||
|
|
err = self.teachable.load_error
|
||
|
|
if err:
|
||
|
|
return f"error: {err}"
|
||
|
|
return f"not loaded - put keras_model.h5 + labels.txt in {self.teachable.model_dir}"
|
||
|
|
|
||
|
|
def set_from_bgr(self, b: int, g: int, r: int) -> None:
|
||
|
|
pixel = np.uint8([[[b, g, r]]])
|
||
|
|
hsv = cv2.cvtColor(pixel, cv2.COLOR_BGR2HSV)[0, 0]
|
||
|
|
self.set_hsv_center(int(hsv[0]), int(hsv[1]), int(hsv[2]))
|
||
|
|
|
||
|
|
def set_hsv_center(self, h: int, s: int, v: int) -> None:
|
||
|
|
with self._lock:
|
||
|
|
self.hsv_center = np.array(
|
||
|
|
[np.clip(h, 0, 179), np.clip(s, 0, 255), np.clip(v, 0, 255)],
|
||
|
|
dtype=np.int32,
|
||
|
|
)
|
||
|
|
|
||
|
|
def set_tolerance(
|
||
|
|
self,
|
||
|
|
h_tol: int | None = None,
|
||
|
|
s_tol: int | None = None,
|
||
|
|
v_tol: int | None = None,
|
||
|
|
) -> None:
|
||
|
|
with self._lock:
|
||
|
|
if h_tol is not None:
|
||
|
|
self.h_tol = int(np.clip(h_tol, 1, 90))
|
||
|
|
if s_tol is not None:
|
||
|
|
self.s_tol = int(np.clip(s_tol, 1, 255))
|
||
|
|
if v_tol is not None:
|
||
|
|
self.v_tol = int(np.clip(v_tol, 1, 255))
|
||
|
|
self._blob_detector = self._make_blob_detector(
|
||
|
|
self.min_area, self.min_circularity
|
||
|
|
)
|
||
|
|
|
||
|
|
def _snapshot(
|
||
|
|
self,
|
||
|
|
) -> tuple[np.ndarray, np.ndarray, float, float, float, int, str, int, int]:
|
||
|
|
with self._lock:
|
||
|
|
h, s, v = (int(x) for x in self.hsv_center)
|
||
|
|
lo = np.array(
|
||
|
|
[
|
||
|
|
max(0, h - self.h_tol),
|
||
|
|
max(0, s - self.s_tol),
|
||
|
|
max(0, v - self.v_tol),
|
||
|
|
],
|
||
|
|
dtype=np.uint8,
|
||
|
|
)
|
||
|
|
hi = np.array(
|
||
|
|
[
|
||
|
|
min(179, h + self.h_tol),
|
||
|
|
min(255, s + self.s_tol),
|
||
|
|
min(255, v + self.v_tol),
|
||
|
|
],
|
||
|
|
dtype=np.uint8,
|
||
|
|
)
|
||
|
|
return (
|
||
|
|
lo,
|
||
|
|
hi,
|
||
|
|
self.min_area,
|
||
|
|
self.min_circularity,
|
||
|
|
self.min_rectangularity,
|
||
|
|
self.max_detections,
|
||
|
|
self.method,
|
||
|
|
h,
|
||
|
|
self.h_tol,
|
||
|
|
)
|
||
|
|
|
||
|
|
def build_mask(self, frame_bgr: np.ndarray) -> np.ndarray:
|
||
|
|
hsv_low, hsv_high, *_rest, h_center, h_tol = self._snapshot()
|
||
|
|
hsv = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2HSV)
|
||
|
|
s_lo, s_hi = int(hsv_low[1]), int(hsv_high[1])
|
||
|
|
v_lo, v_hi = int(hsv_low[2]), int(hsv_high[2])
|
||
|
|
|
||
|
|
# Red hues wrap at 0/179 — merge two ranges when needed
|
||
|
|
if h_center - h_tol < 0:
|
||
|
|
mask_a = cv2.inRange(
|
||
|
|
hsv,
|
||
|
|
np.array([0, s_lo, v_lo], dtype=np.uint8),
|
||
|
|
np.array([h_center + h_tol, s_hi, v_hi], dtype=np.uint8),
|
||
|
|
)
|
||
|
|
wrap_lo = 179 + (h_center - h_tol)
|
||
|
|
mask_b = cv2.inRange(
|
||
|
|
hsv,
|
||
|
|
np.array([wrap_lo, s_lo, v_lo], dtype=np.uint8),
|
||
|
|
np.array([179, s_hi, v_hi], dtype=np.uint8),
|
||
|
|
)
|
||
|
|
mask = cv2.bitwise_or(mask_a, mask_b)
|
||
|
|
elif h_center + h_tol > 179:
|
||
|
|
mask_a = cv2.inRange(
|
||
|
|
hsv,
|
||
|
|
np.array([h_center - h_tol, s_lo, v_lo], dtype=np.uint8),
|
||
|
|
np.array([179, s_hi, v_hi], dtype=np.uint8),
|
||
|
|
)
|
||
|
|
wrap_hi = h_center + h_tol - 179
|
||
|
|
mask_b = cv2.inRange(
|
||
|
|
hsv,
|
||
|
|
np.array([0, s_lo, v_lo], dtype=np.uint8),
|
||
|
|
np.array([wrap_hi, s_hi, v_hi], dtype=np.uint8),
|
||
|
|
)
|
||
|
|
mask = cv2.bitwise_or(mask_a, mask_b)
|
||
|
|
else:
|
||
|
|
mask = cv2.inRange(hsv, hsv_low, hsv_high)
|
||
|
|
|
||
|
|
kernel = np.ones((5, 5), np.uint8)
|
||
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
|
||
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
|
||
|
|
return mask
|
||
|
|
|
||
|
|
def detect(self, frame_bgr: np.ndarray) -> list[Detection]:
|
||
|
|
(
|
||
|
|
_,
|
||
|
|
_,
|
||
|
|
min_area,
|
||
|
|
min_circularity,
|
||
|
|
min_rectangularity,
|
||
|
|
max_detections,
|
||
|
|
method,
|
||
|
|
*_rest,
|
||
|
|
) = self._snapshot()
|
||
|
|
if method == "teachable":
|
||
|
|
found = self.teachable.detect(frame_bgr, max_detections=max_detections)
|
||
|
|
else:
|
||
|
|
mask = self.build_mask(frame_bgr)
|
||
|
|
if method == "hough":
|
||
|
|
found = self._detect_hough(mask, min_area, max_detections)
|
||
|
|
elif method == "blob":
|
||
|
|
found = self._detect_blob(mask, max_detections)
|
||
|
|
elif method == "rect":
|
||
|
|
found = self._detect_rect(
|
||
|
|
mask, min_area, min_rectangularity, max_detections
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
found = self._detect_contour(
|
||
|
|
mask, min_area, min_circularity, max_detections
|
||
|
|
)
|
||
|
|
found.sort(key=lambda d: d.score, reverse=True)
|
||
|
|
return found[:max_detections]
|
||
|
|
|
||
|
|
def _detect_contour(
|
||
|
|
self,
|
||
|
|
mask: np.ndarray,
|
||
|
|
min_area: float,
|
||
|
|
min_circularity: float,
|
||
|
|
max_detections: int,
|
||
|
|
) -> list[Detection]:
|
||
|
|
contours, _ = cv2.findContours(
|
||
|
|
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
|
||
|
|
)
|
||
|
|
found: list[Detection] = []
|
||
|
|
for contour in contours:
|
||
|
|
area = float(cv2.contourArea(contour))
|
||
|
|
if area < min_area:
|
||
|
|
continue
|
||
|
|
peri = cv2.arcLength(contour, True)
|
||
|
|
if peri <= 0:
|
||
|
|
continue
|
||
|
|
circularity = float(4.0 * np.pi * area / (peri * peri))
|
||
|
|
if circularity < min_circularity:
|
||
|
|
continue
|
||
|
|
|
||
|
|
(cx, cy), radius = cv2.minEnclosingCircle(contour)
|
||
|
|
score = max(
|
||
|
|
0.0,
|
||
|
|
min(1.0, 0.5 * circularity + 0.5 * min(1.0, area / 4000.0)),
|
||
|
|
)
|
||
|
|
found.append(
|
||
|
|
Detection(
|
||
|
|
label="contour",
|
||
|
|
cx=float(cx),
|
||
|
|
cy=float(cy),
|
||
|
|
radius=float(radius),
|
||
|
|
score=score,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return found[:max_detections]
|
||
|
|
|
||
|
|
def _detect_rect(
|
||
|
|
self,
|
||
|
|
mask: np.ndarray,
|
||
|
|
min_area: float,
|
||
|
|
min_rectangularity: float,
|
||
|
|
max_detections: int,
|
||
|
|
) -> list[Detection]:
|
||
|
|
contours, _ = cv2.findContours(
|
||
|
|
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
|
||
|
|
)
|
||
|
|
found: list[Detection] = []
|
||
|
|
for contour in contours:
|
||
|
|
area = float(cv2.contourArea(contour))
|
||
|
|
if area < min_area:
|
||
|
|
continue
|
||
|
|
peri = cv2.arcLength(contour, True)
|
||
|
|
if peri <= 0:
|
||
|
|
continue
|
||
|
|
approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
|
||
|
|
if len(approx) != 4:
|
||
|
|
continue
|
||
|
|
if not cv2.isContourConvex(approx):
|
||
|
|
continue
|
||
|
|
|
||
|
|
rect = cv2.minAreaRect(contour)
|
||
|
|
(cx, cy), (w, h), angle = rect
|
||
|
|
box_area = float(w) * float(h)
|
||
|
|
if box_area <= 0:
|
||
|
|
continue
|
||
|
|
rectangularity = area / box_area
|
||
|
|
if rectangularity < min_rectangularity:
|
||
|
|
continue
|
||
|
|
|
||
|
|
score = max(
|
||
|
|
0.0,
|
||
|
|
min(
|
||
|
|
1.0,
|
||
|
|
0.6 * rectangularity + 0.4 * min(1.0, area / 4000.0),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
radius = max(float(w), float(h)) * 0.5
|
||
|
|
found.append(
|
||
|
|
Detection(
|
||
|
|
label="rect",
|
||
|
|
cx=float(cx),
|
||
|
|
cy=float(cy),
|
||
|
|
radius=radius,
|
||
|
|
score=score,
|
||
|
|
width=float(w),
|
||
|
|
height=float(h),
|
||
|
|
angle=float(angle),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return found[:max_detections]
|
||
|
|
|
||
|
|
def _detect_hough(
|
||
|
|
self,
|
||
|
|
mask: np.ndarray,
|
||
|
|
min_area: float,
|
||
|
|
max_detections: int,
|
||
|
|
) -> list[Detection]:
|
||
|
|
# Estimate radius bounds from min_area of a filled circle
|
||
|
|
min_radius = max(3, int(round(np.sqrt(min_area / np.pi))))
|
||
|
|
blurred = cv2.GaussianBlur(mask, (9, 9), 2)
|
||
|
|
circles = cv2.HoughCircles(
|
||
|
|
blurred,
|
||
|
|
cv2.HOUGH_GRADIENT,
|
||
|
|
dp=1.2,
|
||
|
|
minDist=max(12, min_radius * 2),
|
||
|
|
param1=80,
|
||
|
|
param2=18,
|
||
|
|
minRadius=min_radius,
|
||
|
|
maxRadius=0,
|
||
|
|
)
|
||
|
|
found: list[Detection] = []
|
||
|
|
if circles is None:
|
||
|
|
return found
|
||
|
|
for x, y, r in circles[0]:
|
||
|
|
area = float(np.pi * r * r)
|
||
|
|
score = max(0.0, min(1.0, min(1.0, area / 4000.0)))
|
||
|
|
found.append(
|
||
|
|
Detection(
|
||
|
|
label="hough",
|
||
|
|
cx=float(x),
|
||
|
|
cy=float(y),
|
||
|
|
radius=float(r),
|
||
|
|
score=score,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return found[:max_detections]
|
||
|
|
|
||
|
|
def _detect_blob(self, mask: np.ndarray, max_detections: int) -> list[Detection]:
|
||
|
|
with self._lock:
|
||
|
|
detector = self._blob_detector
|
||
|
|
keypoints = detector.detect(mask)
|
||
|
|
found: list[Detection] = []
|
||
|
|
for kp in keypoints:
|
||
|
|
radius = float(kp.size) * 0.5
|
||
|
|
area = float(np.pi * radius * radius)
|
||
|
|
# OpenCV blob response is not always normalized; fold size in
|
||
|
|
response = float(kp.response) if kp.response else 0.0
|
||
|
|
score = max(
|
||
|
|
0.0,
|
||
|
|
min(1.0, 0.4 * min(1.0, response) + 0.6 * min(1.0, area / 4000.0)),
|
||
|
|
)
|
||
|
|
found.append(
|
||
|
|
Detection(
|
||
|
|
label="blob",
|
||
|
|
cx=float(kp.pt[0]),
|
||
|
|
cy=float(kp.pt[1]),
|
||
|
|
radius=radius,
|
||
|
|
score=score,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return found[:max_detections]
|