path: root/plugins/motion_capture/tracker.py

"""
motion_capture/tracker.py

Camera capture + OpenCV CSRT point tracking.

Falls back to simulation mode (Lissajous figure) when OpenCV is not
installed or the requested camera index cannot be opened.

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.
"""

from __future__ import annotations

import math
import threading
import time
from typing import Optional, Tuple

import numpy as np


# ── Simulation constants ──────────────────────────────────────────────────────
_SIM_W, _SIM_H = 640, 480
_BG_COLOR  = (17, 22, 40)    # match app dark theme
_GRID_COLOR = (42, 53, 88)
_DOT_COLOR  = (0, 212, 255)  # cyan accent


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))
    frame[cy - 16:cy + 17, cx - 1:cx + 2] = (0, 255, 0)
    frame[cy - 1:cy + 2,   cx - 16:cx + 17] = (0, 255, 0)
    frame[cy - 6:cy + 7, cx - 6:cx + 7] = (0, 255, 0)
    return frame


# ── Tracker ───────────────────────────────────────────────────────────────────

class CameraTracker:
    """
    Manages camera capture and CSRT 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
    """

    def __init__(self):
        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._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

    # ── Start / stop ──────────────────────────────────────────────────────

    def start(self, camera_index: int = 0) -> bool:
        self.stop()
        self._running = True

        try:
            import cv2
            cap = cv2.VideoCapture(camera_index)
            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._thread = threading.Thread(
                    target=self._run_real, daemon=True)
                self._thread.start()
                return True
        except ImportError:
            pass  # cv2 not installed — fall through to simulation

        with self._lock:
            self._simulated = True
        self._thread = threading.Thread(target=self._run_sim, daemon=True)
        self._thread.start()
        return True   # simulation always succeeds

    def stop(self):
        self._running = False
        if self._thread:
            self._thread.join(timeout=2)
            self._thread = None
        with self._lock:
            if self._cap is not None:
                self._cap.release()
                self._cap = None
            self._ot       = 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`."""
        with self._lock:
            frame = self._frame
            simulated = self._simulated

        if simulated:
            with self._lock:
                self._position = (float(cx), float(cy))
                self._tracking = True
                self._lost     = False
            return

        if frame is None:
            return

        try:
            import cv2
            bgr = cv2.cvtColor(frame, 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:
            return
        ot.init(bgr, (x, y, bw, bh))

        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

    # ── Data access ───────────────────────────────────────────────────────

    def get_frame(self) -> Optional[np.ndarray]:
        with self._lock:
            return self._frame.copy() if self._frame is not None else None

    def get_position(self) -> Optional[Tuple[float, float]]:
        with self._lock:
            return self._position if self._tracking else None

    def is_tracking(self) -> bool:
        with self._lock:
            return self._tracking

    def is_lost(self) -> bool:
        with self._lock:
            return self._lost

    @property
    def fps(self) -> float:
        return self._fps

    @property
    def simulated(self) -> bool:
        with self._lock:
            return self._simulated

    # ── Background loops ──────────────────────────────────────────────────

    def _run_real(self):
        import cv2
        t_prev = time.monotonic()
        while self._running:
            with self._lock:
                cap = self._cap
            if cap is None:
                break

            ret, bgr = cap.read()
            if not ret:
                time.sleep(0.05)
                continue

            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
            now = time.monotonic()
            dt = now - t_prev
            if dt > 0:
                self._fps = 0.9 * self._fps + 0.1 * (1.0 / dt)
            t_prev = now

            with self._lock:
                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
            if dt > 0:
                self._fps = 0.9 * self._fps + 0.1 * (1.0 / dt)
            t_prev = now
            time.sleep(1 / 30)


# ── Helper ────────────────────────────────────────────────────────────────────

def _make_cv_tracker():
    try:
        import cv2
        for factory in (
            lambda: cv2.TrackerCSRT_create(),
            lambda: cv2.legacy.TrackerCSRT_create(),
            lambda: cv2.TrackerKCF_create(),
            lambda: cv2.legacy.TrackerKCF_create(),
        ):
            try:
                return factory()
            except AttributeError:
                continue
    except ImportError:
        pass
    return None