diff options
Diffstat (limited to 'plugins/motion_capture/window.py')
| -rw-r--r-- | plugins/motion_capture/window.py | 444 |
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 |