Merge fix/motion-capture: template matching + frame-freeze selection UI

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

1 files changed, 344 insertions, 100 deletions

diff --git a/plugins/motion_capture/window.py b/plugins/motion_capture/window.py
index 40b436f..bc9a0de 100644
--- a/plugins/motion_capture/window.py
+++ b/plugins/motion_capture/window.py
@@ -5,22 +5,240 @@ Floating tool window for the Motion Capture plugin.
 
 Layout:
   [Title bar]
-  [Video feed — live camera or simulation, click to select tracking point]
-  [Status bar — position, fps, tracking state]
-  [Controls — camera index, connect, click-to-track, clear, calibration]
+  [FrameSelector — live / frozen feed with zoom + rubber-band target selection]
+  [Status bar]
+  [Controls — camera, connect, mode, shape, track/confirm, clear, px/mm]
 """
 
 from __future__ import annotations
 
-from PyQt6.QtCore import Qt, QTimer, pyqtSignal
-from PyQt6.QtGui import QCloseEvent, QImage, QMouseEvent, QPixmap
+from typing import Optional, Tuple
+
+import numpy as np
+
+from PyQt6.QtCore import Qt, QPointF, QTimer, pyqtSignal
+from PyQt6.QtGui import (
+    QCloseEvent, QColor, QFont, QImage, QMouseEvent,
+    QPainter, QPen, QPixmap,
+)
 from PyQt6.QtWidgets import (
     QComboBox, QDoubleSpinBox, QFrame, QHBoxLayout, QLabel, QPushButton,
     QSizePolicy, QVBoxLayout, QWidget,
 )
 
-import numpy as np
 
+# ── Frame display + rubber-band selection widget ──────────────────────────────
+
+class FrameSelector(QWidget):
+    """
+    Live or frozen camera frame display with zoom, pan, and rubber-band
+    target selection.
+
+    Live mode  — set_frame() updates the display in real time.
+    Frozen mode — freeze() locks the frame; user scrolls to zoom (around
+                  cursor), right-drags to pan, left-drags to draw a
+                  rubber-band selection. get_frame_selection() returns the
+                  selected rect in frame pixel coordinates.
+    """
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self._frame_rgb:    Optional[np.ndarray] = None
+        self._frozen_frame: Optional[np.ndarray] = None
+        self._frozen   = False
+        self._zoom     = 1.0
+        self._pan      = QPointF(0.0, 0.0)
+        self._shape    = "rect"
+
+        self._rb_start: Optional[QPointF] = None
+        self._rb_end:   Optional[QPointF] = None
+
+        self._pan_anchor: Optional[QPointF] = None
+        self._pan_origin: Optional[QPointF] = None
+
+        self.setMouseTracking(True)
+        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        self.setMinimumSize(480, 320)
+        self.setStyleSheet("background: #0b0e13;")
+
+    # ── Public API ────────────────────────────────────────────────────────
+
+    def set_frame(self, rgb: np.ndarray):
+        if not self._frozen:
+            self._frame_rgb = rgb
+            self.update()
+
+    def set_shape(self, shape: str):
+        self._shape = shape
+        self.update()
+
+    def freeze(self):
+        """Lock display on current frame and enter selection mode."""
+        self._frozen       = True
+        self._frozen_frame = (
+            self._frame_rgb.copy() if self._frame_rgb is not None else None
+        )
+        self._rb_start = None
+        self._rb_end   = None
+        self.update()
+
+    def unfreeze(self):
+        """Resume live display; reset zoom/pan/selection."""
+        self._frozen       = False
+        self._frozen_frame = None
+        self._zoom         = 1.0
+        self._pan          = QPointF(0.0, 0.0)
+        self._rb_start     = None
+        self._rb_end       = None
+        self.update()
+
+    def get_frozen_frame(self) -> Optional[np.ndarray]:
+        return self._frozen_frame
+
+    def get_frame_selection(self) -> Optional[Tuple[int, int, int, int]]:
+        """Return (x, y, w, h) of selection in frame pixel coords, or None."""
+        if self._rb_start is None or self._rb_end is None:
+            return None
+        frame = self._frozen_frame
+        if frame is None:
+            return None
+        fh, fw = frame.shape[:2]
+        p1 = self._screen_to_frame(self._rb_start)
+        p2 = self._screen_to_frame(self._rb_end)
+        x  = int(max(0,  min(p1[0], p2[0])))
+        y  = int(max(0,  min(p1[1], p2[1])))
+        x2 = int(min(fw, max(p1[0], p2[0])))
+        y2 = int(min(fh, max(p1[1], p2[1])))
+        w, h = x2 - x, y2 - y
+        return (x, y, w, h) if w >= 8 and h >= 8 else None
+
+    # ── Coordinate transforms ─────────────────────────────────────────────
+
+    def _get_transform(self) -> Tuple[float, float, float]:
+        frame = self._frozen_frame if self._frozen else self._frame_rgb
+        if frame is None:
+            return 1.0, 0.0, 0.0
+        fh, fw = frame.shape[:2]
+        W, H   = self.width(), self.height()
+        ds = min(W / fw, H / fh)
+        ts = ds * self._zoom
+        ox = (W - fw * ts) / 2 + self._pan.x()
+        oy = (H - fh * ts) / 2 + self._pan.y()
+        return ts, ox, oy
+
+    def _screen_to_frame(self, pt: QPointF) -> Tuple[float, float]:
+        ts, ox, oy = self._get_transform()
+        return (pt.x() - ox) / ts, (pt.y() - oy) / ts
+
+    # ── Events ────────────────────────────────────────────────────────────
+
+    def wheelEvent(self, e):
+        frame = self._frozen_frame if self._frozen else self._frame_rgb
+        if frame is None:
+            e.ignore()
+            return
+        fh, fw = frame.shape[:2]
+        W, H   = self.width(), self.height()
+        ds     = min(W / fw, H / fh)
+
+        factor   = 1.15 if e.angleDelta().y() > 0 else 1 / 1.15
+        new_zoom = max(1.0, min(12.0, self._zoom * factor))
+        cursor   = QPointF(e.position())
+
+        old_ts = ds * self._zoom
+        new_ts = ds * new_zoom
+        old_ox = (W - fw * old_ts) / 2 + self._pan.x()
+        old_oy = (H - fh * old_ts) / 2 + self._pan.y()
+        fx = (cursor.x() - old_ox) / old_ts
+        fy = (cursor.y() - old_oy) / old_ts
+        self._pan = QPointF(
+            cursor.x() - fx * new_ts - (W - fw * new_ts) / 2,
+            cursor.y() - fy * new_ts - (H - fh * new_ts) / 2,
+        )
+        self._zoom = new_zoom
+        e.accept()
+        self.update()
+
+    def mousePressEvent(self, e: QMouseEvent):
+        if e.button() == Qt.MouseButton.RightButton:
+            self._pan_anchor = QPointF(e.position())
+            self._pan_origin = QPointF(self._pan)
+        elif e.button() == Qt.MouseButton.LeftButton and self._frozen:
+            self._rb_start = QPointF(e.position())
+            self._rb_end   = None
+            self.update()
+
+    def mouseMoveEvent(self, e: QMouseEvent):
+        if e.buttons() & Qt.MouseButton.RightButton and self._pan_anchor is not None:
+            self._pan = self._pan_origin + (QPointF(e.position()) - self._pan_anchor)
+            self.update()
+        elif (e.buttons() & Qt.MouseButton.LeftButton
+              and self._frozen and self._rb_start is not None):
+            self._rb_end = QPointF(e.position())
+            self.update()
+
+    def mouseReleaseEvent(self, e: QMouseEvent):
+        if e.button() == Qt.MouseButton.RightButton:
+            self._pan_anchor = None
+            self._pan_origin = None
+        elif e.button() == Qt.MouseButton.LeftButton and self._frozen:
+            if self._rb_start is not None:
+                self._rb_end = QPointF(e.position())
+                self.update()
+
+    def paintEvent(self, _):
+        p = QPainter(self)
+        p.fillRect(self.rect(), QColor(11, 14, 19))
+
+        frame = self._frozen_frame if self._frozen else self._frame_rgb
+
+        if frame is None:
+            p.setPen(QColor(120, 130, 150))
+            p.setFont(QFont("sans-serif", 11))
+            p.drawText(self.rect(), Qt.AlignmentFlag.AlignCenter,
+                       "Connect a camera to begin.")
+            return
+
+        # Draw frame with zoom/pan applied
+        ts, ox, oy = self._get_transform()
+        fh, fw = frame.shape[:2]
+        arr = np.ascontiguousarray(frame)
+        qi  = QImage(arr.data, fw, fh, 3 * fw, QImage.Format.Format_RGB888)
+        pix = QPixmap.fromImage(qi).scaled(
+            max(1, int(fw * ts)), max(1, int(fh * ts)),
+            Qt.AspectRatioMode.IgnoreAspectRatio,
+            Qt.TransformationMode.FastTransformation,
+        )
+        p.drawPixmap(int(ox), int(oy), pix)
+
+        # Rubber-band selection overlay
+        if self._frozen and self._rb_start is not None and self._rb_end is not None:
+            pen = QPen(QColor(0, 220, 100))
+            pen.setWidth(2)
+            pen.setStyle(Qt.PenStyle.DashLine)
+            p.setPen(pen)
+            sx = min(self._rb_start.x(), self._rb_end.x())
+            sy = min(self._rb_start.y(), self._rb_end.y())
+            rw = abs(self._rb_end.x() - self._rb_start.x())
+            rh = abs(self._rb_end.y() - self._rb_start.y())
+            if self._shape == "circle":
+                p.drawEllipse(int(sx), int(sy), int(rw), int(rh))
+            else:
+                p.drawRect(int(sx), int(sy), int(rw), int(rh))
+
+        # Instruction banner (frozen mode)
+        if self._frozen:
+            p.fillRect(0, 0, self.width(), 26, QColor(0, 0, 0, 170))
+            p.setPen(QColor(0, 212, 255))
+            p.setFont(QFont("sans-serif", 9))
+            p.drawText(
+                0, 0, self.width(), 26,
+                Qt.AlignmentFlag.AlignCenter,
+                "Scroll = zoom  ·  Right-drag = pan  ·  Left-drag = select target",
+            )
+
+
+# ── Main window ───────────────────────────────────────────────────────────────
 
 class MotionCaptureWindow(QWidget):
 
@@ -30,17 +248,17 @@ class MotionCaptureWindow(QWidget):
         super().__init__(parent, Qt.WindowType.Window | Qt.WindowType.Tool)
         self._tracker   = tracker
         self._connected = False
-        self._arming    = False   # waiting for click-to-set-ROI
+        self._selecting = False   # True while frame is frozen for selection
 
         self.setWindowTitle("Motion Capture")
         self.setMinimumSize(680, 560)
-        self.resize(720, 580)
+        self.resize(760, 600)
 
         self._build()
 
         self._refresh = QTimer(self)
         self._refresh.timeout.connect(self._update)
-        self._refresh.start(33)   # ~30 fps display
+        self._refresh.start(33)
 
     # ── UI ────────────────────────────────────────────────────────────────
 
@@ -68,15 +286,7 @@ class MotionCaptureWindow(QWidget):
         root.addWidget(div)
 
         # Feed
-        self._feed = QLabel()
-        self._feed.setObjectName("motionFeed")
-        self._feed.setAlignment(Qt.AlignmentFlag.AlignCenter)
-        self._feed.setSizePolicy(
-            QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
-        self._feed.setMinimumSize(480, 320)
-        self._feed.setStyleSheet("background: #0b0e13;")
-        self._feed.setText("Connect a camera to begin.")
-        self._feed.mousePressEvent = self._on_feed_click
+        self._feed = FrameSelector()
         root.addWidget(self._feed, 1)
 
         # Status bar
@@ -91,44 +301,62 @@ class MotionCaptureWindow(QWidget):
         ctrl = QHBoxLayout(ctrl_frame)
         ctrl.setContentsMargins(10, 6, 10, 6); ctrl.setSpacing(6)
 
+        # Camera selector
         ctrl.addWidget(QLabel("Camera:"))
         self._cam_combo = QComboBox()
         self._cam_combo.setObjectName("channelPickerCb")
-        self._cam_combo.setMinimumWidth(200)
+        self._cam_combo.setMinimumWidth(180)
         ctrl.addWidget(self._cam_combo)
-
         self._refresh_btn = QPushButton("⟳")
         self._refresh_btn.setObjectName("configButton")
         self._refresh_btn.setFixedWidth(28)
         self._refresh_btn.setToolTip("Rescan cameras")
         self._refresh_btn.clicked.connect(self._scan_cameras)
         ctrl.addWidget(self._refresh_btn)
-
-        self._scan_cameras()   # populate on open
+        self._scan_cameras()
 
         self._conn_btn = QPushButton("Connect")
         self._conn_btn.setObjectName("applyButton")
         self._conn_btn.clicked.connect(self._toggle_connect)
         ctrl.addWidget(self._conn_btn)
 
-        sep = QFrame(); sep.setFrameShape(QFrame.Shape.VLine)
-        sep.setObjectName("devWindowDivider"); ctrl.addWidget(sep)
+        ctrl.addWidget(_vsep())
+
+        # Tracking mode
+        ctrl.addWidget(QLabel("Mode:"))
+        self._mode_cb = QComboBox()
+        self._mode_cb.setObjectName("channelPickerCb")
+        self._mode_cb.addItems(["Template Match", "CSRT"])
+        self._mode_cb.currentIndexChanged.connect(self._on_mode_changed)
+        ctrl.addWidget(self._mode_cb)
+
+        # Shape (template mode only)
+        self._shape_lbl = QLabel("Shape:")
+        ctrl.addWidget(self._shape_lbl)
+        self._shape_cb = QComboBox()
+        self._shape_cb.setObjectName("channelPickerCb")
+        self._shape_cb.addItems(["Rectangle", "Circle"])
+        self._shape_cb.currentIndexChanged.connect(self._on_shape_changed)
+        ctrl.addWidget(self._shape_cb)
+
+        ctrl.addWidget(_vsep())
 
+        # Track / confirm + clear
         self._track_btn = QPushButton("🎯  Click to Track")
         self._track_btn.setObjectName("toolbarSectionBtn")
         self._track_btn.setEnabled(False)
-        self._track_btn.clicked.connect(self._arm_tracking)
+        self._track_btn.clicked.connect(self._on_track_btn)
         ctrl.addWidget(self._track_btn)
 
         self._clear_btn = QPushButton("✕  Clear")
         self._clear_btn.setObjectName("configButton")
         self._clear_btn.setEnabled(False)
-        self._clear_btn.clicked.connect(self._clear_tracking)
+        self._clear_btn.clicked.connect(self._on_clear)
         ctrl.addWidget(self._clear_btn)
 
-        sep2 = QFrame(); sep2.setFrameShape(QFrame.Shape.VLine)
-        sep2.setObjectName("devWindowDivider"); ctrl.addWidget(sep2)
+        ctrl.addWidget(_vsep())
 
+        # px/mm calibration
         ctrl.addWidget(QLabel("px/mm:"))
         self._pxmm_spin = QDoubleSpinBox()
         self._pxmm_spin.setRange(0.01, 100_000)
@@ -149,10 +377,9 @@ class MotionCaptureWindow(QWidget):
         current_idx = self.get_camera_index()
         self._cam_combo.blockSignals(True)
         self._cam_combo.clear()
-        self._cameras = list_cameras()   # [(index, label), ...]
+        self._cameras = list_cameras()
         for _, label in self._cameras:
             self._cam_combo.addItem(label)
-        # Restore selection by index if still present
         for i, (idx, _) in enumerate(self._cameras):
             if idx == current_idx:
                 self._cam_combo.setCurrentIndex(i)
@@ -161,109 +388,117 @@ class MotionCaptureWindow(QWidget):
 
     def _toggle_connect(self):
         if self._connected:
+            if self._selecting:
+                self._feed.unfreeze()
+                self._selecting = False
             self._tracker.stop()
             self._connected = False
-            self._arming    = False
-            self._conn_btn.setText("Connect")
-            self._track_btn.setEnabled(False)
-            self._track_btn.setText("🎯  Click to Track")
-            self._clear_btn.setEnabled(False)
-            self._feed.setPixmap(QPixmap())
-            self._feed.setText("Connect a camera to begin.")
             self._sim_badge.setVisible(False)
             self._status.setText("Disconnected")
         else:
             ok = self._tracker.start(self.get_camera_index())
             if ok:
                 self._connected = True
-                self._conn_btn.setText("Disconnect")
-                self._track_btn.setEnabled(True)
-                self._clear_btn.setEnabled(True)
                 self._sim_badge.setVisible(self._tracker.simulated)
             else:
                 self._status.setText("Could not open camera")
+        self._sync_controls()
+
+    def _on_track_btn(self):
+        if self._selecting:
+            sel = self._feed.get_frame_selection()
+            if sel is None:
+                self._status.setText(
+                    "Draw a selection box first, then click Confirm")
+                return
+            x, y, w, h = sel
+            self._tracker.set_roi_rect(x, y, w, h,
+                                        frame_rgb=self._feed.get_frozen_frame())
+            self._feed.unfreeze()
+            self._selecting = False
+        else:
+            self._feed.freeze()
+            self._selecting = True
+        self._sync_controls()
+
+    def _on_clear(self):
+        if self._selecting:
+            self._feed.unfreeze()
+            self._selecting = False
+        else:
+            self._tracker.clear_roi()
+        self._sync_controls()
+
+    def _on_mode_changed(self, idx: int):
+        mode = "csrt" if idx == 1 else "template"
+        self._tracker.set_track_mode(mode)
+        is_template = mode == "template"
+        self._shape_lbl.setVisible(is_template)
+        self._shape_cb.setVisible(is_template)
+        if self._selecting:
+            self._feed.unfreeze()
+            self._selecting = False
+        self._sync_controls()
+
+    def _on_shape_changed(self, idx: int):
+        shape = "circle" if idx == 1 else "rect"
+        self._tracker.set_template_shape(shape)
+        self._feed.set_shape(shape)
+
+    def _sync_controls(self):
+        connected = self._connected
+        selecting = self._selecting
+        tracking  = self._tracker.is_tracking() if self._tracker else False
+
+        self._conn_btn.setText("Disconnect" if connected else "Connect")
+        self._track_btn.setEnabled(connected)
+        self._clear_btn.setEnabled(connected and (selecting or tracking))
+
+        if selecting:
+            self._track_btn.setText("✓  Confirm Selection")
+        else:
+            self._track_btn.setText("🎯  Click to Track")
 
-    def _arm_tracking(self):
-        self._arming = True
-        self._track_btn.setText("⊹  Click on target…")
-        self._feed.setCursor(Qt.CursorShape.CrossCursor)
-
-    def _clear_tracking(self):
-        self._arming = False
-        self._tracker.clear_roi()
-        self._track_btn.setText("🎯  Click to Track")
-        self._feed.unsetCursor()
-
-    def _on_feed_click(self, event: QMouseEvent):
-        if not self._arming or not self._connected:
-            return
-
-        frame = self._tracker.get_frame()
-        if frame is None:
-            return
-
-        lw = self._feed.width()
-        lh = self._feed.height()
-        fh, fw = frame.shape[:2]
-        scale  = min(lw / fw, lh / fh)
-        off_x  = (lw - int(fw * scale)) // 2
-        off_y  = (lh - int(fh * scale)) // 2
-
-        fx = int((event.position().x() - off_x) / scale)
-        fy = int((event.position().y() - off_y) / scale)
-
-        if 0 <= fx < fw and 0 <= fy < fh:
-            self._tracker.set_roi(fx, fy)
-
-        self._arming = False
-        self._track_btn.setText("🎯  Click to Track")
-        self._feed.unsetCursor()
-
-    # ── Display update ────────────────────────────────────────────────────
+    # ── Display update ─────────────────────────────────────────────────────
 
     def _update(self):
         if not self._connected:
             return
 
-        frame = self._tracker.get_frame()
-        if frame is not None:
-            self._show_frame(frame)
+        if not self._selecting:
+            frame = self._tracker.get_frame()
+            if frame is not None:
+                self._feed.set_frame(frame)
 
-        pos = self._tracker.get_position()
-        fps = self._tracker.fps
+        pos       = self._tracker.get_position()
+        fps       = self._tracker.fps
         px_per_mm = self._pxmm_spin.value()
 
-        if self._tracker.is_lost():
+        if self._selecting:
+            sel = self._feed.get_frame_selection()
+            if sel:
+                x, y, w, h = sel
+                self._status.setText(
+                    f"Selected: {w} × {h} px at ({x}, {y})  ·  Click Confirm to track")
+            else:
+                self._status.setText(
+                    "Left-drag to select target  ·  Scroll to zoom  ·  Right-drag to pan")
+        elif self._tracker.is_lost():
             self._status.setText(
                 f"Tracking lost — click 'Click to Track' to reselect  |  {fps:.0f} FPS")
+            self._sync_controls()
         elif pos is not None:
             xmm = pos[0] / px_per_mm
             ymm = pos[1] / px_per_mm
             self._status.setText(
                 f"X: {pos[0]:.1f} px  ({xmm:.2f} mm)   "
                 f"Y: {pos[1]:.1f} px  ({ymm:.2f} mm)   |   {fps:.0f} FPS")
-        elif self._arming:
-            self._status.setText("Click on the object you want to track")
         else:
             sim = "  [SIMULATION]" if self._tracker.simulated else ""
             self._status.setText(
                 f"Live{sim}  |  {fps:.0f} FPS  |  Click 'Click to Track' to begin")
 
-    def _show_frame(self, rgb: np.ndarray):
-        h, w = rgb.shape[:2]
-        bytes_per_line = 3 * w
-        qi = QImage(
-            rgb.data, w, h, bytes_per_line,
-            QImage.Format.Format_RGB888,
-        )
-        pix = QPixmap.fromImage(qi).scaled(
-            self._feed.width(), self._feed.height(),
-            Qt.AspectRatioMode.KeepAspectRatio,
-            Qt.TransformationMode.FastTransformation,
-        )
-        self._feed.setPixmap(pix)
-
-    # ── State accessors (used by plugin for save/restore) ─────────────────
+    # ── State accessors ────────────────────────────────────────────────────
 
     def get_camera_index(self) -> int:
         i = self._cam_combo.currentIndex()
@@ -291,3 +526,12 @@ class MotionCaptureWindow(QWidget):
         self._refresh.stop()
         self.closed.emit()
         e.accept()
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _vsep() -> QFrame:
+    sep = QFrame()
+    sep.setFrameShape(QFrame.Shape.VLine)
+    sep.setObjectName("devWindowDivider")
+    return sep