From ce633b7325c9f5ed4f5e7314142e5491c0bad070 Mon Sep 17 00:00:00 2001 From: Glen Beane <356266+gbeane@users.noreply.github.com> Date: Thu, 18 Jun 2026 01:33:38 -0400 Subject: [PATCH 1/2] Add optional pose/segmentation overlay copy to frame export --- docs/user-guide/gui.md | 2 +- src/jabs/resources/docs/user_guide/gui.md | 2 +- src/jabs/ui/main_window/menu_handlers.py | 76 +++++--- .../ui/player_widget/overlays/pose_overlay.py | 84 ++------- src/jabs/ui/player_widget/player_widget.py | 98 ++++++++++- src/jabs/ui/player_widget/pose_drawing.py | 109 ++++++++++++ tests/ui/test_menu_handlers.py | 164 +++++++++--------- tests/ui/test_pose_drawing.py | 125 +++++++++++++ 8 files changed, 484 insertions(+), 176 deletions(-) create mode 100644 src/jabs/ui/player_widget/pose_drawing.py create mode 100644 tests/ui/test_pose_drawing.py diff --git a/docs/user-guide/gui.md b/docs/user-guide/gui.md index 19058412..f303cfa7 100644 --- a/docs/user-guide/gui.md +++ b/docs/user-guide/gui.md @@ -115,7 +115,7 @@ Clicking the video or moving the mouse off the video frame will dismiss the slid - **JABS→Quit JABS:** Quit Program - **File→Open Project:** Select a project directory to open. If a project is already opened, it will be closed and the newly selected project will be opened. - **File→Open Recent:** Submenu to open recently opened projects. -- **File→Export Frame:** Export the current paused video frame as a PNG image. The export uses the original video resolution with no overlays, is disabled during playback, and remembers the last export directory. Shortcut: Ctrl+E / ⌘E. +- **File→Export Frame:** Export the current paused video frame as a PNG image at the original video resolution. A checkbox on the save dialog optionally saves a second copy, suffixed `-overlay.png`, with the pose and (when available) segmentation overlays. The export is disabled during playback and remembers the last export directory and the overlay-copy setting. Shortcut: Ctrl+E / ⌘E. - **File→Export Training Data:** Create a file with the information needed to share a classifier. This exported file is written to the project directory and has the form `_training_.h5`. This file is used as one input for the `jabs-classify` script. - **File→Archive Behavior:** Remove behavior and its labels from project. Labels are archived in the `jabs/archive` directory. - **File→Prune Project:** Remove videos and pose files that are not labeled. diff --git a/src/jabs/resources/docs/user_guide/gui.md b/src/jabs/resources/docs/user_guide/gui.md index aec7c611..a53a654d 100644 --- a/src/jabs/resources/docs/user_guide/gui.md +++ b/src/jabs/resources/docs/user_guide/gui.md @@ -115,7 +115,7 @@ Clicking the video or moving the mouse off the video frame will dismiss the slid - **JABS→Quit JABS:** Quit Program - **File→Open Project:** Select a project directory to open. If a project is already opened, it will be closed and the newly selected project will be opened. - **File→Open Recent:** Submenu to open recently opened projects. -- **File→Export Frame:** Export the current paused video frame as a PNG image. The export uses the original video resolution with no overlays, is disabled during playback, and remembers the last export directory. +- **File→Export Frame:** Export the current paused video frame as a PNG image at the original video resolution. A checkbox on the save dialog optionally saves a second copy, suffixed `-overlay.png`, with the pose and (when available) segmentation overlays. The export is disabled during playback and remembers the last export directory and the overlay-copy setting. - **File→Export Training Data:** Create a file with the information needed to share a classifier. This exported file is written to the project directory and has the form `_training_.h5`. This file is used as one input for the `jabs-classify` script. - **File→Archive Behavior:** Remove behavior and its labels from project. Labels are archived in the `jabs/archive` directory. - **File→Prune Project:** Remove videos and pose files that are not labeled. diff --git a/src/jabs/ui/main_window/menu_handlers.py b/src/jabs/ui/main_window/menu_handlers.py index d57b8714..50aed8fb 100644 --- a/src/jabs/ui/main_window/menu_handlers.py +++ b/src/jabs/ui/main_window/menu_handlers.py @@ -45,6 +45,8 @@ # Keep the existing key so users retain the last-used directory across the rename. _SETTINGS_EXPORT_FRAME_DIR = "ui/save_frame_last_dir" +# Remembers whether the "also save overlay copy" checkbox was last checked. +_SETTINGS_EXPORT_OVERLAY = "ui/save_frame_overlay_copy" class UpdateCheckThread(QtCore.QThread): @@ -104,10 +106,11 @@ def export_frame(self) -> None: """Export the current video frame as a PNG image. Opens a file-save dialog pre-populated with a suggested filename derived from - the current video name and frame number. The last-used directory is persisted - in QSettings and restored on the next invocation; if that directory no longer - exists the dialog falls back to OS-default behaviour. If the user cancels, no - action is taken. + the current video name and frame number. The dialog includes a checkbox to also + save a second copy, suffixed ``-overlay.png``, with the pose and (when available) + segmentation overlays drawn at native resolution. The last-used directory and the + checkbox state are persisted in QSettings and restored on the next invocation. If + the user cancels, no action is taken. """ # noinspection PyProtectedMember player = self.window._central_widget._player_widget @@ -118,41 +121,68 @@ def export_frame(self) -> None: else: suggested_name = f"frame{frame_number:06d}.png" + # A checkbox is added to the dialog below, which native OS dialogs cannot host, + # so Qt's own (non-native) save dialog is used here regardless of the global + # native-dialog preference. + dialog = QtWidgets.QFileDialog(self.window, "Export Frame") + dialog.setOption(QtWidgets.QFileDialog.Option.DontUseNativeDialog, True) + dialog.setAcceptMode(QtWidgets.QFileDialog.AcceptMode.AcceptSave) + dialog.setNameFilter("PNG Images (*.png)") + dialog.setDefaultSuffix("png") last_dir = self.window._settings.value(_SETTINGS_EXPORT_FRAME_DIR, "", type=str) if last_dir and Path(last_dir).is_dir(): - initial_path = str(Path(last_dir) / suggested_name) - else: - initial_path = suggested_name + dialog.setDirectory(last_dir) + dialog.selectFile(suggested_name) - save_path, _ = QtWidgets.QFileDialog.getSaveFileName( - self.window, - "Export Frame", - initial_path, - "PNG Images (*.png)", - options=( - QtWidgets.QFileDialog.Option(0) - if USE_NATIVE_FILE_DIALOG - else QtWidgets.QFileDialog.Option.DontUseNativeDialog - ), + overlay_checkbox = QtWidgets.QCheckBox( + "Also save a copy with pose/segmentation overlay (…-overlay.png)" + ) + overlay_checkbox.setChecked( + self.window._settings.value(_SETTINGS_EXPORT_OVERLAY, False, type=bool) ) + layout = dialog.layout() + if isinstance(layout, QtWidgets.QGridLayout): + layout.addWidget(overlay_checkbox, layout.rowCount(), 0, 1, layout.columnCount()) - if not save_path: + if dialog.exec() != QtWidgets.QFileDialog.DialogCode.Accepted: return # user cancelled - - pixmap = player.get_raw_frame(frame_number) - if pixmap is None: - MessageDialog.warning(self.window, message="No frame available to export.") + selected_files = dialog.selectedFiles() + if not selected_files: return + save_path = selected_files[0] + save_overlay = overlay_checkbox.isChecked() if not save_path.lower().endswith(".png"): save_path += ".png" + pixmap = player.get_raw_frame(frame_number) + if pixmap is None: + MessageDialog.warning(self.window, message="No frame available to export.") + return if not pixmap.save(save_path, "PNG"): MessageDialog.error(self.window, message=f"Failed to export frame to:\n{save_path}") return + overlay_saved = False + if save_overlay: + overlay_path = str(Path(save_path).with_name(f"{Path(save_path).stem}-overlay.png")) + overlay_pixmap = player.get_overlay_frame(frame_number) + if overlay_pixmap is None: + MessageDialog.warning(self.window, message="No overlay frame available to export.") + elif not overlay_pixmap.save(overlay_path, "PNG"): + MessageDialog.error( + self.window, + message=f"Failed to export overlay frame to:\n{overlay_path}", + ) + else: + overlay_saved = True + self.window._settings.setValue(_SETTINGS_EXPORT_FRAME_DIR, str(Path(save_path).parent)) - self.window.display_status_message(f"Frame exported: {save_path}", 5000) + self.window._settings.setValue(_SETTINGS_EXPORT_OVERLAY, save_overlay) + status = f"Frame exported: {save_path}" + if overlay_saved: + status += " (+ overlay copy)" + self.window.display_status_message(status, 5000) def export_training_data(self) -> None: """Export training data for the current classifier in a background thread.""" diff --git a/src/jabs/ui/player_widget/overlays/pose_overlay.py b/src/jabs/ui/player_widget/overlays/pose_overlay.py index df847aa0..2cf04a81 100644 --- a/src/jabs/ui/player_widget/overlays/pose_overlay.py +++ b/src/jabs/ui/player_widget/overlays/pose_overlay.py @@ -1,22 +1,14 @@ from typing import TYPE_CHECKING -import numpy as np from PySide6 import QtCore, QtGui -from jabs.core.utils.pose_util import gen_line_fragments -from jabs.pose_estimation import PoseEstimation -from jabs.ui.colors import KEYPOINT_COLOR_MAP - +from ..pose_drawing import KEYPOINT_SIZE, draw_identity_pose from .overlay import Overlay if TYPE_CHECKING: from jabs.ui.player_widget.frame_with_overlays import FrameWithOverlaysWidget -_LINE_SEGMENT_COLOR = QtGui.QColor(255, 255, 255, 128) # color for the pose line segments -_KEYPOINT_SIZE = 3 # size of the keypoint circles - - class PoseOverlay(Overlay): """Overlay for displaying pose keypoints and connecting line segments on the video frame.""" @@ -58,71 +50,21 @@ def _overlay_pose( if self.parent.pose is None: return + # Keypoint size scales with the on-screen zoom so markers stay a sensible size. zoom = self.parent.scaled_pix_width / max(crop_rect.width(), 1) - keypoint_size = max(1, round(_KEYPOINT_SIZE * zoom**0.8)) + keypoint_size = max(1, round(KEYPOINT_SIZE * zoom**0.8)) - # draw the pose estimation skeletons for identity in self.parent.pose.identities: if not all_identities and identity != self.parent.active_identity: continue - points, mask = self.parent.pose.get_points(self.parent.current_frame, identity) - - if points is None or mask is None: - continue - - # Adjust alpha for non-active identities - if identity != self.parent.active_identity: - line_color = QtGui.QColor(_LINE_SEGMENT_COLOR) - line_color.setAlpha(line_color.alpha() // 3) # More translucent - else: - line_color = _LINE_SEGMENT_COLOR - - pen = QtGui.QPen(line_color) - pen.setWidth(3) - painter.setPen(pen) - - for seg in gen_line_fragments( - self.parent.pose.get_connected_segments(), np.flatnonzero(mask == 0) - ): - segment_points = [ - self.parent.image_to_widget_coords_cropped(p[0], p[1], crop_rect) - for p in points[seg] - ] - # Filter out points outside the crop - segment_points = [pt for pt in segment_points if pt is not None] - - # draw lines - if len(segment_points) >= 2: - for i in range(len(segment_points) - 1): - painter.drawLine( - segment_points[i][0], - segment_points[i][1], - segment_points[i + 1][0], - segment_points[i + 1][1], - ) - - # draw points at each keypoint of the pose (if it exists at this frame) - painter.setPen(QtCore.Qt.PenStyle.NoPen) - for keypoint in PoseEstimation.KeypointIndex: - point_index = keypoint.value - if mask[point_index]: - widget_coords = self.parent.image_to_widget_coords_cropped( - points[point_index][0], points[point_index][1], crop_rect - ) - if widget_coords is None: - continue - - widget_x, widget_y = widget_coords - color = KEYPOINT_COLOR_MAP[keypoint] - if identity != self.parent.active_identity: - # Make keypoints translucent for non-active identities - translucent_color = QtGui.QColor(color) - translucent_color.setAlpha(96) - painter.setBrush(translucent_color) - else: - painter.setBrush(color) - - painter.drawEllipse( - QtCore.QPoint(widget_x, widget_y), keypoint_size, keypoint_size - ) + draw_identity_pose( + painter, + self.parent.pose, + self.parent.current_frame, + identity, + to_output=lambda x, y: self.parent.image_to_widget_coords_cropped(x, y, crop_rect), + keypoint_size=keypoint_size, + line_width=3, + active=(identity == self.parent.active_identity), + ) diff --git a/src/jabs/ui/player_widget/player_widget.py b/src/jabs/ui/player_widget/player_widget.py index 99745aca..8487c642 100644 --- a/src/jabs/ui/player_widget/player_widget.py +++ b/src/jabs/ui/player_widget/player_widget.py @@ -5,12 +5,13 @@ import numpy as np from PySide6 import QtCore, QtGui, QtWidgets -from jabs.pose_estimation import PoseEstimation +from jabs.pose_estimation import PoseEstimation, PoseEstimationV6 from jabs.project import VideoLabels -from jabs.video_reader import VideoReader +from jabs.video_reader import VideoReader, overlay_segmentation from .frame_with_overlays import FrameWithOverlaysWidget from .player_thread import PlayerThread +from .pose_drawing import draw_identity_pose, native_pose_sizes _SPEED_VALUES = [0.5, 1, 2, 4] @@ -412,6 +413,99 @@ def get_raw_frame(self, frame_number: int | None = None) -> QtGui.QPixmap | None ) ) + def get_overlay_frame(self, frame_number: int | None = None) -> QtGui.QPixmap | None: + """Return a frame at original resolution with pose and segmentation overlays. + + Like :meth:`get_raw_frame`, but additionally renders the segmentation contours + (when the pose file provides them) and the pose keypoints/skeleton for every + identity, all at native video resolution. The overlays are always drawn, + independent of the on-screen overlay toggles. Unlike the live view, the export + has no active identity: every identity is drawn the same way, with pose skeletons + at full opacity and segmentation contours in a single color. Must only be called + when the video is not playing. + + Args: + frame_number: Frame index to export. Defaults to the currently selected frame. + + Returns: + A QPixmap at native video resolution with overlays drawn, or None if no + video is loaded. + + Raises: + RuntimeError: If called while the video is playing. + """ + if self._playing: + raise RuntimeError("get_overlay_frame() must not be called while the video is playing") + if self._video_stream is None: + return None + target_frame = self.current_frame if frame_number is None else frame_number + self._video_stream.seek(target_frame) + frame = self._video_stream.load_next_frame() + if frame["data"] is None: + return None + + img = frame["data"] + if img.dtype != np.uint8: + img = img.astype(np.uint8) + + # Bake segmentation contours into the BGR frame (a no-op when the pose file has + # no segmentation data for this identity/frame). + if isinstance(self._pose_est, PoseEstimationV6): + for ident in range(self._pose_est.num_identities): + # active=True for every identity: the export has no active identity, so + # all contours are drawn in the same (active) color rather than singling + # one out. + overlay_segmentation( + img, + self._pose_est, + identity=ident, + frame_index=target_frame, + active=True, + ) + + img_rgb = np.ascontiguousarray(img[..., ::-1]) # BGR → RGB + height, width, channels = img_rgb.shape + # QPixmap.fromImage copies the buffer, so it is safe to paint on afterwards + # without aliasing the numpy array backing the source QImage. + pixmap = QtGui.QPixmap.fromImage( + QtGui.QImage( + img_rgb.data, + width, + height, + channels * width, + QtGui.QImage.Format.Format_RGB888, + ) + ) + + # Draw the pose keypoints/skeleton on top, at native resolution, for every identity. + if self._pose_est is not None: + keypoint_size, line_width = native_pose_sizes(width, height) + + def to_native(x: float, y: float) -> tuple[int, int]: + # Coords come from numpy; round to whole native pixels as Python ints. + return round(float(x)), round(float(y)) + + painter = QtGui.QPainter(pixmap) + painter.setRenderHint(QtGui.QPainter.RenderHint.Antialiasing, True) + try: + for identity in self._pose_est.identities: + # active=True for every identity: all skeletons are drawn at full + # opacity, with no active-identity emphasis. + draw_identity_pose( + painter, + self._pose_est, + target_frame, + identity, + to_output=to_native, + keypoint_size=keypoint_size, + line_width=line_width, + active=True, + ) + finally: + painter.end() + + return pixmap + def _set_overlay_attr( self, attr: str, signal: QtCore.Signal | None, enabled: bool | None ) -> None: diff --git a/src/jabs/ui/player_widget/pose_drawing.py b/src/jabs/ui/player_widget/pose_drawing.py new file mode 100644 index 00000000..2eb0775f --- /dev/null +++ b/src/jabs/ui/player_widget/pose_drawing.py @@ -0,0 +1,109 @@ +"""Shared pose-skeleton drawing. + +Both the on-screen :class:`~jabs.ui.player_widget.overlays.pose_overlay.PoseOverlay` +(which draws at the scaled/cropped display resolution) and the full-resolution frame +export use :func:`draw_identity_pose`. The only thing that differs between the two is +how image coordinates map to the painter's coordinate space, so that mapping is passed +in as ``to_output``. +""" + +from collections.abc import Callable + +import numpy as np +from PySide6 import QtCore, QtGui + +from jabs.core.utils.pose_util import gen_line_fragments +from jabs.pose_estimation import PoseEstimation +from jabs.ui.colors import KEYPOINT_COLOR_MAP + +# Color for the pose line segments (white, semi-transparent). +LINE_SEGMENT_COLOR = QtGui.QColor(255, 255, 255, 128) + +# Base keypoint circle radius at a 1:1 (zoom 1.0) display scale. +KEYPOINT_SIZE = 3 + +# Alpha applied to keypoints/line segments of non-active identities. +_INACTIVE_KEYPOINT_ALPHA = 96 + + +def native_pose_sizes(width: int, height: int) -> tuple[int, int]: + """Return ``(keypoint_size, line_width)`` for drawing pose at native video resolution. + + The marker and line sizes scale with the frame's larger dimension so the overlay + stays legible across videos of different resolutions. + + Args: + width: Frame width in pixels. + height: Frame height in pixels. + + Returns: + Tuple of ``(keypoint_size, line_width)`` in pixels. + """ + reference = max(width, height, 1) + keypoint_size = max(3, round(reference / 250)) + line_width = max(2, round(reference / 400)) + return keypoint_size, line_width + + +def draw_identity_pose( + painter: QtGui.QPainter, + pose: PoseEstimation, + frame_index: int, + identity: int, + *, + to_output: Callable[[float, float], tuple[int, int] | None], + keypoint_size: int, + line_width: int, + active: bool, +) -> None: + """Draw one identity's pose skeleton and keypoints with ``painter``. + + Args: + painter: The painter to draw with. + pose: Pose estimation data for the video. + frame_index: Frame to draw the pose for. + identity: Identity whose pose is drawn. + to_output: Maps an image-space ``(x, y)`` to the painter's coordinate space, or + returns ``None`` to skip a point (e.g. a point outside a display crop). + keypoint_size: Radius (pixels) of each keypoint circle. + line_width: Width (pixels) of the skeleton line segments. + active: Whether this is the active identity. Inactive identities are drawn more + translucent. + """ + points, mask = pose.get_points(frame_index, identity) + if points is None or mask is None: + return + + line_color = QtGui.QColor(LINE_SEGMENT_COLOR) + if not active: + line_color.setAlpha(line_color.alpha() // 3) + pen = QtGui.QPen(line_color) + pen.setWidth(line_width) + painter.setPen(pen) + + for seg in gen_line_fragments(pose.get_connected_segments(), np.flatnonzero(mask == 0)): + segment_points = [to_output(p[0], p[1]) for p in points[seg]] + segment_points = [pt for pt in segment_points if pt is not None] + if len(segment_points) >= 2: + for i in range(len(segment_points) - 1): + painter.drawLine( + segment_points[i][0], + segment_points[i][1], + segment_points[i + 1][0], + segment_points[i + 1][1], + ) + + painter.setPen(QtCore.Qt.PenStyle.NoPen) + for keypoint in PoseEstimation.KeypointIndex: + point_index = keypoint.value + if not mask[point_index]: + continue + out = to_output(points[point_index][0], points[point_index][1]) + if out is None: + continue + color = KEYPOINT_COLOR_MAP[keypoint] + if not active: + color = QtGui.QColor(color) + color.setAlpha(_INACTIVE_KEYPOINT_ALPHA) + painter.setBrush(color) + painter.drawEllipse(QtCore.QPoint(out[0], out[1]), keypoint_size, keypoint_size) diff --git a/tests/ui/test_menu_handlers.py b/tests/ui/test_menu_handlers.py index 27d6bfdf..2f036a6b 100644 --- a/tests/ui/test_menu_handlers.py +++ b/tests/ui/test_menu_handlers.py @@ -8,7 +8,11 @@ from PySide6.QtWidgets import QApplication import jabs.ui.main_window.menu_handlers as menu_handlers_module - from jabs.ui.main_window.menu_handlers import _SETTINGS_EXPORT_FRAME_DIR, MenuHandlers + from jabs.ui.main_window.menu_handlers import ( + _SETTINGS_EXPORT_FRAME_DIR, + _SETTINGS_EXPORT_OVERLAY, + MenuHandlers, + ) SKIP_UI_TESTS = False SKIP_REASON = None @@ -36,6 +40,7 @@ def handler_setup(): """Create a MenuHandlers instance with a lightweight fake window.""" player = SimpleNamespace( get_raw_frame=MagicMock(), + get_overlay_frame=MagicMock(), current_frame=42, current_video_path=Path("video.mp4"), ) @@ -48,6 +53,35 @@ def handler_setup(): return MenuHandlers(window), window, player +def _patch_export_dialog( + monkeypatch, + *, + accepted: bool = True, + selected: tuple[str, ...] = ("/tmp/video_frame000042.png",), + overlay_checked: bool = False, +): + """Replace the QFileDialog and QCheckBox used by export_frame with controllable mocks. + + Returns the mock dialog instance so tests can assert on dialog interactions + (e.g. ``selectFile`` / ``setDirectory``). + """ + fake_qfiledialog = MagicMock() + dialog = fake_qfiledialog.return_value + dialog.exec.return_value = ( + fake_qfiledialog.DialogCode.Accepted if accepted else fake_qfiledialog.DialogCode.Rejected + ) + dialog.selectedFiles.return_value = list(selected) + dialog.layout.return_value = None # not a QGridLayout -> checkbox is not inserted + monkeypatch.setattr(menu_handlers_module.QtWidgets, "QFileDialog", fake_qfiledialog) + + fake_checkbox = MagicMock() + fake_checkbox.isChecked.return_value = overlay_checked + monkeypatch.setattr( + menu_handlers_module.QtWidgets, "QCheckBox", MagicMock(return_value=fake_checkbox) + ) + return dialog + + def test_export_frame_success(monkeypatch, handler_setup): """Export uses the suggested filename, persists the directory, and reports success.""" handler, window, player = handler_setup @@ -55,27 +89,15 @@ def test_export_frame_success(monkeypatch, handler_setup): pixmap.save.return_value = True player.get_raw_frame.return_value = pixmap - captured = {} - - def fake_get_save_file_name(parent, title, initial_path, file_filter, options=None): - captured["title"] = title - captured["initial_path"] = initial_path - return "/tmp/video_frame000042.png", "PNG Images (*.png)" - - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - fake_get_save_file_name, - ) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) + dialog = _patch_export_dialog(monkeypatch, selected=("/tmp/video_frame000042.png",)) handler.export_frame() - assert captured["title"] == "Export Frame" - assert captured["initial_path"] == "video_frame000042.png" + dialog.selectFile.assert_called_once_with("video_frame000042.png") player.get_raw_frame.assert_called_once_with(42) pixmap.save.assert_called_once_with("/tmp/video_frame000042.png", "PNG") - window._settings.setValue.assert_called_once_with(_SETTINGS_EXPORT_FRAME_DIR, "/tmp") + player.get_overlay_frame.assert_not_called() + window._settings.setValue.assert_any_call(_SETTINGS_EXPORT_FRAME_DIR, "/tmp") window.display_status_message.assert_called_once_with( "Frame exported: /tmp/video_frame000042.png", 5000, @@ -88,12 +110,7 @@ def test_export_frame_cancelled_does_not_save(monkeypatch, handler_setup): pixmap = MagicMock() player.get_raw_frame.return_value = pixmap - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - lambda *args, **kwargs: ("", ""), - ) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) + _patch_export_dialog(monkeypatch, accepted=False) handler.export_frame() @@ -109,13 +126,8 @@ def test_export_frame_no_frame_uses_warning_message(monkeypatch, handler_setup): player.get_raw_frame.return_value = None warning = MagicMock() - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - lambda *args, **kwargs: ("/tmp/video_frame000042.png", "PNG Images (*.png)"), - ) + _patch_export_dialog(monkeypatch) monkeypatch.setattr(menu_handlers_module.MessageDialog, "warning", warning) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) handler.export_frame() @@ -130,12 +142,7 @@ def test_export_frame_appends_extension(monkeypatch, handler_setup): pixmap.save.return_value = True player.get_raw_frame.return_value = pixmap - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - lambda *args, **kwargs: ("/tmp/custom_name", "PNG Images (*.png)"), - ) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) + _patch_export_dialog(monkeypatch, selected=("/tmp/custom_name",)) handler.export_frame() @@ -150,13 +157,8 @@ def test_export_frame_write_failure_uses_error_message(monkeypatch, handler_setu player.get_raw_frame.return_value = pixmap error = MagicMock() - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - lambda *args, **kwargs: ("/tmp/video_frame000042.png", "PNG Images (*.png)"), - ) + _patch_export_dialog(monkeypatch, selected=("/tmp/video_frame000042.png",)) monkeypatch.setattr(menu_handlers_module.MessageDialog, "error", error) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) handler.export_frame() @@ -168,6 +170,32 @@ def test_export_frame_write_failure_uses_error_message(monkeypatch, handler_setu window.display_status_message.assert_not_called() +def test_export_frame_saves_overlay_copy(monkeypatch, handler_setup): + """With the checkbox enabled, a second -overlay.png copy is written and reported.""" + handler, window, player = handler_setup + pixmap = MagicMock() + pixmap.save.return_value = True + player.get_raw_frame.return_value = pixmap + overlay_pixmap = MagicMock() + overlay_pixmap.save.return_value = True + player.get_overlay_frame.return_value = overlay_pixmap + + _patch_export_dialog( + monkeypatch, selected=("/tmp/video_frame000042.png",), overlay_checked=True + ) + + handler.export_frame() + + pixmap.save.assert_called_once_with("/tmp/video_frame000042.png", "PNG") + player.get_overlay_frame.assert_called_once_with(42) + overlay_pixmap.save.assert_called_once_with("/tmp/video_frame000042-overlay.png", "PNG") + window._settings.setValue.assert_any_call(_SETTINGS_EXPORT_OVERLAY, True) + window.display_status_message.assert_called_once_with( + "Frame exported: /tmp/video_frame000042.png (+ overlay copy)", + 5000, + ) + + def test_export_frame_restores_existing_directory(monkeypatch, tmp_path, handler_setup): """The last successful export directory is reused when it still exists.""" handler, window, player = handler_setup @@ -176,22 +204,28 @@ def test_export_frame_restores_existing_directory(monkeypatch, tmp_path, handler player.get_raw_frame.return_value = pixmap window._settings.value.return_value = str(tmp_path) - captured = {} + dialog = _patch_export_dialog(monkeypatch, selected=(str(tmp_path / "chosen_frame.png"),)) + + handler.export_frame() - def fake_get_save_file_name(parent, title, initial_path, file_filter, options=None): - captured["initial_path"] = initial_path - return str(tmp_path / "chosen_frame.png"), "PNG Images (*.png)" + dialog.setDirectory.assert_called_once_with(str(tmp_path)) + dialog.selectFile.assert_called_once_with("video_frame000042.png") - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - fake_get_save_file_name, - ) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) + +def test_export_frame_missing_directory_falls_back(monkeypatch, handler_setup): + """A missing saved directory is not applied to the dialog.""" + handler, window, player = handler_setup + pixmap = MagicMock() + pixmap.save.return_value = True + player.get_raw_frame.return_value = pixmap + window._settings.value.return_value = "/path/that/does/not/exist" + + dialog = _patch_export_dialog(monkeypatch, selected=("/tmp/video_frame000042.png",)) handler.export_frame() - assert captured["initial_path"] == str(tmp_path / "video_frame000042.png") + dialog.setDirectory.assert_not_called() + dialog.selectFile.assert_called_once_with("video_frame000042.png") def test_handle_select_all_delegates_to_central_widget(): @@ -210,29 +244,3 @@ def test_handle_select_current_bout_delegates_to_central_widget(): handler = MenuHandlers(window) handler.handle_select_current_bout() central.select_current_bout.assert_called_once_with() - - -def test_export_frame_missing_directory_falls_back(monkeypatch, handler_setup): - """A missing saved directory falls back to the bare suggested filename.""" - handler, window, player = handler_setup - pixmap = MagicMock() - pixmap.save.return_value = True - player.get_raw_frame.return_value = pixmap - window._settings.value.return_value = "/path/that/does/not/exist" - - captured = {} - - def fake_get_save_file_name(parent, title, initial_path, file_filter, options=None): - captured["initial_path"] = initial_path - return "/tmp/video_frame000042.png", "PNG Images (*.png)" - - monkeypatch.setattr( - menu_handlers_module.QtWidgets.QFileDialog, - "getSaveFileName", - fake_get_save_file_name, - ) - monkeypatch.setattr(menu_handlers_module, "USE_NATIVE_FILE_DIALOG", True) - - handler.export_frame() - - assert captured["initial_path"] == "video_frame000042.png" diff --git a/tests/ui/test_pose_drawing.py b/tests/ui/test_pose_drawing.py new file mode 100644 index 00000000..150149ae --- /dev/null +++ b/tests/ui/test_pose_drawing.py @@ -0,0 +1,125 @@ +"""Tests for the shared pose-drawing helper used by the overlay and frame export.""" + +from unittest.mock import MagicMock + +import numpy as np +import pytest + +try: + from PySide6.QtWidgets import QApplication # noqa: F401 + + from jabs.pose_estimation import PoseEstimation + from jabs.ui.player_widget.pose_drawing import draw_identity_pose, native_pose_sizes + + SKIP_UI_TESTS = False + SKIP_REASON = "" +except ImportError as e: + SKIP_UI_TESTS = True + SKIP_REASON = f"Qt/UI dependencies not available: {e}" + +pytestmark = pytest.mark.skipif(SKIP_UI_TESTS, reason=SKIP_REASON) + + +class _StubPose: + """Minimal stand-in for PoseEstimation exposing what draw_identity_pose needs.""" + + def __init__(self, num_keypoints: int, *, points_present: bool) -> None: + self._n = num_keypoints + self._present = points_present + + def get_points(self, frame_index: int, identity: int): + """Return zeroed points with an all-present (or no) mask.""" + if not self._present: + return None, None + points = np.zeros((self._n, 2), dtype=np.float32) + mask = np.ones(self._n, dtype=np.uint8) + return points, mask + + def get_connected_segments(self): + """No skeleton segments, so only keypoint circles are drawn.""" + return [] + + +def _identity_transform(x, y): + return int(x), int(y) + + +@pytest.mark.parametrize( + ("width", "height", "expected"), + [ + (100, 100, (3, 2)), # min bounds + (800, 600, (3, 2)), + (1920, 1080, (8, 5)), + ], + ids=["min", "sd", "hd"], +) +def test_native_pose_sizes(width, height, expected) -> None: + """Keypoint/line sizes honor their minimums and scale with resolution.""" + assert native_pose_sizes(width, height) == expected + + +def test_native_pose_sizes_monotonic() -> None: + """Larger frames never produce smaller markers.""" + small_kp, small_lw = native_pose_sizes(640, 480) + large_kp, large_lw = native_pose_sizes(3840, 2160) + assert large_kp > small_kp + assert large_lw > small_lw + + +def test_draw_identity_pose_noop_when_points_missing() -> None: + """Nothing is drawn when the identity has no pose at this frame.""" + painter = MagicMock() + pose = _StubPose(len(PoseEstimation.KeypointIndex), points_present=False) + + draw_identity_pose( + painter, + pose, + frame_index=0, + identity=0, + to_output=_identity_transform, + keypoint_size=4, + line_width=2, + active=True, + ) + + painter.drawEllipse.assert_not_called() + painter.drawLine.assert_not_called() + + +def test_draw_identity_pose_draws_each_visible_keypoint() -> None: + """One keypoint circle is drawn per present, in-bounds keypoint.""" + painter = MagicMock() + num_keypoints = len(PoseEstimation.KeypointIndex) + pose = _StubPose(num_keypoints, points_present=True) + + draw_identity_pose( + painter, + pose, + frame_index=0, + identity=0, + to_output=_identity_transform, + keypoint_size=4, + line_width=2, + active=True, + ) + + assert painter.drawEllipse.call_count == num_keypoints + + +def test_draw_identity_pose_skips_points_filtered_by_transform() -> None: + """Keypoints whose transform returns None (e.g. outside a crop) are not drawn.""" + painter = MagicMock() + pose = _StubPose(len(PoseEstimation.KeypointIndex), points_present=True) + + draw_identity_pose( + painter, + pose, + frame_index=0, + identity=0, + to_output=lambda x, y: None, + keypoint_size=4, + line_width=2, + active=True, + ) + + painter.drawEllipse.assert_not_called() From 93c8975b4b5e01eb55ae8085ac292c878f659627 Mon Sep 17 00:00:00 2001 From: Glen Beane <356266+gbeane@users.noreply.github.com> Date: Thu, 18 Jun 2026 09:35:32 -0400 Subject: [PATCH 2/2] Use pose.identities for overlay export segmentation loop --- src/jabs/ui/player_widget/player_widget.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/jabs/ui/player_widget/player_widget.py b/src/jabs/ui/player_widget/player_widget.py index 8487c642..9689b57e 100644 --- a/src/jabs/ui/player_widget/player_widget.py +++ b/src/jabs/ui/player_widget/player_widget.py @@ -451,14 +451,14 @@ def get_overlay_frame(self, frame_number: int | None = None) -> QtGui.QPixmap | # Bake segmentation contours into the BGR frame (a no-op when the pose file has # no segmentation data for this identity/frame). if isinstance(self._pose_est, PoseEstimationV6): - for ident in range(self._pose_est.num_identities): + for identity in self._pose_est.identities: # active=True for every identity: the export has no active identity, so # all contours are drawn in the same (active) color rather than singling # one out. overlay_segmentation( img, self._pose_est, - identity=ident, + identity=identity, frame_index=target_frame, active=True, )