Merge pull request #179 from KumarLabJax/session-tracker

initial implementation of jabs labeling session tracking

Author: gbeane

pyproject.toml

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

src/jabs/classifier/classifier.py

@@ -622,32 +622,38 @@ def count_label_threshold(all_counts: dict):
 
             all_counts is a dict with the following form
             {
-                '<video name>': [
-                    (
-                        <identity>,
-                        (behavior frame count  - fragmented, not behavior frame count  - fragmented),
-                        (behavior bout count  - fragmented, not behavior bout count  - fragmented),
-                        (behavior frame count  - unfragmented, not behavior frame count  - unfragmented),
-                        (behavior bout count  - unfragmented, not behavior bout count  - unfragmented),
-                    ),
-                ]
+                '<video name>': {
+                    <identity>: {
+                        "fragmented_frame_counts": (
+                            behavior frame count: fragmented,
+                            not behavior frame count: fragmented),
+                        "fragmented_bout_counts": (
+                            behavior bout count: fragmented,
+                            not behavior bout count: fragmented
+                        ),
+                        "unfragmented_frame_counts": (
+                            behavior frame count: unfragmented,
+                            not behavior frame count: unfragmented
+                        ),
+                        "unfragmented_bout_counts": (
+                            behavior bout count: unfragmented,
+                            not behavior bout count: unfragmented
+                        ),
+                    },
+                }
             }
 
         Returns:
             number of groups that meet label criteria
 
         Note: uses "fragmented" label counts, since these reflect the counts of labels that are usable for training
-        fragmented counts are:
-
-           count[1][0] - behavior frame count
-           count[1][1] - not behavior frame count
         """
         group_count = 0
-        for _, counts in all_counts.items():
-            for count in counts:
+        for video in all_counts:
+            for identity_count in all_counts[video].values():
                 if (
-                    count[1][0] >= Classifier.LABEL_THRESHOLD
-                    and count[1][1] >= Classifier.LABEL_THRESHOLD
+                    identity_count["fragmented_frame_counts"][0] >= Classifier.LABEL_THRESHOLD
+                    and identity_count["fragmented_frame_counts"][1] >= Classifier.LABEL_THRESHOLD
                 ):
                     group_count += 1
         return group_count

src/jabs/project/project.py

@@ -20,6 +20,7 @@
 from .prediction_manager import PredictionManager
 from .project_paths import ProjectPaths
 from .project_utils import to_safe_name
+from .session_tracker import SessionTracker
 from .settings_manager import SettingsManager
 from .track_labels import TrackLabels
 from .video_labels import VideoLabels
@@ -52,7 +53,9 @@ class Project:
         project_paths: ProjectPaths instance for this project.
     """
 
-    def __init__(self, project_path, use_cache=True, enable_video_check=True):
+    def __init__(
+        self, project_path, use_cache=True, enable_video_check=True, enable_session_tracker=True
+    ):
         self._paths = ProjectPaths(Path(project_path), use_cache=use_cache)
         self._paths.create_directories()
         self._total_project_identities = 0
@@ -62,10 +65,16 @@ def __init__(self, project_path, use_cache=True, enable_video_check=True):
         self._video_manager = VideoManager(self._paths, self._settings_manager, enable_video_check)
         self._feature_manager = FeatureManager(self._paths, self._video_manager.videos)
         self._prediction_manager = PredictionManager(self)
+        self._session_tracker = SessionTracker(self, tracking_enabled=enable_session_tracker)
 
         # write out the defaults to the project file
         self._settings_manager.save_project_file({"defaults": self.get_project_defaults()})
 
+        # Start a session tracker for this project.
+        # Since the session has a reference to the Project, the Project should be fully initialized before starting
+        # the session tracker.
+        self._session_tracker.start_session()
+
     def _validate_pose_files(self):
         """Ensure all videos have corresponding pose files."""
         err = False
@@ -143,6 +152,11 @@ def labeler(self) -> str:
         """
         return getpass.getuser()
 
+    @property
+    def session_tracker(self) -> SessionTracker | None:
+        """get the session tracker for this project"""
+        return self._session_tracker
+
     def load_pose_est(self, video_path: Path) -> PoseEstimation:
         """return a PoseEstimation object for a given video path
 
@@ -391,7 +405,7 @@ def counts(self, behavior):
         """
         counts = {}
         for video in self._video_manager.videos:
-            counts[video] = self.read_counts(video, behavior)
+            counts[video] = self.load_counts(video, behavior)
         return counts
 
     def get_labeled_features(
@@ -544,19 +558,32 @@ def __has_pose(self, vid: str):
             return False
         return True
 
-    def read_counts(self, video, behavior) -> list[tuple]:
-        """read labeled frame and bout counts from json file
+    def load_counts(self, video, behavior) -> dict[str, tuple[int, int]]:
+        """load labeled frame and bout counts from json file
 
         Returns:
-            list of labeled frame and bout counts for each identity for
-            the specified behavior. Each element in the list is a tuple of the form
-            (
-                identity,
-                (fragmented behavior frame count, fragmented not behavior frame count),
-                (fragmented behavior bout count, fragmented not behavior bout count),
-                (unfragmented behavior frame count, unfragmented not behavior frame count)
-                (unfragmented behavior bout count, unfragmented not behavior bout count)
-            )
+            dict of labeled frame and bout counts for each identity for
+            the specified behavior.
+            {
+                identity: {
+                    "fragmented_frame_counts": (
+                        fragmented behavior frame count,
+                        fragmented not behavior frame count
+                    ),
+                    "fragmented_bout_counts": (
+                        fragmented behavior bout count,
+                        fragmented not behavior bout count
+                    ),
+                    "unfragmented_frame_counts": (
+                        unfragmented behavior frame count,
+                        unfragmented not behavior frame count
+                    )
+                    "unfragmented_bout_counts": (
+                        unfragmented behavior bout count,
+                        unfragmented not behavior bout count
+                    )
+                }
+            }
 
         Note: "unfragmented" counts labels where identity drops out. "fragmented" does not,
             so if an identity drops out during a bout, the bout will be split in the fragmented counts but will
@@ -585,7 +612,7 @@ def count_labels(
 
         video_filename = Path(video).name
         path = self._paths.annotations_dir / Path(video_filename).with_suffix(".json")
-        counts = []
+        counts = {}
 
         if path.exists():
             with path.open() as f:
@@ -605,13 +632,13 @@ def count_labels(
                         # unless the user creates some new labels over frames without identity
                         unfragmented_counts = fragmented_counts
 
-                    counts.append(
-                        (
-                            identity,
-                            fragmented_counts[0],
-                            fragmented_counts[1],
-                            unfragmented_counts[0],
-                            unfragmented_counts[1],
-                        )
-                    )
+                    # identity is stored as a string in the JSON file because it's used as a key. Turn it back
+                    # into an int as used internally by JABS
+                    counts[int(identity)] = {
+                        "fragmented_frame_counts": fragmented_counts[0],
+                        "fragmented_bout_counts": fragmented_counts[1],
+                        "unfragmented_frame_counts": unfragmented_counts[0],
+                        "unfragmented_bout_counts": unfragmented_counts[1],
+                    }
+
         return counts

src/jabs/project/project_paths.py

@@ -16,6 +16,7 @@ def __init__(self, base_path: Path, use_cache: bool = True):
         self._prediction_dir = self._jabs_dir / "predictions"
         self._classifier_dir = self._jabs_dir / "classifiers"
         self._archive_dir = self._jabs_dir / "archive"
+        self._session_dir = self._jabs_dir / "session"
         self._cache_dir = self._jabs_dir / "cache" if use_cache else None
 
         self._project_file = self._jabs_dir / self.__PROJECT_FILE
@@ -65,13 +66,19 @@ def cache_dir(self) -> Path | None:
         """Get the path to the cache directory."""
         return self._cache_dir
 
+    @property
+    def session_dir(self) -> Path:
+        """Get the path to the session directory."""
+        return self._session_dir
+
     def create_directories(self):
         """Create all necessary directories for the project."""
         self._annotations_dir.mkdir(parents=True, exist_ok=True)
         self._feature_dir.mkdir(parents=True, exist_ok=True)
         self._prediction_dir.mkdir(parents=True, exist_ok=True)
         self._classifier_dir.mkdir(parents=True, exist_ok=True)
         self._archive_dir.mkdir(parents=True, exist_ok=True)
+        self._session_dir.mkdir(parents=True, exist_ok=True)
 
         if self._cache_dir:
             self._cache_dir.mkdir(parents=True, exist_ok=True)

src/jabs/project/project_pruning.py

@@ -22,9 +22,13 @@ def get_videos_to_prune(project: Project, behavior: str | None = None) -> list[V
         behavior (str | None): The behavior to check for labels. If None, checks all behaviors.
     """
 
-    def check_label_counts(label_counts: list[tuple[str, tuple[int, int]]]) -> bool:
-        """Check if there are any labels for the given counts."""
-        return any(count[1][0] > 0 or count[1][1] > 0 for count in label_counts)
+    def check_label_counts(label_counts: dict[str, dict[str, tuple[int, int]]]) -> bool:
+        """Return True if any count in label_counts is greater than zero."""
+        for identity_counts in label_counts.values():
+            for counts in identity_counts.values():
+                if any(count > 0 for count in counts):
+                    return True
+        return False
 
     videos_to_remove = []
     for video in project.video_manager.videos:
@@ -34,11 +38,11 @@ def check_label_counts(label_counts: list[tuple[str, tuple[int, int]]]) -> bool:
 
         has_labels = False
         if behavior:
-            counts = project.read_counts(video, behavior)
+            counts = project.load_counts(video, behavior)
             has_labels = check_label_counts(counts)
         else:
             for b in project.settings_manager.behavior_names:
-                counts = project.read_counts(video, b)
+                counts = project.load_counts(video, b)
                 has_labels = check_label_counts(counts)
 
                 # found labels for at least one behavior, so we can stop checking