diff options
Diffstat (limited to 'plugins/motion_capture/tracker.py')
| -rw-r--r-- | plugins/motion_capture/tracker.py | 306 |
1 files changed, 201 insertions, 105 deletions
diff --git a/plugins/motion_capture/tracker.py b/plugins/motion_capture/tracker.py index d65aca1..ffc9f25 100644 --- a/plugins/motion_capture/tracker.py +++ b/plugins/motion_capture/tracker.py @@ -1,15 +1,15 @@ """ motion_capture/tracker.py -Camera capture + OpenCV CSRT point tracking. +Camera capture + two tracking modes: + "template" — snapshot + shape-masked template matching (default, good for markers) + "csrt" — OpenCV CSRT object tracker (texture-based, good for complex targets) -Falls back to simulation mode (Lissajous figure) when OpenCV is not -installed or the requested camera index cannot be opened. +Falls back to simulation (Lissajous) when OpenCV is absent or camera fails. Thread model: - _thread runs continuously, updating _frame (RGB numpy array) and - _position (x, y floats). All public methods are thread-safe. - get_frame() / get_position() are safe to call from the Qt main thread. + _thread runs continuously, updating _frame and _position. + All public methods are thread-safe via _lock. """ from __future__ import annotations @@ -24,30 +24,25 @@ import numpy as np # ── Simulation constants ────────────────────────────────────────────────────── _SIM_W, _SIM_H = 640, 480 -_BG_COLOR = (17, 22, 40) # match app dark theme +_BG_COLOR = (17, 22, 40) _GRID_COLOR = (42, 53, 88) -_DOT_COLOR = (0, 212, 255) # cyan accent +_DOT_COLOR = (0, 212, 255) def _make_sim_frame(cx: int, cy: int) -> np.ndarray: frame = np.full((_SIM_H, _SIM_W, 3), _BG_COLOR, dtype=np.uint8) - # Grid frame[::60, :] = _GRID_COLOR frame[:, ::80] = _GRID_COLOR - # Crosshair cx = max(20, min(_SIM_W - 21, cx)) cy = max(20, min(_SIM_H - 21, cy)) frame[cy - 12:cy + 13, cx - 1:cx + 2] = _DOT_COLOR frame[cy - 1:cy + 2, cx - 12:cx + 13] = _DOT_COLOR - # Dot frame[cy - 5:cy + 6, cx - 5:cx + 6] = _DOT_COLOR - # Label frame[10:18, 10:260] = (30, 38, 65) return frame def _draw_tracking(frame: np.ndarray, cx: int, cy: int) -> np.ndarray: - """Draw crosshair + circle on an RGB frame (modifies in-place).""" h, w = frame.shape[:2] cx = max(20, min(w - 21, cx)) cy = max(20, min(h - 21, cy)) @@ -61,41 +56,79 @@ def _draw_tracking(frame: np.ndarray, cx: int, cy: int) -> np.ndarray: class CameraTracker: """ - Manages camera capture and CSRT tracking in a background thread. + Manages camera capture and tracking in a background thread. Public API (all thread-safe): - start(camera_index) → bool open camera / start sim - stop() close camera / stop thread - set_roi(cx, cy, size) begin tracking at pixel (cx, cy) - clear_roi() stop tracking (keep camera running) - get_frame() → ndarray | None latest RGB frame with overlay - get_position() → (x, y) | None tracked point, or None - is_tracking() → bool - is_lost() → bool - fps → float - simulated → bool + start(camera_index, resolution) open camera / start sim + stop() close camera / stop thread + set_track_mode(mode) "template" or "csrt" + set_template_shape(shape) "rect" or "circle" + set_roi(cx, cy, size) point-click target init (legacy) + set_roi_rect(x, y, w, h, frame) exact-rect target init + clear_roi() stop tracking + get_frame() → ndarray | None + get_position() → (x, y) | None + is_tracking() → bool + is_lost() → bool + fps → float + simulated → bool + track_mode → str + template_shape → str """ def __init__(self): - self._lock = threading.Lock() - self._running = False + self._lock = threading.Lock() + self._running = False self._thread: Optional[threading.Thread] = None - self._frame: Optional[np.ndarray] = None - self._position: Tuple[float, float] = (0.0, 0.0) + self._frame: Optional[np.ndarray] = None + self._position: Tuple[float, float] = (0.0, 0.0) self._tracking = False self._lost = False self._fps = 0.0 self._simulated = True - # OpenCV objects — only set when cv2 available - self._cap = None - self._ot = None # OpenCV tracker - self._roi: Optional[Tuple] = None # pending ROI to init + self._track_mode = "template" + self._template_shape = "rect" + + self._cap = None + self._ot = None + + self._template_bgr: Optional[np.ndarray] = None + self._template_mask: Optional[np.ndarray] = None + + # ── Configuration ───────────────────────────────────────────────────── + + def set_track_mode(self, mode: str): + with self._lock: + self._track_mode = mode + self._ot = None + self._template_bgr = None + self._template_mask = None + self._tracking = False + self._lost = False + + def set_template_shape(self, shape: str): + with self._lock: + self._template_shape = shape + self._template_bgr = None + self._template_mask = None + self._tracking = False + + @property + def track_mode(self) -> str: + with self._lock: + return self._track_mode + + @property + def template_shape(self) -> str: + with self._lock: + return self._template_shape # ── Start / stop ────────────────────────────────────────────────────── - def start(self, camera_index: int = 0) -> bool: + def start(self, camera_index: int = 0, + resolution: Optional[Tuple[int, int]] = None) -> bool: self.stop() self._running = True @@ -109,26 +142,30 @@ class CameraTracker: try: import cv2 cap = cv2.VideoCapture(camera_index) + if resolution: + cap.set(cv2.CAP_PROP_FRAME_WIDTH, resolution[0]) + cap.set(cv2.CAP_PROP_FRAME_HEIGHT, resolution[1]) if cap.isOpened(): with self._lock: - self._cap = cap - self._simulated = False - self._tracking = False - self._lost = False - self._ot = None - self._roi = None + self._cap = cap + self._simulated = False + self._tracking = False + self._lost = False + self._ot = None + self._template_bgr = None + self._template_mask = None self._thread = threading.Thread( target=self._run_real, daemon=True) self._thread.start() return True except ImportError: - pass # cv2 not installed — fall through to simulation + pass with self._lock: self._simulated = True self._thread = threading.Thread(target=self._run_sim, daemon=True) self._thread.start() - return True # simulation always succeeds + return True def stop(self): self._running = False @@ -139,18 +176,32 @@ class CameraTracker: if self._cap is not None: self._cap.release() self._cap = None - self._ot = None - self._tracking = False - self._lost = False - self._frame = None + self._ot = None + self._template_bgr = None + self._template_mask = None + self._tracking = False + self._lost = False + self._frame = None # ── ROI control ─────────────────────────────────────────────────────── def set_roi(self, cx: int, cy: int, size: int = 70): - """Begin tracking at pixel coordinate (cx, cy) with box of `size`.""" + """Point-click init — delegates to set_roi_rect with a centred square.""" + half = size // 2 + self.set_roi_rect(cx - half, cy - half, size, size) + + def set_roi_rect(self, x: int, y: int, w: int, h: int, + frame_rgb: Optional[np.ndarray] = None): + """ + Initialise tracking from an exact rectangle. + frame_rgb: frozen RGB frame to extract template from; falls back to + the latest live frame if None. + """ with self._lock: - frame = self._frame simulated = self._simulated + mode = self._track_mode + + cx, cy = x + w / 2, y + h / 2 if simulated: with self._lock: @@ -159,37 +210,48 @@ class CameraTracker: self._lost = False return - if frame is None: + src = frame_rgb + if src is None: + with self._lock: + src = self._frame + if src is None: return try: import cv2 - bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + bgr = cv2.cvtColor(src, cv2.COLOR_RGB2BGR) except Exception: return - h, w = bgr.shape[:2] - half = size // 2 - x = max(0, cx - half) - y = max(0, cy - half) - bw = min(w - x, size) - bh = min(h - y, size) - - ot = _make_cv_tracker() - if ot is None: + fh, fw = bgr.shape[:2] + x1 = max(0, x); y1 = max(0, y) + x2 = min(fw, x + w); y2 = min(fh, y + h) + rw, rh = x2 - x1, y2 - y1 + if rw < 4 or rh < 4: return - ot.init(bgr, (x, y, bw, bh)) - with self._lock: - self._ot = ot - self._tracking = True - self._lost = False + if mode == "template": + self._build_template(bgr, x1, y1, rw, rh) + with self._lock: + self._tracking = True + self._lost = False + else: + ot = _make_cv_tracker() + if ot is None: + return + ot.init(bgr, (x1, y1, rw, rh)) + with self._lock: + self._ot = ot + self._tracking = True + self._lost = False def clear_roi(self): with self._lock: - self._ot = None - self._tracking = False - self._lost = False + self._ot = None + self._template_bgr = None + self._template_mask = None + self._tracking = False + self._lost = False # ── Data access ─────────────────────────────────────────────────────── @@ -218,6 +280,47 @@ class CameraTracker: with self._lock: return self._simulated + # ── Template helpers ────────────────────────────────────────────────── + + def _build_template(self, bgr: np.ndarray, x1: int, y1: int, rw: int, rh: int): + """Extract ROI from bgr and build the shape-masked template (under lock).""" + import cv2 + template = bgr[y1:y1 + rh, x1:x1 + rw].copy() + th, tw = template.shape[:2] + + with self._lock: + shape = self._template_shape + + if shape == "circle": + mask = np.zeros((th, tw), dtype=np.uint8) + cv2.circle(mask, (tw // 2, th // 2), min(tw, th) // 2, 255, -1) + else: + mask = None + + with self._lock: + self._template_bgr = template + self._template_mask = mask + + def _match_template(self, bgr: np.ndarray) -> Optional[Tuple[float, float]]: + import cv2 + with self._lock: + tmpl = self._template_bgr + mask = self._template_mask + + if tmpl is None: + return None + th, tw = tmpl.shape[:2] + if bgr.shape[0] < th or bgr.shape[1] < tw: + return None + + try: + method = cv2.TM_CCORR_NORMED if mask is not None else cv2.TM_CCOEFF_NORMED + result = cv2.matchTemplate(bgr, tmpl, method, mask=mask) + _, _, _, max_loc = cv2.minMaxLoc(result) + return float(max_loc[0] + tw / 2), float(max_loc[1] + th / 2) + except Exception: + return None + # ── Background loops ────────────────────────────────────────────────── def _run_real(self): @@ -237,26 +340,36 @@ class CameraTracker: rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) with self._lock: - ot = self._ot - - if ot is not None: - ok, bbox = ot.update(bgr) - if ok: - cx = bbox[0] + bbox[2] / 2 - cy = bbox[1] + bbox[3] / 2 - with self._lock: - self._position = (float(cx), float(cy)) - self._lost = False - _draw_tracking(rgb, int(cx), int(cy)) - else: - with self._lock: - self._lost = True - self._tracking = False - self._ot = None - - # FPS + mode = self._track_mode + ot = self._ot + has_tmpl = self._template_bgr is not None + tracking = self._tracking + + if tracking: + if mode == "template" and has_tmpl: + pos = self._match_template(bgr) + if pos is not None: + with self._lock: + self._position = pos + self._lost = False + _draw_tracking(rgb, int(pos[0]), int(pos[1])) + elif mode == "csrt" and ot is not None: + ok, bbox = ot.update(bgr) + if ok: + cx = bbox[0] + bbox[2] / 2 + cy = bbox[1] + bbox[3] / 2 + with self._lock: + self._position = (float(cx), float(cy)) + self._lost = False + _draw_tracking(rgb, int(cx), int(cy)) + else: + with self._lock: + self._lost = True + self._tracking = False + self._ot = None + now = time.monotonic() - dt = now - t_prev + dt = now - t_prev if dt > 0: self._fps = 0.9 * self._fps + 0.1 * (1.0 / dt) t_prev = now @@ -265,29 +378,21 @@ class CameraTracker: self._frame = rgb def _run_sim(self): - """Lissajous simulation — no camera hardware needed.""" t0 = time.monotonic() t_prev = t0 - cx, cy = float(_SIM_W // 2), float(_SIM_H // 2) - while self._running: t = time.monotonic() - t0 cx = _SIM_W / 2 + (_SIM_W / 2 - 60) * 0.9 * math.sin(t * 0.31) cy = _SIM_H / 2 + (_SIM_H / 2 - 50) * 0.9 * math.sin(t * 0.47 + 0.5) - with self._lock: tracking = self._tracking - frame = _make_sim_frame(int(cx), int(cy)) - with self._lock: if tracking: - # Sim tracking: dot follows Lissajous self._position = (cx, cy) self._frame = frame - - now = time.monotonic() - dt = now - t_prev + now = time.monotonic() + dt = now - t_prev if dt > 0: self._fps = 0.9 * self._fps + 0.1 * (1.0 / dt) t_prev = now @@ -297,16 +402,8 @@ class CameraTracker: # ── Camera enumeration ──────────────────────────────────────────────────────── def list_cameras(): - """ - Return list of (index, label) for all detected camera devices. - - On Linux: reads /dev/video* and pulls human-readable names from sysfs. - Falls back to probing cv2.VideoCapture indices 0-7 on other platforms. - Always appends a Simulation entry. - """ cameras = [] - - import os, glob + import glob video_nodes = sorted(glob.glob("/dev/video*")) if video_nodes: @@ -323,7 +420,6 @@ def list_cameras(): name = path cameras.append((idx, f"{name} [/dev/video{idx}]")) else: - # Non-Linux fallback: probe indices try: import cv2 for idx in range(8): |