Merge pull request #199 from KumarLabJax/add-identity-to-stacked-timeline-rows

add the identity as a label to each section of the stacked timeline

Author: gbeane

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jabs-behavior-classifier"
-version = "0.34.2"
+version = "0.34.3"
 license = "Proprietary"
 repository = "https://github.com/KumarLabJax/JABS-behavior-classifier"
 description = ""

src/jabs/project/session_tracker.py

@@ -332,5 +332,9 @@ def resume_session(self):
     def _flush_session(self):
         """Flush the current session data to the session file."""
         if self._session_file and self._session:
-            with open(self._session_file, "w") as f:
-                json.dump(self._session, f, indent=4)
+            try:
+                with open(self._session_file, "w") as f:
+                    json.dump(self._session, f, indent=4)
+            except NameError:
+                # Python is shutting down, and open is no longer available
+                pass

src/jabs/ui/central_widget.py

@@ -290,8 +290,7 @@ def load_video(self, path: Path) -> None:
             # open poses and any labels that might exist for this video
             self._pose_est = self._project.load_pose_est(path)
             self._labels = self._project.video_manager.load_video_labels(path)
-            self._stacked_timeline.num_identities = self._pose_est.num_identities
-            self._stacked_timeline.num_frames = self._pose_est.num_frames
+            self._stacked_timeline.pose = self._pose_est
 
             # if no saved labels exist, initialize a new VideoLabels object
             if self._labels is None:

src/jabs/ui/stacked_timeline_widget/stacked_timeline_widget.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 from PySide6.QtCore import Slot
-from PySide6.QtWidgets import QApplication, QFrame, QSizePolicy, QVBoxLayout, QWidget
+from PySide6.QtWidgets import QApplication, QFrame, QLabel, QSizePolicy, QVBoxLayout, QWidget
 
 from jabs.behavior_search import (
     BehaviorSearchQuery,
@@ -11,6 +11,7 @@
     SearchHit,
     TimelineAnnotationSearchQuery,
 )
+from jabs.pose_estimation import PoseEstimation
 from jabs.project import TrackLabels
 
 from .frame_labels_widget import FrameLabelsWidget
@@ -63,6 +64,7 @@ def __init__(self, *args, **kwargs) -> None:
         self._prediction_overview_widgets = []
         self._identity_frames = []
         self._frame_labels = FrameLabelsWidget(self)
+        self._pose: PoseEstimation | None = None
 
         self._layout = QVBoxLayout(self)
         self._layout.setContentsMargins(0, 0, 0, 0)
@@ -87,18 +89,6 @@ def _prediction_overview_widget_factory(self, parent) -> PredictionOverviewWidge
         widget.framerate = self.framerate
         return widget
 
-    @property
-    def num_identities(self) -> int:
-        """Get the number of identities."""
-        return self._num_identities
-
-    @num_identities.setter
-    def num_identities(self, value: int) -> None:
-        """Set the number of identities and reset the layout."""
-        if value != self._num_identities:
-            self._num_identities = value
-            self._reset_layout()
-
     @property
     def view_mode(self) -> ViewMode:
         """Get the current view mode."""
@@ -123,6 +113,23 @@ def identity_mode(self, value: IdentityMode) -> None:
             self._identity_mode = value
             self._update_widget_visibility()
 
+    @property
+    def pose(self) -> PoseEstimation | None:
+        """Get the PoseEstimation object used by the label overview widgets."""
+        return self._pose
+
+    @pose.setter
+    def pose(self, pose_est: PoseEstimation) -> None:
+        """Set the PoseEstimation object used by the label overview widgets.
+
+        Args:
+            pose_est: PoseEstimation object to set.
+        """
+        self._pose = pose_est
+        self._num_identities = pose_est.num_identities
+        self._num_frames = pose_est.num_frames
+        self._reset_layout()
+
     def _reset_layout(self) -> None:
         """Recreate the layout and child widgets for all identities.
 
@@ -143,7 +150,12 @@ def _reset_layout(self) -> None:
         self._identity_frames = []
 
         # Create new frames and widgets
-        for _ in range(self._num_identities):
+        for identity_index in range(self._num_identities):
+            if self._pose and self._pose.external_identities is not None:
+                identity_display_name = self._pose.external_identities[identity_index]
+            else:
+                identity_display_name = identity_index
+
             frame = QFrame(self)
             frame.setFrameShape(QFrame.Shape.NoFrame)
             frame.setStyleSheet("QFrame {border: none; padding: 2px;}")
@@ -156,6 +168,7 @@ def _reset_layout(self) -> None:
             label_widget.setVisible(False)
             prediction_widget.setVisible(False)
 
+            vbox.addWidget(QLabel(f"{identity_display_name}:"))
             vbox.addWidget(label_widget)
             vbox.addWidget(prediction_widget)
 
@@ -189,10 +202,12 @@ def _set_active_frame_border(self, active_index: int | None = None) -> None:
         for i, frame in enumerate(self._identity_frames):
             if i == active_index:
                 frame.setStyleSheet(
-                    f"QFrame {{border: 2px solid {accent_color}; border-radius: 8px;}}"
+                    f"QFrame > QWidget {{border: none;}} QFrame {{border: 2px solid {accent_color}; border-radius: 8px; padding: 2px;}}"
                 )
             else:
-                frame.setStyleSheet("QFrame {border: none; padding: 2px;}")
+                frame.setStyleSheet(
+                    "QFrame > QWidget {border: none;} QFrame {border: 2px solid transparent; border-radius: 8px; padding: 2px;}"
+                )
 
     @staticmethod
     def _get_accent_color() -> str:
@@ -205,12 +220,6 @@ def num_frames(self) -> int:
         """Get the number of frames."""
         return self._num_frames
 
-    @num_frames.setter
-    def num_frames(self, value: int) -> None:
-        """Set the number of frames."""
-        self._num_frames = value
-        self._frame_labels.set_num_frames(value)
-
     @property
     def framerate(self) -> int:
         """Get the framerate."""