Merge branch 'main' into divya/slumbr_base

Author: gitttt-1234

.claude/commands/coverage.md

@@ -2,13 +2,17 @@ Run tests with coverage.
 
 Command to run:
 ```
-pytest -q --maxfail=1 --cov=sleap_nn --cov-branch tests/ && rm -f .coverage.* && python -m coverage annotate
+uv run pytest -q --maxfail=1 --cov --cov-branch && rm .coverage.* && uv run coverage annotate
 ```
 
-The result will be the terminal output and the line-by-line coverage will be in files sitting next to each module with the file naming `{module_name.py},cover`. 
+This generates a coverage annotation file next to each module with the name `{module_name.py},cover`, as well as a simple summary.
 
-If you are working on a PR, figure out which files were changed and look for coverage specifically in those. If you don't know which files to look for coverage in, use this:
+To get the final actionable summary, run this script:
 
 ```
-git diff --name-only $(git merge-base origin/main HEAD) | jq -R . | jq -s .
-```
+uv run python scripts/cov_summary.py --only-pr-diff-lines
+```
+
+This will output one module per line with line number ranges for missing coverage. Importantly, it will filter it by diffs in the PR.
+
+Use this summary together with the corresponding `,cover` file to describe each miss to inform subsequent test development.

.claude/commands/pr-description.md

@@ -2,6 +2,14 @@ Update PR description.
 
 Use the `gh` CLI to fetch the current PR description, then update it with a comprehensive description of the changes made in this PR.
 
+
+Command to fetch PR info:
+
+```
+gh pr view PR_NUMBER --json
+      number,title,body,url,state,closingIssuesReferences
+```
+
 If there is an associated issue (linked in the PR metadata or mentioned in the PR description), then use the `gh` CLI to fetch that too to contextualize the work done in the PR.
 
 Include a summary, example usage (for enhancements), API changes, and other notes for future consideration (including reasoning behind design decisions).

.gitignore

@@ -173,4 +173,7 @@ wandb/
 .serena/
 
 # macOS
-.DS_Store
+.DS_Store
+
+# Development scratch folder
+scratch/

pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-    "sleap-io>=0.2.0",
+    "sleap-io>=0.5.7",
     "numpy",
     "lightning",
     "kornia",
@@ -54,19 +54,19 @@ readme = {file = ["README.md"], content-type="text/markdown"}
 [project.optional-dependencies]
 torch = [
     "torch",
-    "torchvision<0.24.0",
+    "torchvision>=0.20.0,<0.24.0",
 ]
 torch-cpu = [
     "torch",
-    "torchvision<0.24.0",
+    "torchvision>=0.20.0,<0.24.0",
 ]
 torch-cuda118 = [
     "torch",
-    "torchvision<0.24.0",
+    "torchvision>=0.20.0,<0.24.0",
 ]
 torch-cuda128 = [
     "torch",
-    "torchvision<0.24.0",
+    "torchvision>=0.20.0,<0.24.0",
 ]
 dev = [
     "pytest",

sleap_nn/__init__.py

@@ -48,4 +48,4 @@ def _safe_print(msg):
     format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {name}:{function}:{line} | {message}",
 )
 
-__version__ = "0.0.3"
+__version__ = "0.0.4"

sleap_nn/data/custom_datasets.py

@@ -799,7 +799,9 @@ def __getitem__(self, index) -> Dict:
 
         instances = []
         for inst in instances_list:
-            instances.append(inst.numpy())
+            instances.append(
+                inst.numpy()
+            )  # no need to filter empty instances; handled while creating instance_idx_list
         instances = np.stack(instances, axis=0)
 
         # Add singleton time dimension for single frames.
@@ -1045,8 +1047,9 @@ def __getitem__(self, index) -> Dict:
 
         instances = []
         for inst in instances_list:
-            if not inst.is_empty:
-                instances.append(inst.numpy())
+            instances.append(
+                inst.numpy()
+            )  # no need to filter empty instance (handled while creating instance_idx)
         instances = np.stack(instances, axis=0)
 
         # Add singleton time dimension for single frames.

sleap_nn/data/instance_cropping.py

@@ -44,15 +44,16 @@ def find_instance_crop_size(
     max_length = 0.0
     for lf in labels:
         for inst in lf.instances:
-            pts = inst.numpy()
-            pts *= input_scaling
-            diff_x = np.nanmax(pts[:, 0]) - np.nanmin(pts[:, 0])
-            diff_x = 0 if np.isnan(diff_x) else diff_x
-            max_length = np.maximum(max_length, diff_x)
-            diff_y = np.nanmax(pts[:, 1]) - np.nanmin(pts[:, 1])
-            diff_y = 0 if np.isnan(diff_y) else diff_y
-            max_length = np.maximum(max_length, diff_y)
-            max_length = np.maximum(max_length, min_crop_size_no_pad)
+            if not inst.is_empty:  # only if at least one point is not nan
+                pts = inst.numpy()
+                pts *= input_scaling
+                diff_x = np.nanmax(pts[:, 0]) - np.nanmin(pts[:, 0])
+                diff_x = 0 if np.isnan(diff_x) else diff_x
+                max_length = np.maximum(max_length, diff_x)
+                diff_y = np.nanmax(pts[:, 1]) - np.nanmin(pts[:, 1])
+                diff_y = 0 if np.isnan(diff_y) else diff_y
+                max_length = np.maximum(max_length, diff_y)
+                max_length = np.maximum(max_length, min_crop_size_no_pad)
 
     max_length += float(padding)
     crop_size = math.ceil(max_length / float(maximum_stride)) * maximum_stride

sleap_nn/tracking/tracker.py

@@ -9,6 +9,14 @@
 from datetime import datetime
 from loguru import logger
 import functools
+import rich
+from rich.progress import (
+    Progress,
+    BarColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+    MofNCompleteColumn,
+)
 
 import sleap_io as sio
 from sleap_nn.evaluation import compute_oks
@@ -716,6 +724,17 @@ def connect_single_breaks(
     return lfs
 
 
+class RateColumn(rich.progress.ProgressColumn):
+    """Renders the progress rate."""
+
+    def render(self, task: "Task") -> rich.progress.Text:
+        """Show progress rate."""
+        speed = task.speed
+        if speed is None:
+            return rich.progress.Text("?", style="progress.data.speed")
+        return rich.progress.Text(f"{speed:.1f} frames/s", style="progress.data.speed")
+
+
 def run_tracker(
     untracked_frames: List[sio.LabeledFrame],
     window_size: int = 5,
@@ -810,35 +829,66 @@ def run_tracker(
         tracking_pre_cull_to_target=tracking_pre_cull_to_target,
         tracking_pre_cull_iou_threshold=tracking_pre_cull_iou_threshold,
     )
-    tracked_lfs = []
-    for lf in untracked_frames:
-        # prefer user instances over predicted instance
-        instances = []
-        if lf.has_user_instances:
-            instances_to_track = lf.user_instances
-            if lf.has_predicted_instances:
-                instances = lf.predicted_instances
-        else:
-            instances_to_track = lf.predicted_instances
 
-        instances.extend(
-            tracker.track(
-                untracked_instances=instances_to_track,
-                frame_idx=lf.frame_idx,
-                image=lf.image,
-            )
-        )
-        tracked_lfs.append(
-            sio.LabeledFrame(
-                video=lf.video, frame_idx=lf.frame_idx, instances=instances
-            )
-        )
+    try:
+        with Progress(
+            "{task.description}",
+            BarColumn(),
+            "[progress.percentage]{task.percentage:>3.0f}%",
+            MofNCompleteColumn(),
+            "ETA:",
+            TimeRemainingColumn(),
+            "Elapsed:",
+            TimeElapsedColumn(),
+            RateColumn(),
+            auto_refresh=False,
+            refresh_per_second=4,
+            speed_estimate_period=5,
+        ) as progress:
+            task = progress.add_task("Tracking...", total=len(untracked_frames))
+            last_report = time()
+
+            tracked_lfs = []
+            for lf in untracked_frames:
+                # prefer user instances over predicted instance
+                instances = []
+                if lf.has_user_instances:
+                    instances_to_track = lf.user_instances
+                    if lf.has_predicted_instances:
+                        instances = lf.predicted_instances
+                else:
+                    instances_to_track = lf.predicted_instances
+
+                instances.extend(
+                    tracker.track(
+                        untracked_instances=instances_to_track,
+                        frame_idx=lf.frame_idx,
+                        image=lf.image,
+                    )
+                )
+                tracked_lfs.append(
+                    sio.LabeledFrame(
+                        video=lf.video, frame_idx=lf.frame_idx, instances=instances
+                    )
+                )
+
+                progress.update(task, advance=1)
+
+                if time() - last_report > 0.25:
+                    progress.refresh()
+                    last_report = time()
+
+    except KeyboardInterrupt:
+        logger.info("Tracking interrupted by user")
+        raise KeyboardInterrupt
 
     if tracking_clean_instance_count > 0:
+        logger.info("Post-processing: Culling instances...")
         tracked_lfs = cull_instances(
             tracked_lfs, tracking_clean_instance_count, tracking_clean_iou_threshold
         )
         if not post_connect_single_breaks:
+            logger.info("Post-processing: Connecting single breaks...")
             tracked_lfs = connect_single_breaks(
                 tracked_lfs, tracking_clean_instance_count
             )