Add optional pose/segmentation overlay copy to frame export

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 `<Behavior_Name>_training_<YYYYMMDD_hhmmss>.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.

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 `<Behavior_Name>_training_<YYYYMMDD_hhmmss>.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.

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."""

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),
+            )

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 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=identity,
+                    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: