""" motion_capture/window.py Floating tool window for the Motion Capture plugin. Layout: [Title bar] [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 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, ) # ── 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): closed = pyqtSignal() def __init__(self, tracker, parent=None): super().__init__(parent, Qt.WindowType.Window) self._tracker = tracker self._selecting = False self.setWindowTitle("Motion Capture") self.setMinimumSize(680, 560) self.resize(760, 600) self._build() # Reflect tracker's actual state — it may already be running self._connected = tracker.is_running() self._sim_badge.setVisible(self._connected and tracker.simulated) self._sync_controls() self._refresh = QTimer(self) self._refresh.timeout.connect(self._update) self._refresh.start(33) # ── UI ──────────────────────────────────────────────────────────────── def _build(self): root = QVBoxLayout(self) root.setContentsMargins(0, 0, 0, 0) root.setSpacing(0) # Title bar hdr = QWidget(); hdr.setObjectName("devWindowTitleBar") hdr.setFixedHeight(44) hl = QHBoxLayout(hdr); hl.setContentsMargins(14, 0, 14, 0) self._title_lbl = QLabel("MOTION CAPTURE") self._title_lbl.setObjectName("devWindowTitle") hl.addWidget(self._title_lbl) hl.addStretch() self._sim_badge = QLabel("SIMULATION") self._sim_badge.setObjectName("traceSource") self._sim_badge.setVisible(False) hl.addWidget(self._sim_badge) root.addWidget(hdr) div = QFrame(); div.setFrameShape(QFrame.Shape.HLine) div.setObjectName("devWindowDivider") root.addWidget(div) # Feed self._feed = FrameSelector() root.addWidget(self._feed, 1) # Status bar self._status = QLabel("Not connected") self._status.setObjectName("traceSource") self._status.setAlignment(Qt.AlignmentFlag.AlignCenter) self._status.setFixedHeight(24) root.addWidget(self._status) # Controls bar ctrl_frame = QWidget(); ctrl_frame.setObjectName("cfgBottomBar") 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(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() self._conn_btn = QPushButton("Connect") self._conn_btn.setObjectName("applyButton") self._conn_btn.clicked.connect(self._toggle_connect) ctrl.addWidget(self._conn_btn) 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._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._on_clear) ctrl.addWidget(self._clear_btn) ctrl.addWidget(_vsep()) # px/mm calibration ctrl.addWidget(QLabel("px/mm:")) self._pxmm_spin = QDoubleSpinBox() self._pxmm_spin.setRange(0.01, 100_000) self._pxmm_spin.setValue(10.0) self._pxmm_spin.setDecimals(2) self._pxmm_spin.setSuffix(" px/mm") self._pxmm_spin.setObjectName("cfgGlobalSpin") self._pxmm_spin.setFixedWidth(130) ctrl.addWidget(self._pxmm_spin) ctrl.addStretch() root.addWidget(ctrl_frame) # ── Actions ─────────────────────────────────────────────────────────── def _scan_cameras(self): from tracker import list_cameras current_idx = self.get_camera_index() self._cam_combo.blockSignals(True) self._cam_combo.clear() self._cameras = list_cameras() for _, label in self._cameras: self._cam_combo.addItem(label) for i, (idx, _) in enumerate(self._cameras): if idx == current_idx: self._cam_combo.setCurrentIndex(i) break self._cam_combo.blockSignals(False) def _toggle_connect(self): if self._connected: if self._selecting: self._feed.unfreeze() self._selecting = False self._tracker.stop() self._connected = False self._sim_badge.setVisible(False) self._status.setText("Disconnected") else: ok = self._tracker.start(self.get_camera_index()) if ok: self._connected = 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") # ── Display update ───────────────────────────────────────────────────── def _update(self): if not self._connected: return 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 px_per_mm = self._pxmm_spin.value() 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") else: sim = " [SIMULATION]" if self._tracker.simulated else "" self._status.setText( f"Live{sim} | {fps:.0f} FPS | Click 'Click to Track' to begin") # ── State accessors ──────────────────────────────────────────────────── def restore_ui_state(self, mode: str, shape: str): """Sync mode/shape UI to tracker state without triggering tracker resets.""" self._mode_cb.blockSignals(True) self._shape_cb.blockSignals(True) self._mode_cb.setCurrentIndex(1 if mode == "csrt" else 0) self._shape_cb.setCurrentIndex(1 if shape == "circle" else 0) is_template = mode == "template" self._shape_lbl.setVisible(is_template) self._shape_cb.setVisible(is_template) self._mode_cb.blockSignals(False) self._shape_cb.blockSignals(False) self._feed.set_shape(shape) def get_camera_index(self) -> int: i = self._cam_combo.currentIndex() cameras = getattr(self, "_cameras", []) if 0 <= i < len(cameras): return cameras[i][0] return 0 def set_camera_index(self, v: int): cameras = getattr(self, "_cameras", []) for i, (idx, _) in enumerate(cameras): if idx == v: self._cam_combo.setCurrentIndex(i) return def get_px_per_mm(self) -> float: return self._pxmm_spin.value() def set_px_per_mm(self, v: float): self._pxmm_spin.setValue(v) # ── Close ───────────────────────────────────────────────────────────── def closeEvent(self, e: QCloseEvent): self._refresh.stop() self.closed.emit() e.accept() # ── Helpers ─────────────────────────────────────────────────────────────────── def _vsep() -> QFrame: sep = QFrame() sep.setFrameShape(QFrame.Shape.VLine) sep.setObjectName("devWindowDivider") return sep