diff --git a/src/glider/gui/behavior/analysis_window.py b/src/glider/gui/behavior/analysis_window.py index 307503ef..167cd1e9 100644 --- a/src/glider/gui/behavior/analysis_window.py +++ b/src/glider/gui/behavior/analysis_window.py @@ -630,14 +630,15 @@ def _build_top_bar(self) -> QFrame: ("Retrieval tables (CSV)…", self._export_retrieval), ("Nest distance per frame (CSV)…", self._export_nest_distance), ("Heatmap (PNG + CSV)…", self._export_heatmap), + ("Current frame (PNG)…", self._export_frame), ], ) self._export_window_action = export_actions[0] - self._export_retrieval_action, self._export_nest_action = export_actions[-3:-1] + self._export_retrieval_action, self._export_nest_action = export_actions[-4:-2] # Shown under a pup-retrieval preset only (_apply_experiment_type). for action in (self._export_retrieval_action, self._export_nest_action): action.setVisible(False) - self._export_heatmap_btn = export_actions[-1] + self._export_heatmap_btn = export_actions[-2] self._export_window_action.setEnabled(False) self._export_heatmap_btn.setEnabled(False) self._tour_btn, _ = self._menu_button( @@ -3613,6 +3614,30 @@ def _export_heatmap(self) -> None: note = "" if png_path is not None else " (no figure: matplotlib is not installed)" self.statusBar().showMessage(f"Wrote {wrote}{note}", 10000) + def _export_frame(self) -> None: + """Save the frame on screen, with the overlays that are switched on, as a PNG.""" + image = self._canvas.render_frame() + if image is None: + QMessageBox.information( + self, "Export frame", "Open a session with a video or poses first." + ) + return + frame = self._canvas.frame + if self._ethogram_csv is not None: + session_dir = ( + self._ethogram_csv if self._ethogram_csv.is_dir() else self._ethogram_csv.parent + ) + default = str(session_dir / f"{session_dir.name}_frame{frame}.png") + else: + default = f"frame{frame}.png" + path, _ = QFileDialog.getSaveFileName(self, "Export frame", default, "PNG Files (*.png)") + if not path: + return + if not image.save(path, "PNG"): + QMessageBox.critical(self, "Export frame", f"Could not write {path}") + return + self.statusBar().showMessage(f"Wrote {path}", 10000) + def _describe(self, stats) -> str: span = ( f"{stats.start_frame}–{stats.end_frame} " diff --git a/src/glider/gui/dialogs/arena_dialog.py b/src/glider/gui/dialogs/arena_dialog.py index e1b1f245..2498672d 100644 --- a/src/glider/gui/dialogs/arena_dialog.py +++ b/src/glider/gui/dialogs/arena_dialog.py @@ -1,8 +1,9 @@ """Arena Dialog - outline the floor to calibrate scale and perspective. The operator clicks the four floor corners of an arena of known size, or, for -a round arena, its centre and the ends of two radii, or, for an elevated plus -maze, the far ends of its four arms. That +a round arena, drags a circle out from its centre and then moves and stretches +it onto the rim, or, for an elevated plus maze, clicks the far ends of its four +arms. That fixes a homography (see :mod:`glider.vision.arena`), which in turn gives a centre zone meaning the same patch of floor in every video, and a scale that varies correctly across the floor instead of being one number for all of it. @@ -31,7 +32,16 @@ import cv2 import numpy as np from PyQt6.QtCore import QPoint, QRect, Qt, pyqtSignal -from PyQt6.QtGui import QColor, QFont, QImage, QMouseEvent, QPainter, QPen, QPixmap +from PyQt6.QtGui import ( + QColor, + QFont, + QImage, + QMouseEvent, + QPainter, + QPen, + QPixmap, + QWheelEvent, +) from PyQt6.QtWidgets import ( QComboBox, QDialog, @@ -81,6 +91,20 @@ "plus": tuple(f"end of the {name}" for name in ARM_END_NAMES), } +#: Radius a round arena starts at when placed with a click rather than dragged +#: out, as a fraction of the frame's shorter side: most rigs frame the arena +#: closely, so this lands near the rim and the handles finish the fit. +_CIRCLE_DEFAULT_RADIUS = 0.4 + +#: Canvas pixels a press must travel before it is a drag rather than a click. +_DRAG_THRESHOLD = 5 + +#: How much one wheel notch scales a round arena. +_WHEEL_STEP = 1.03 + +#: Rim handles of a round arena: which axis each moves, and which end of it. +_RIM_HANDLES = (("u", 1), ("v", 1), ("u", -1), ("v", -1)) + #: Size a plus maze starts at when the operator has not typed one: a standard #: mouse EPM is two 35 cm arms either side of a 5 cm centre. _PLUS_DEFAULT_CM = 75.0 @@ -107,6 +131,12 @@ def __init__(self, parent=None): self._pixmap: QPixmap | None = None self._corners: list[tuple[float, float]] = [] self._dragging: int | None = None + # A round arena is dragged as a whole: "create", "move", or a rim + # handle from _RIM_HANDLES. Kept apart from _dragging, which indexes a + # clicked point. + self._ellipse_drag: str | tuple[str, int] | None = None + self._press_norm: tuple[float, float] | None = None + self._press_corners: list[tuple[float, float]] = [] self._zone_cm = 10.0 self._arena_cm = (30.0, 30.0) self._shape = "rectangle" @@ -162,11 +192,15 @@ def is_complete(self) -> bool: def clear(self) -> None: self._corners = [] self._dragging = None + self._ellipse_drag = None self.update() self.corners_changed.emit() def undo(self) -> None: - if self._corners: + if self._shape == "ellipse": + # A circle is placed in one gesture, so it is undone in one too. + self.clear() + elif self._corners: self._corners.pop() self.update() self.corners_changed.emit() @@ -243,9 +277,155 @@ def _to_canvas(self, point) -> QPoint: # -- interaction --------------------------------------------------- + # -- round arena --------------------------------------------------- + # + # Geometry is done in frame pixels, not normalised coordinates: those are + # stretched by the frame's aspect ratio, so a circle there is an ellipse on + # screen, and "perpendicular" means nothing. + + def _px(self, norm) -> np.ndarray: + fh, fw = self._frame.shape[:2] + return np.array([norm[0] * fw, norm[1] * fh], dtype=np.float64) + + def _from_px(self, px) -> tuple[float, float]: + fh, fw = self._frame.shape[:2] + return (float(px[0] / fw), float(px[1] / fh)) + + def _ellipse_axes(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Centre and the two *principal* semi-axes, in frame pixels. + + The stored handles may be any conjugate pair -- an arena drawn by the + old three-click tool, or one passed to :meth:`set_corners` -- which + describe the same ellipse but put handles in odd places. An SVD of the + pair gives the principal axes of that same ellipse. ``u`` is the more + horizontal one, pointing right, and ``v`` turns clockwise from it on + screen, so the corners from_ellipse derives keep their TL-TR-BR-BL + order and the exploration grid is not mirrored. + """ + c, a, b = (self._px(p) for p in self._corners) + left, scale, _ = np.linalg.svd(np.column_stack([a - c, b - c])) + u, v = left[:, 0] * scale[0], left[:, 1] * scale[1] + if abs(v[0]) > abs(u[0]): + u, v = v, u + if u[0] < 0: + u = -u + if u[0] * v[1] - u[1] * v[0] < 0: + v = -v + return c, u, v + + def _set_ellipse(self, c, u, v) -> None: + self._corners = [self._from_px(c), self._from_px(c + u), self._from_px(c + v)] + self.update() + self.corners_changed.emit() + + def _ellipse_hit(self, pos: QPoint): + """What a press at *pos* grabs: a rim handle, "move", or None.""" + c, u, v = self._ellipse_axes() + for axis, sign in _RIM_HANDLES: + end = c + sign * (u if axis == "u" else v) + if (self._to_canvas(self._from_px(end)) - pos).manhattanLength() <= _GRAB_RADIUS: + return (axis, sign) + # Inside the ellipse, or on the centre of one too small to get inside. + p = self._px(self._to_norm(pos)) - c + inside = (p @ u / (u @ u)) ** 2 + (p @ v / (v @ v)) ** 2 <= 1 + near_centre = (self._to_canvas(self._corners[0]) - pos).manhattanLength() <= _GRAB_RADIUS + return "move" if inside or near_centre else None + + def _ellipse_press(self, pos: QPoint) -> None: + norm = self._to_norm(pos) + if norm is None: + return + if self.is_complete: + # A miss grabs nothing: Clear starts over, and a stray click beside + # a fitted circle should not throw the fit away. + self._ellipse_drag = self._ellipse_hit(pos) + else: + fh, fw = self._frame.shape[:2] + r = _CIRCLE_DEFAULT_RADIUS * min(fw, fh) + self._set_ellipse(self._px(norm), np.array([r, 0.0]), np.array([0.0, r])) + self._ellipse_drag = "create" + self._press_norm = norm + self._press_corners = list(self._corners) + if self._ellipse_drag == "move": + self.setCursor(Qt.CursorShape.ClosedHandCursor) + + def _ellipse_move(self, event: QMouseEvent) -> None: + norm = self._to_norm(event.pos()) + if norm is None: + return + drag = self._ellipse_drag + mods = event.modifiers() + if drag == "create": + if ( + event.pos() - self._to_canvas(self._press_norm) + ).manhattanLength() < _DRAG_THRESHOLD: + return # still a click: keep the default circle + c = self._px(self._corners[0]) + r = max(float(np.linalg.norm(self._px(norm) - c)), 2.0) + self._set_ellipse(c, np.array([r, 0.0]), np.array([0.0, r])) + elif drag == "move": + delta = self._px(norm) - self._px(self._press_norm) + self._corners = [self._from_px(self._px(p) + delta) for p in self._press_corners] + self.update() + self.corners_changed.emit() + else: + self._stretch(*drag, self._px(norm), mods) + + def _stretch(self, axis: str, sign: int, p: np.ndarray, mods) -> None: + """Drag one rim handle to *p*, keeping the opposite rim point fixed. + + The handle moves along its own axis; Ctrl lets it turn the ellipse too, + and Shift makes the other axis follow, keeping the arena round. + """ + c, u, v = self._ellipse_axes() + a, b = (u, v) if axis == "u" else (v, u) + fixed = c - sign * a + direction = sign * a / np.linalg.norm(a) # from the fixed point outwards + if mods & Qt.KeyboardModifier.ControlModifier: + span = p - fixed + if np.linalg.norm(span) > 1e-9: + direction = span / np.linalg.norm(span) + reach = max(float((p - fixed) @ direction), 4.0) # never through the far rim + a_new = sign * direction * reach / 2 + b_len = reach / 2 if mods & Qt.KeyboardModifier.ShiftModifier else np.linalg.norm(b) + # b turns with a, staying on the side it was on. + perp = np.array([-direction[1], direction[0]]) + if perp @ b < 0: + perp = -perp + b_new = perp * b_len + c_new = fixed + sign * a_new + if axis == "u": + self._set_ellipse(c_new, a_new, b_new) + else: + self._set_ellipse(c_new, b_new, a_new) + + def _ellipse_hover(self, pos: QPoint) -> None: + hit = self._ellipse_hit(pos) if self.is_complete else None + if hit == "move": + self.setCursor(Qt.CursorShape.OpenHandCursor) + elif hit is not None: + self.setCursor(Qt.CursorShape.SizeAllCursor) + else: + self.setCursor(Qt.CursorShape.CrossCursor) + + def wheelEvent(self, event: QWheelEvent) -> None: # noqa: N802 - Qt override + """Over a round arena, the wheel scales it about its centre.""" + if self._shape != "ellipse" or not self.is_complete or self._frame is None: + return super().wheelEvent(event) + notches = event.angleDelta().y() / 120 + if not notches: + return super().wheelEvent(event) + c, u, v = self._ellipse_axes() + factor = _WHEEL_STEP**notches + self._set_ellipse(c, u * factor, v * factor) + event.accept() + def mousePressEvent(self, event: QMouseEvent) -> None: if self._frame is None or event.button() != Qt.MouseButton.LeftButton: return super().mousePressEvent(event) + if self._shape == "ellipse": + self._ellipse_press(event.pos()) + return pos = event.pos() for i, corner in enumerate(self._corners): if (self._to_canvas(corner) - pos).manhattanLength() <= _GRAB_RADIUS: @@ -259,21 +439,25 @@ def mousePressEvent(self, event: QMouseEvent) -> None: self.corners_changed.emit() def mouseMoveEvent(self, event: QMouseEvent) -> None: + if self._shape == "ellipse" and self._frame is not None: + if self._ellipse_drag is not None: + self._ellipse_move(event) + else: + self._ellipse_hover(event.pos()) + return if self._dragging is None: return super().mouseMoveEvent(event) norm = self._to_norm(event.pos()) if norm is not None: - if self._shape == "ellipse" and self._dragging == 0: - # Moving the centre carries the whole ellipse with it. - dx = norm[0] - self._corners[0][0] - dy = norm[1] - self._corners[0][1] - self._corners = [(x + dx, y + dy) for x, y in self._corners] self._corners[self._dragging] = norm self.update() self.corners_changed.emit() def mouseReleaseEvent(self, event: QMouseEvent) -> None: self._dragging = None + if self._ellipse_drag is not None: + self._ellipse_drag = None + self._ellipse_hover(event.pos()) super().mouseReleaseEvent(event) # -- painting ------------------------------------------------------ @@ -347,6 +531,9 @@ def _paint_outline(self, painter: QPainter) -> None: painter.drawLine(zone[i], zone[(i + 1) % len(zone)]) def _paint_corners(self, painter: QPainter) -> None: + if self._shape == "ellipse" and self.is_complete: + self._paint_ellipse_handles(painter) + return font = QFont() font.setPointSize(9) font.setBold(True) @@ -358,14 +545,25 @@ def _paint_corners(self, painter: QPainter) -> None: painter.drawLine(point.x() - 9, point.y(), point.x() + 9, point.y()) painter.drawLine(point.x(), point.y() - 9, point.x(), point.y() + 9) painter.setPen(QColor(colors.TEXT_PRIMARY)) - if self._shape == "ellipse": - label = ("C", "1", "2")[i] - elif self._shape == "plus": + if self._shape == "plus": label = ("T", "R", "B", "L")[i] else: label = str(i + 1) painter.drawText(point.x() + 10, point.y() - 8, label) + def _paint_ellipse_handles(self, painter: QPainter) -> None: + """A crosshair at the centre and a square grip at each axis end.""" + c, u, v = self._ellipse_axes() + centre = self._to_canvas(self._corners[0]) + painter.setPen(QPen(_CORNER_COLOUR, 2)) + painter.drawLine(centre.x() - 9, centre.y(), centre.x() + 9, centre.y()) + painter.drawLine(centre.x(), centre.y() - 9, centre.x(), centre.y() + 9) + painter.setBrush(_CORNER_COLOUR) + for axis, sign in _RIM_HANDLES: + end = self._to_canvas(self._from_px(c + sign * (u if axis == "u" else v))) + painter.drawRect(end.x() - 5, end.y() - 5, 10, 10) + painter.setBrush(Qt.BrushStyle.NoBrush) + class ArenaDialog(QDialog): """Outline one video's arena floor: four corners, a circle's rim, or a plus maze.""" @@ -480,9 +678,10 @@ def _shape_changed(self) -> None: ) elif round_arena: self.instructions.setText( - "Click the centre of the floor, then a point on its edge, then a second " - "edge point about a quarter-turn round. Click where the bedding meets " - "the wall. Drag the centre to move the circle, an edge point to fit it." + "Press at the centre of the floor and drag out to its edge, where the " + "bedding meets the wall. Then drag inside the circle to move it, drag a " + "square handle to stretch it (Shift keeps it round, Ctrl turns it), or " + "use the mouse wheel to resize it." ) else: self.instructions.setText( @@ -530,6 +729,11 @@ def _refresh(self) -> None: self.status.setStyleSheet(f"color: {colors.ERROR};") ok_button.setEnabled(False) return + if cal is None and self.canvas.shape == "ellipse": + self.status.setText("Press at the centre of the floor and drag out to its edge.") + self.status.setStyleSheet(f"color: {colors.TEXT_SECONDARY};") + ok_button.setEnabled(False) + return if cal is None: names = _POINT_NAMES[self.canvas.shape] placed = len(self.canvas.corners) diff --git a/src/glider/gui/pose_batch/arena_actions.py b/src/glider/gui/pose_batch/arena_actions.py index 16c9d824..2f26bcaa 100644 --- a/src/glider/gui/pose_batch/arena_actions.py +++ b/src/glider/gui/pose_batch/arena_actions.py @@ -132,7 +132,10 @@ def regate_videos( else: gated += 1 if on_log: - on_log(f"{video.name}: blanked {report.blanked_fraction:.1%}") + on_log( + f"{video.name}: blanked {report.blanked_fraction:.1%}, " + f"{report.keypoints_jumped} jumped keypoints" + ) if on_progress: on_progress(index + 1, len(videos)) return gated, skipped diff --git a/src/glider/gui/review/viewer.py b/src/glider/gui/review/viewer.py index 09e8ceae..3532e0eb 100644 --- a/src/glider/gui/review/viewer.py +++ b/src/glider/gui/review/viewer.py @@ -22,6 +22,7 @@ QPainter, QPainterPath, QPen, + QPolygonF, ) from PyQt6.QtWidgets import ( QCheckBox, @@ -124,6 +125,8 @@ def __init__(self, parent=None): self._grid_n = 0 self._visited = None self._show_arena = True + # Set only while render_frame paints an export image. + self._paint_size: tuple[int, int] | None = None self._reader = None # VideoFileSource, opened lazily self._video_failed = False # this view's video would not open; retried on set_view self._cached: tuple[int, QImage] | None = None @@ -293,22 +296,54 @@ def set_frame(self, frame: int) -> None: self._frame = int(frame) self.update() + @property + def frame(self) -> int: + return self._frame + def set_trail(self, seconds: float, enabled: bool) -> None: self._trail_s, self._show_trail = float(seconds), bool(enabled) self.update() + def _size(self) -> tuple[int, int]: + """What is being painted on: the widget, or an export image.""" + return self._paint_size or (self.width(), self.height()) + def _transform(self): """Scale and offset mapping video pixels onto the widget, or None.""" if self._view is None or not self._view.resolution: return None width, height = self._view.resolution - scale = min(self.width() / width, self.height() / height) - return scale, (self.width() - width * scale) / 2, (self.height() - height * scale) / 2 + w, h = self._size() + scale = min(w / width, h / height) + return scale, (w - width * scale) / 2, (h - height * scale) / 2 + + def render_frame(self) -> QImage | None: + """The current frame with every overlay that is switched on, at video resolution. + + Painted by the same code as the screen, so what is exported is what the + reviewer was looking at -- less the frame-number caption, which is + viewer chrome, not part of the picture. + """ + if self._view is None or not self._view.resolution: + return None + width, height = (int(v) for v in self._view.resolution) + image = QImage(width, height, QImage.Format.Format_RGB32) + self._paint_size = (width, height) + painter = QPainter(image) + try: + self._paint(painter, chrome=False) + finally: + painter.end() + self._paint_size = None + return image def paintEvent(self, _event): painter = QPainter(self) + self._paint(painter) + + def _paint(self, painter: QPainter, *, chrome: bool = True) -> None: painter.setRenderHint(QPainter.RenderHint.Antialiasing) - painter.fillRect(self.rect(), QColor(colors.CANVAS)) + painter.fillRect(QRectF(0, 0, *self._size()), QColor(colors.CANVAS)) transform = self._transform() if ( @@ -318,7 +353,7 @@ def paintEvent(self, _event): ): painter.setPen(QPen(QColor(colors.TEXT_MUTED))) painter.drawText( - self.rect(), + QRectF(0, 0, *self._size()), Qt.AlignmentFlag.AlignCenter | Qt.TextFlag.TextWordWrap, self._why_blank(), ) @@ -333,8 +368,9 @@ def paintEvent(self, _event): painter.drawImage(arena, image) if self._heatmap is not None: painter.drawImage(arena, self._heatmap) - painter.setPen(QPen(QColor(colors.BORDER), 1)) - painter.drawRect(arena) + if chrome: + painter.setPen(QPen(QColor(colors.BORDER), 1)) + painter.drawRect(arena) if self._show_arena and self._arena is not None: self._paint_arena(painter, arena) if self._show_zones: @@ -378,10 +414,12 @@ def to_widget(point): painter.restore() if self._hud is None: + if not chrome: + return label = self._view.label_at(self._frame) painter.setPen(QPen(QColor(colors.TEXT_PRIMARY))) painter.drawText( - QRectF(8, 6, self.width() - 16, 20), + QRectF(8, 6, self._size()[0] - 16, 20), Qt.AlignmentFlag.AlignLeft, f"frame {self._frame} {label or '(unscored)'}", ) @@ -417,20 +455,27 @@ def to_widget(points): colour = QColor(colors.SUCCESS) try: + width, height = self._view.resolution + outline = [(x / width, y / height) for x, y in self._arena.outline(width, height)] n = self._grid_n if n >= 1 and self._visited is not None and self._visited.shape == (n, n): + # Cells are squares of the arena's bounding box; clipped to the + # floor so a round arena or a plus maze is not filled square. + floor = QPainterPath() + floor.addPolygon(QPolygonF(to_widget(outline))) + painter.save() + painter.setClipPath(floor) painter.setPen(Qt.PenStyle.NoPen) painter.setBrush(colors.qcolor_with_alpha(colour, 0.18)) for row, col in zip(*np.nonzero(self._visited), strict=True): painter.drawPolygon( *to_widget(exploration.cell_polygon(self._arena, n, int(col), int(row))) ) + painter.restore() if n >= 2: painter.setPen(QPen(colors.qcolor_with_alpha(colour, 0.9), 1.5)) for line in exploration.grid_segments(self._arena, n): painter.drawPolyline(*to_widget(line)) - width, height = self._view.resolution - outline = [(x / width, y / height) for x, y in self._arena.outline(width, height)] except ValueError: # a degenerate arena: nothing sensible to draw return painter.setPen(QPen(colour, 2)) @@ -685,12 +730,12 @@ def _paint_hud(self, painter: QPainter) -> None: self._chip(painter, 12.0, 10.0, *behavior, right=False, large=True) y = 10.0 for text, colour in chips[: self._MAX_CHIPS]: - self._chip(painter, self.width() - 12.0, y, text, colour, right=True) + self._chip(painter, self._size()[0] - 12.0, y, text, colour, right=True) y += 24 if len(chips) > self._MAX_CHIPS: extra = f"+{len(chips) - self._MAX_CHIPS}" self._chip( - painter, self.width() - 12.0, y, extra, QColor(colors.TEXT_MUTED), right=True + painter, self._size()[0] - 12.0, y, extra, QColor(colors.TEXT_MUTED), right=True ) def _chip_font(self, large: bool) -> QFont: diff --git a/src/glider/vision/arena_gate.py b/src/glider/vision/arena_gate.py index 713919df..78a0c911 100644 --- a/src/glider/vision/arena_gate.py +++ b/src/glider/vision/arena_gate.py @@ -18,6 +18,13 @@ the quad, so a bare containment test would delete real rearing -- invisibly, which is worse than leaving a visible glitch. +The margin also means containment alone cannot catch a keypoint that jumps +*within* the arena -- to a reflection, an object, the far corner -- and back. +The jump gate does that, on the same floor coordinates: a keypoint's track is +cut wherever it would have to move faster than ``max_speed_cm_s``, and short +islands between cuts are blanked, shortest first, so the track the animal +spends most of its time on is the one that survives. + Qt-free on purpose: ``run_batch`` and a GUI button both drive this, and a notebook can too. """ @@ -50,12 +57,26 @@ #: catching the bench-floor detections this exists for. _DEFAULT_MARGIN_FRACTION = 0.25 +#: Above any real rodent locomotion -- a mouse sprints at roughly 1 m/s, a rat +#: not much faster in an open field -- and far below a detector relocating +#: across a 30-60 cm arena in one frame, which is 900+ cm/s at 30 fps. +_DEFAULT_MAX_SPEED_CM_S = 200.0 + +#: Islands up to this long are blanked. Longer than any single-frame glitch +#: and short enough that a sustained relocation is left visible rather than +#: decided by a majority vote over several seconds. +_DEFAULT_MAX_EXCURSION_S = 1.0 + @dataclass(frozen=True) class ArenaGateSettings: margin_cm: float | None = None min_inside_fraction: float = 0.5 min_detected_fraction: float = 0.0 + # None turns the jump gate off. It is also what a block written before the + # jump gate existed reads back as -- see settings_from_block. + max_speed_cm_s: float | None = _DEFAULT_MAX_SPEED_CM_S + max_excursion_s: float = _DEFAULT_MAX_EXCURSION_S def margin_for(self, arena: ArenaCalibration) -> float: if self.margin_cm is not None: @@ -72,6 +93,9 @@ class GateReport: masked_by_keypoint: dict[str, int] = field(default_factory=dict) settings: ArenaGateSettings = field(default_factory=ArenaGateSettings) arena_corners: list = field(default_factory=list) + #: Keypoints the jump gate blanked, in frames the arena test kept. + keypoints_jumped: int = 0 + jumped_by_keypoint: dict[str, int] = field(default_factory=dict) @property def blanked_fraction(self) -> float: @@ -155,6 +179,14 @@ def _outside(arena, xy_px, resolution, margin_cm) -> np.ndarray: keypoint is not "outside" -- it is simply not present, which the caller accounts for separately. """ + x, y = _floor_cm(arena, xy_px, resolution) + with np.errstate(invalid="ignore"): + # An ellipse lies inside its bounding box, so the argument above holds. + return np.asarray(arena.outside_cm(x, y, margin_cm), dtype=bool) + + +def _floor_cm(arena, xy_px, resolution) -> tuple[np.ndarray, np.ndarray]: + """``(T, K)`` floor coordinates in centimetres; see :func:`_outside` on ``w``.""" flat = np.asarray(xy_px, dtype=np.float64).reshape(-1, 2) width, height = resolution homogeneous = np.stack( @@ -164,9 +196,44 @@ def _outside(arena, xy_px, resolution, margin_cm) -> np.ndarray: w = projected[2] with np.errstate(invalid="ignore", divide="ignore"): x, y = projected[0] / w, projected[1] / w - # An ellipse lies inside its bounding box, so the argument above holds. - beyond = arena.outside_cm(x, y, margin_cm) - return beyond.reshape(np.asarray(xy_px).shape[:2]) + shape = np.asarray(xy_px).shape[:2] + return x.reshape(shape), y.reshape(shape) + + +def _jumps(x, y, valid, fps: float, max_speed_cm_s: float, max_excursion_s: float) -> np.ndarray: + """Boolean ``(T, K)``: valid keypoints on a short island between jumps. + + Per keypoint, the valid samples are cut wherever the step to the next one + would need more than ``max_speed_cm_s``. The allowance grows with the gap + in frames, so a keypoint that was missing for a while may reappear further + away. The shortest island is then blanked, provided it spans no more than + ``max_excursion_s``, and the cuts are recomputed: once a glitch is gone, + the frames either side of it are usually reachable from each other, and + the track rejoins. Shortest first is what makes the result independent of + which end of the video the glitch is at -- a teleport on frame 0 loses the + same way a teleport mid-session does. + """ + out = np.zeros(valid.shape, dtype=bool) + max_frames = max_excursion_s * fps + for k in range(valid.shape[1]): + frames = np.flatnonzero(valid[:, k]) + keep = np.ones(len(frames), dtype=bool) + while True: + live = frames[keep] + if len(live) < 2: + break + step = np.hypot(np.diff(x[live, k]), np.diff(y[live, k])) + cuts = np.flatnonzero(step > max_speed_cm_s * np.diff(live) / fps) + if not len(cuts): + break + bounds = np.concatenate(([0], cuts + 1, [len(live)])) + spans = live[bounds[1:] - 1] - live[bounds[:-1]] + 1 + shortest = int(np.argmin(spans)) + if spans[shortest] > max_frames: + break + keep[np.flatnonzero(keep)[bounds[shortest] : bounds[shortest + 1]]] = False + out[frames[~keep], k] = True + return out def gate_to_arena(pose, arena, *, settings=None, resolution=None): @@ -219,6 +286,23 @@ def gate_to_arena(pose, arena, *, settings=None, resolution=None): # the frame verdict, not by their own position. blank is (T,), so the # trailing axis is added to broadcast against stray's (T, K). counted = stray & ~blank[:, None] + + # Jumps last, over what containment kept: a stray is already gone and + # would otherwise cut the track around itself for no reason. + jumped = np.zeros_like(stray) + if settings.max_speed_cm_s is not None: + x, y = _floor_cm(arena, pose.xy, resolution) + jumped = _jumps( + x, + y, + detected & ~stray & ~blank[:, None], + float(pose.fps), + float(settings.max_speed_cm_s), + float(settings.max_excursion_s), + ) + out.xy[jumped] = np.nan + out.confidence[jumped] = 0.0 + names = list(pose.keypoint_names) return out, GateReport( frames_total=int(pose.n_frames), @@ -228,6 +312,8 @@ def gate_to_arena(pose, arena, *, settings=None, resolution=None): masked_by_keypoint={n: int(counted[:, i].sum()) for i, n in enumerate(names)}, settings=settings, arena_corners=corners, + keypoints_jumped=int(jumped.sum()), + jumped_by_keypoint={n: int(jumped[:, i].sum()) for i, n in enumerate(names)}, ) @@ -241,8 +327,18 @@ def gate_to_arena(pose, arena, *, settings=None, resolution=None): "keypoints_masked", "masked_by_keypoint", "arena_corners", + "keypoints_jumped", + "jumped_by_keypoint", ) +#: Every field a gate block recorded before the jump gate existed. Writers have +#: always stored ``asdict(settings)``, so a block spelling out exactly these is +#: from a build that had no jump gate -- its silence on ``max_speed_cm_s`` means +#: "none", not today's default. Without this an old file would claim a jump gate +#: it never had and pool with jump-gated sessions. A terse block that records +#: only what was changed is still read as defaults, as before. +_PRE_JUMP_FIELDS = frozenset({"margin_cm", "min_inside_fraction", "min_detected_fraction"}) + def settings_from_block(block) -> dict: """The settings a stored ``arena_gate`` block records, defaults filled in. @@ -258,7 +354,9 @@ def settings_from_block(block) -> dict: understand is still a difference between two blocks, and silently ignoring it is how a real mismatch would pass. """ - return {**asdict(ArenaGateSettings()), **((block or {}).get("settings") or {})} + stored = (block or {}).get("settings") or {} + legacy = {"max_speed_cm_s": None} if set(stored) == _PRE_JUMP_FIELDS else {} + return {**asdict(ArenaGateSettings()), **legacy, **stored} def _report_from_block(block) -> GateReport: diff --git a/tests/unit/gui/behavior/test_analysis_window.py b/tests/unit/gui/behavior/test_analysis_window.py index 78e3c431..c7657e61 100644 --- a/tests/unit/gui/behavior/test_analysis_window.py +++ b/tests/unit/gui/behavior/test_analysis_window.py @@ -2619,6 +2619,7 @@ def test_the_export_menu(self, qtbot, tmp_path): "Retrieval tables (CSV)…", "Nest distance per frame (CSV)…", "Heatmap (PNG + CSV)…", + "Current frame (PNG)…", ] def test_the_markers_tab_opens_the_epoch_table(self, qtbot, tmp_path): diff --git a/tests/unit/gui/pose_batch/test_window_arena.py b/tests/unit/gui/pose_batch/test_window_arena.py index 6695dc9a..44970234 100644 --- a/tests/unit/gui/pose_batch/test_window_arena.py +++ b/tests/unit/gui/pose_batch/test_window_arena.py @@ -671,7 +671,7 @@ def flaky(csv, arena, settings=None): calls.append(csv) if len(calls) == 1: raise ValueError("refused") - return SimpleNamespace(blanked_fraction=0.0) + return SimpleNamespace(blanked_fraction=0.0, keypoints_jumped=0) monkeypatch.setattr("glider.vision.arena_gate.gate_pose_csv", flaky) gated, skipped = regate_videos(videos, calibrations, on_log=lambda m: None) diff --git a/tests/unit/gui/review/test_viewer.py b/tests/unit/gui/review/test_viewer.py index fd67555c..0d998823 100644 --- a/tests/unit/gui/review/test_viewer.py +++ b/tests/unit/gui/review/test_viewer.py @@ -86,6 +86,57 @@ def test_the_hud_draws_the_behaviour_chip(canvas): assert canvas.grab().toImage().pixelColor(*centre) != colour +def _round_arena(): + from glider.vision.arena import ArenaCalibration + + # A 400 px circle filling the middle of the 640x480 frame. + return ArenaCalibration.from_ellipse( + (0.5, 0.5), + (0.5 + 200 / 640, 0.5), + (0.5, 0.5 + 200 / 480), + width_cm=30.0, + frame_size=(640, 480), + ) + + +class TestRenderFrame: + def test_it_is_the_video_resolution(self, canvas): + canvas.set_view(_view()) + image = canvas.render_frame() + assert (image.width(), image.height()) == (640, 480) + + def test_it_has_the_overlays_and_not_the_frame_caption(self, canvas): + canvas.set_view(_view()) + canvas.set_frame(10) + image = canvas.render_frame() + # The pose dot at (10, 100) is in the export... + assert image.pixelColor(10, 100) != QColor(colors.CANVAS) + # ...and the "frame N" caption, which the widget draws at top-left, is not. + caption = [image.pixelColor(x, 12) for x in range(8, 120)] + assert all(c == QColor(colors.CANVAS) for c in caption) + + def test_it_does_not_disturb_the_widget(self, canvas): + canvas.set_view(_view()) + before = canvas.grab().toImage() + canvas.render_frame() + assert canvas.grab().toImage() == before + + def test_without_a_session_there_is_nothing(self, canvas): + assert canvas.render_frame() is None + + +def test_visited_cells_are_filled_only_inside_a_round_arena(canvas): + canvas.set_view(_view()) + canvas.set_show_poses(False) + canvas.set_arena(_round_arena()) + canvas.set_grid(3) + canvas.set_visited(np.ones((3, 3), dtype=bool)) + image = canvas.render_frame() + # The corner of the bounding box is in a visited cell but off the floor. + assert image.pixelColor(130, 50) == QColor(colors.CANVAS) + assert image.pixelColor(320, 240) != QColor(colors.CANVAS) + + def test_the_behaviour_chip_is_larger_than_an_output_chip(canvas): """It is what a reviewer reads while scrubbing; the output chips are secondary.""" assert BEHAVIOUR_CHIP_H >= 1.5 * CHIP_H diff --git a/tests/unit/gui/test_arena_dialog.py b/tests/unit/gui/test_arena_dialog.py index 28e493c8..fa65f1b6 100644 --- a/tests/unit/gui/test_arena_dialog.py +++ b/tests/unit/gui/test_arena_dialog.py @@ -254,3 +254,139 @@ def test_the_canvas_paints_a_finished_maze(self, canvas): canvas.set_corners(ARM_ENDS) canvas.repaint() # must not raise assert canvas.calibration().shape == "plus" + + +def _mouse(kind, pos: QPoint, modifiers=None): + from PyQt6.QtCore import QPointF + from PyQt6.QtCore import Qt as _Qt + from PyQt6.QtGui import QMouseEvent + + types = { + "press": QMouseEvent.Type.MouseButtonPress, + "move": QMouseEvent.Type.MouseMove, + "release": QMouseEvent.Type.MouseButtonRelease, + } + button = _Qt.MouseButton.NoButton if kind == "move" else _Qt.MouseButton.LeftButton + return QMouseEvent( + types[kind], + QPointF(pos), + button, + _Qt.MouseButton.LeftButton, + modifiers or _Qt.KeyboardModifier.NoModifier, + ) + + +def _drag(canvas, start, end, modifiers=None): + a, b = canvas._to_canvas(start), canvas._to_canvas(end) + canvas.mousePressEvent(_mouse("press", a)) + canvas.mouseMoveEvent(_mouse("move", b, modifiers)) + canvas.mouseReleaseEvent(_mouse("release", b)) + + +def _axes(canvas): + c, u, v = canvas._ellipse_axes() + return c, np.linalg.norm(u), np.linalg.norm(v) + + +#: A 100 px circle at the centre of the 640x480 frame, dragged out. +_CENTRE, _EDGE = (0.5, 0.5), (0.5 + 100 / 640, 0.5) + + +class TestCircleTool: + """A round arena is dragged out, then moved and stretched as one shape.""" + + @pytest.fixture + def round_canvas(self, canvas): + canvas.set_shape("ellipse") + return canvas + + def test_press_and_drag_draws_a_circle_out_from_the_centre(self, round_canvas): + _drag(round_canvas, _CENTRE, _EDGE) + assert round_canvas.is_complete + c, ru, rv = _axes(round_canvas) + np.testing.assert_allclose(c, [320, 240], atol=2) + assert ru == pytest.approx(100, abs=2) and rv == pytest.approx(100, abs=2) + + def test_a_click_places_a_default_circle(self, round_canvas): + _drag(round_canvas, _CENTRE, _CENTRE) + _, ru, rv = _axes(round_canvas) + assert ru == pytest.approx(0.4 * 480) and rv == pytest.approx(0.4 * 480) + + def test_dragging_inside_moves_it(self, round_canvas): + _drag(round_canvas, _CENTRE, _EDGE) + _drag(round_canvas, (0.45, 0.45), (0.45 + 32 / 640, 0.45 + 24 / 480)) + c, ru, _ = _axes(round_canvas) + np.testing.assert_allclose(c, [352, 264], atol=2) + assert ru == pytest.approx(100, abs=2) + + def test_a_rim_handle_stretches_one_axis_with_the_far_rim_fixed(self, round_canvas): + _drag(round_canvas, _CENTRE, _EDGE) + _drag(round_canvas, (420 / 640, 0.5), (460 / 640, 0.5)) + c, ru, rv = _axes(round_canvas) + assert ru == pytest.approx(120, abs=2) and rv == pytest.approx(100, abs=2) + assert c[0] - ru == pytest.approx(220, abs=2) # the left rim did not move + + def test_shift_keeps_it_round(self, round_canvas): + from PyQt6.QtCore import Qt as _Qt + + _drag(round_canvas, _CENTRE, _EDGE) + _drag(round_canvas, (420 / 640, 0.5), (460 / 640, 0.5), _Qt.KeyboardModifier.ShiftModifier) + _, ru, rv = _axes(round_canvas) + assert ru == pytest.approx(120, abs=2) and rv == pytest.approx(ru, abs=1) + + def test_a_press_outside_a_fitted_circle_does_nothing(self, round_canvas): + _drag(round_canvas, _CENTRE, (0.5 + 50 / 640, 0.5)) + before = round_canvas.corners + _drag(round_canvas, (0.05, 0.05), (0.1, 0.1)) + assert round_canvas.corners == before + + def test_the_wheel_scales_it(self, round_canvas): + from PyQt6.QtCore import QPointF + from PyQt6.QtCore import Qt as _Qt + from PyQt6.QtGui import QWheelEvent + + _drag(round_canvas, _CENTRE, _EDGE) + at = QPointF(round_canvas._to_canvas(_CENTRE)) + round_canvas.wheelEvent( + QWheelEvent( + at, + at, + QPoint(0, 0), + QPoint(0, 120), + _Qt.MouseButton.NoButton, + _Qt.KeyboardModifier.NoModifier, + _Qt.ScrollPhase.NoScrollPhase, + False, + ) + ) + _, ru, _ = _axes(round_canvas) + assert ru == pytest.approx(103, abs=1) + + def test_an_old_three_click_arena_reopens_as_the_same_ellipse(self, round_canvas): + from glider.vision.arena import ArenaCalibration + + # Conjugate, not principal: the second edge is not a quarter-turn round. + old_handles = [(0.5, 0.5), (0.7, 0.5), (0.6, 0.8)] + old = ArenaCalibration.from_ellipse(*old_handles, width_cm=30.0, frame_size=(640, 480)) + round_canvas.set_corners(old_handles) + c, u, v = round_canvas._ellipse_axes() + assert abs(u @ v) < 1e-6 # principal axes + # Every rim point of the re-expressed ellipse lies on the old one. + t = np.linspace(0, 2 * np.pi, 24, endpoint=False) + rim = [round_canvas._from_px(c + u * np.cos(a) + v * np.sin(a)) for a in t] + radii = np.hypot(*(np.array(old.to_arena_cm(rim)) - 15.0).T) + np.testing.assert_allclose(radii, 15.0, atol=1e-6) + + def test_the_grid_is_not_mirrored(self, round_canvas): + _drag(round_canvas, _CENTRE, _EDGE) + tl, tr, br, bl = round_canvas.calibration().corners + assert tl[0] < tr[0] and tl[1] < bl[1] + + def test_undo_clears_the_circle(self, round_canvas): + _drag(round_canvas, _CENTRE, _EDGE) + round_canvas.undo() + assert round_canvas.corners == [] + + def test_the_canvas_paints_a_fitted_circle(self, round_canvas): + _drag(round_canvas, _CENTRE, _EDGE) + round_canvas.repaint() # must not raise diff --git a/tests/unit/vision/test_arena_gate.py b/tests/unit/vision/test_arena_gate.py index c8e2fe8b..57ab0d41 100644 --- a/tests/unit/vision/test_arena_gate.py +++ b/tests/unit/vision/test_arena_gate.py @@ -429,3 +429,63 @@ def test_a_sidecar_less_csv_still_gates(self, tmp_path): meta_path(csv).unlink() gate_pose_csv(csv, _arena(), settings=ArenaGateSettings(margin_cm=7.5)) assert (tmp_path / f"{csv.stem}_ungated.csv").exists() + + +def _walk(n=60, jump_at=(), to=(500.0, 60.0)): + """Keypoint ``a`` strolling across SQUARE's centre, teleporting on *jump_at*.""" + xy = np.full((n, 4, 2), 320.0) + xy[:, :, 1] = 240.0 + xy[:, 0, 0] = 300.0 + np.arange(n) * 0.5 # ~1 cm/s: well under any limit + for t in jump_at: + xy[t, 0] = to + return xy + + +class TestJumps: + """Containment cannot see a keypoint that jumps somewhere still inside the arena.""" + + def test_a_short_teleport_inside_the_arena_is_blanked(self): + # (500, 60) is inside the square, so containment alone keeps it. + gated, report = gate_to_arena(_pose(_walk(jump_at=(20, 21, 22))), _arena()) + assert np.isnan(gated.xy[20:23, 0]).all() + assert np.isfinite(gated.xy[:20, 0]).all() and np.isfinite(gated.xy[23:, 0]).all() + assert report.keypoints_jumped == 3 + assert report.jumped_by_keypoint == {"a": 3, "b": 0, "c": 0, "d": 0} + assert report.frames_blanked == 0 # the other keypoints were fine + + def test_a_teleport_on_the_first_frames_loses_too(self): + gated, _ = gate_to_arena(_pose(_walk(jump_at=(0, 1))), _arena()) + assert np.isnan(gated.xy[:2, 0]).all() + assert np.isfinite(gated.xy[2:, 0]).all() + + def test_a_relocation_longer_than_the_excursion_is_kept(self): + settings = ArenaGateSettings(max_excursion_s=0.5) # 15 frames + gated, report = gate_to_arena( + _pose(_walk(jump_at=range(20, 40))), _arena(), settings=settings + ) + assert report.keypoints_jumped == 0 + assert np.isfinite(gated.xy[:, 0]).all() + + def test_none_turns_the_jump_gate_off(self): + settings = ArenaGateSettings(max_speed_cm_s=None) + gated, report = gate_to_arena(_pose(_walk(jump_at=(20,))), _arena(), settings=settings) + assert report.keypoints_jumped == 0 + assert np.isfinite(gated.xy[20, 0]).all() + + def test_fast_but_real_running_is_kept(self): + xy = _walk() + # 100 cm/s at 30 fps on a 13.3 px/cm floor is 44 px a frame. + xy[:, 0, 0] = 150.0 + np.minimum(np.arange(60), 8) * 44.0 + _, report = gate_to_arena(_pose(xy), _arena()) + assert report.keypoints_jumped == 0 + + def test_a_block_from_before_the_jump_gate_reads_as_no_jump_gate(self): + from dataclasses import asdict + + from glider.vision.arena_gate import settings_from_block + + old = asdict(ArenaGateSettings()) + del old["max_speed_cm_s"], old["max_excursion_s"] + assert settings_from_block({"gated": True, "settings": old})["max_speed_cm_s"] is None + new = {"gated": True, "settings": asdict(ArenaGateSettings())} + assert settings_from_block(new)["max_speed_cm_s"] == ArenaGateSettings().max_speed_cm_s