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 Optional, Tuple

import numpy as np


# ── Simulation constants ──────────────────────────────────────────────────────
_SIM_W, _SIM_H = 640, 480
_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) -> bool:
        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()
            return True

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

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

    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

    @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):
        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:
                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

    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