path: root/plugins/motion_capture/tracker.py

"""
motion_capture/tracker.py

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 (Lissajous) when OpenCV is absent or camera fails.

Thread model:
    _thread runs continuously, updating _frame and _position.
    All public methods are thread-safe via _lock.
"""

from __future__ import annotations

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

import numpy as np


# ── Simulation constants ──────────────────────────────────────────────────────
_SIM_W, _SIM_H = 640, 480

# Caps how often _run_real processes a captured frame, independent of the
# camera's configured capture fps (camera_panel.py allows up to 120 fps at
# 2560x1440 with no validation) — otherwise cvtColor + matchTemplate/CSRT
# update on every frame at high settings saturates a CPU core and starves
# the GUI thread's ability to service the Qt event loop (sluggish/unresponsive
# controls app-wide, not just in the camera window).
_MAX_TRACK_HZ = 20

# Consecutive failed cap.read() calls before giving up on a real camera and
# falling back to simulation — ~1.5s at the 50ms failed-read retry spacing
# below. Some backends (Windows MSMF/obsensor observed) report a camera as
# opened even when the stream never actually delivers frames.
_MAX_READ_FAILURES = 30
_BG_COLOR   = (17, 22, 40)
_GRID_COLOR = (42, 53, 88)
_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)
    frame[::60, :] = _GRID_COLOR
    frame[:, ::80] = _GRID_COLOR
    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
    frame[cy - 5:cy + 6, cx - 5:cx + 6] = _DOT_COLOR
    frame[10:18, 10:260] = (30, 38, 65)
    return frame


def _draw_tracking(frame: np.ndarray, cx: int, cy: int) -> np.ndarray:
    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 tracking in a background thread.

    Public API (all thread-safe):
        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._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

        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,
              resolution: Optional[Tuple[int, int]] = None,
              fps: Optional[int] = None,
              on_result: Optional[Callable[[bool], None]] = None) -> bool:
        """Start capture. Returns True immediately — always non-blocking.

        cv2.VideoCapture(camera_index) can block for seconds (driver/backend
        enumeration, especially on Windows), so opening the camera happens
        on the background thread, not here — calling this synchronously
        during e.g. profile load used to freeze the whole app. `on_result`,
        if given, is called from that thread once the outcome is known:
        on_result(True) for simulated/fallback, on_result(False) for a real
        camera that opened successfully."""
        self.stop()
        self._running = True

        if camera_index == -1:
            with self._lock:
                self._simulated = True
            self._thread = threading.Thread(target=self._run_sim, daemon=True)
            self._thread.start()
            if on_result:
                on_result(True)
            return True

        self._thread = threading.Thread(
            target=self._open_and_run,
            args=(camera_index, resolution, fps, on_result),
            daemon=True,
        )
        self._thread.start()
        return True

    def _open_and_run(self, camera_index: int,
                      resolution: Optional[Tuple[int, int]],
                      fps: Optional[int],
                      on_result: Optional[Callable[[bool], None]]):
        """Runs on the background thread: opens the camera (the slow part),
        then falls into the real or simulated read loop in place."""
        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 fps:
                cap.set(cv2.CAP_PROP_FPS, fps)

            # isOpened() alone isn't enough — some backends (seen with
            # Windows MSMF/obsensor) report a device as open even when it
            # can't actually deliver frames, e.g. a stale camera_index from
            # a profile saved on a different machine. Confirm with a real
            # read before committing, or _run_real would otherwise loop
            # forever hammering a broken handle.
            opened = False
            if cap.isOpened():
                for _ in range(5):
                    ret, _ = cap.read()
                    if ret:
                        opened = True
                        break
                    time.sleep(0.05)

            if opened:
                with self._lock:
                    self._cap           = cap
                    self._simulated     = False
                    self._tracking      = False
                    self._lost          = False
                    self._ot            = None
                    self._template_bgr  = None
                    self._template_mask = None
                if on_result:
                    on_result(False)
                self._run_real(on_result)
                return
            cap.release()
        except Exception:
            pass

        with self._lock:
            self._simulated = True
        if on_result:
            on_result(True)
        self._run_sim()

    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._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):
        """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:
            simulated = self._simulated
            mode      = self._track_mode

        cx, cy = x + w / 2, y + h / 2

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

        src = frame_rgb
        if src is None:
            with self._lock:
                src = self._frame
        if src is None:
            return

        try:
            import cv2
            bgr = cv2.cvtColor(src, cv2.COLOR_RGB2BGR)
        except Exception:
            return

        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

        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._template_bgr  = None
            self._template_mask = 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

    def is_running(self) -> bool:
        return self._running and self._thread is not None and self._thread.is_alive()

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

    @property
    def simulated(self) -> bool:
        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, on_result: Optional[Callable[[bool], None]] = None):
        import cv2
        t_prev = time.monotonic()
        fail_count = 0
        while self._running:
            t_iter_start = time.monotonic()

            with self._lock:
                cap = self._cap
            if cap is None:
                break

            ret, bgr = cap.read()
            if not ret:
                fail_count += 1
                # The initial test-read in _open_and_run already screens out
                # cameras that never work, but a device can also die mid-
                # session (unplugged, driver fault) — without this, a broken
                # handle gets hammered forever, spamming OpenCV errors and
                # burning CPU/GIL time the whole app needs.
                if fail_count >= _MAX_READ_FAILURES:
                    with self._lock:
                        if self._cap is not None:
                            self._cap.release()
                            self._cap = None
                        self._simulated = True
                    if on_result:
                        on_result(True)
                    self._run_sim()
                    return
                time.sleep(0.05)
                continue
            fail_count = 0

            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)

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

            with self._lock:
                self._frame = rgb

            elapsed = time.monotonic() - t_iter_start
            remaining = (1.0 / _MAX_TRACK_HZ) - elapsed
            if remaining > 0:
                time.sleep(remaining)

    def _run_sim(self):
        t0     = time.monotonic()
        t_prev = t0
        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:
                    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)


# ── Camera enumeration ────────────────────────────────────────────────────────

def list_cameras():
    cameras = []
    import glob

    video_nodes = sorted(glob.glob("/dev/video*"))
    if video_nodes:
        for path in video_nodes:
            try:
                idx = int(path.replace("/dev/video", ""))
            except ValueError:
                continue
            name_path = f"/sys/class/video4linux/video{idx}/name"
            try:
                with open(name_path) as f:
                    name = f.read().strip()
            except OSError:
                name = path
            cameras.append((idx, f"{name}  [/dev/video{idx}]"))
    else:
        try:
            import cv2
            for idx in range(8):
                cap = cv2.VideoCapture(idx)
                if cap.isOpened():
                    cameras.append((idx, f"Camera {idx}"))
                    cap.release()
        except ImportError:
            pass

    cameras.append((-1, "Simulation (no camera)"))
    return cameras


# ── 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