sleap-nn 0.1.0a0__tar.gz → 0.1.0a2__tar.gz

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.1.0a0
+Version: 0.1.0a2
 Summary: Neural network backend for training and inference for animal pose estimation.
 Author-email: Divya Seshadri Murali <dimurali@salk.edu>, Elizabeth Berrigan <eberrigan@salk.edu>, Vincent Tu <vitu@ucsd.edu>, Liezl Maree <lmaree@salk.edu>, David Samy <davidasamy@gmail.com>, Talmo Pereira <talmo@salk.edu>
 License: BSD-3-Clause
@@ -13,7 +13,7 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: sleap-io<0.7.0,>=0.6.0
+Requires-Dist: sleap-io<0.7.0,>=0.6.2
 Requires-Dist: numpy
 Requires-Dist: lightning
 Requires-Dist: kornia

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/docs/config.md

@@ -845,6 +845,7 @@ trainer_config:
     - `wandb_mode`: (str) "offline" if only local logging is required. **Default**: `"None"`
     - `prv_runid`: (str) Previous run ID if training should be resumed from a previous ckpt. **Default**: `None`
     - `group`: (str) Group name for the run.
+    - `delete_local_logs`: (bool, optional) If `True`, delete local wandb logs folder (`wandb/`) after training completes. If `False`, keep the folder. If `None` (default), automatically delete if logging online (`wandb_mode` != "offline") and keep if logging offline. This can save significant disk space since wandb local logs can be several GB. **Default**: `None`
 
 **Example WandB configurations:**
 
@@ -876,6 +877,18 @@ trainer_config:
     group: "continued_experiments"
 ```
 
+**Keep local wandb logs (override auto-delete):**
+```yaml
+trainer_config:
+  use_wandb: true
+  wandb:
+    entity: "your_username"
+    project: "sleap_nn_experiments"
+    name: "training_run"
+    wandb_mode: "online"
+    delete_local_logs: false  # Keep local logs even when syncing online
+```
+
 **No WandB logging:**
 ```yaml
 trainer_config:

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-    "sleap-io>=0.6.0,<0.7.0",
+    "sleap-io>=0.6.2,<0.7.0",
     "numpy",
     "lightning",
     "kornia",

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/__init__.py

@@ -41,14 +41,16 @@ def _safe_print(msg):
 
 
 # Add logger with the custom filter
+# Disable colorization to avoid ANSI codes in captured output
 logger.add(
     _safe_print,
     level="DEBUG",
     filter=_should_log,
-    format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {name}:{function}:{line} | {message}",
+    format="{time:YYYY-MM-DD HH:mm:ss} | {message}",
+    colorize=False,
 )
 
-__version__ = "0.1.0a0"
+__version__ = "0.1.0a2"
 
 # Public API
 from sleap_nn.evaluation import load_metrics

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/config/get_config.py

@@ -677,6 +677,7 @@ def get_trainer_config(
     wandb_save_viz_imgs_wandb: bool = False,
     wandb_resume_prv_runid: Optional[str] = None,
     wandb_group_name: Optional[str] = None,
+    wandb_delete_local_logs: Optional[bool] = None,
     optimizer: str = "Adam",
     learning_rate: float = 1e-3,
     amsgrad: bool = False,
@@ -746,6 +747,9 @@ def get_trainer_config(
         wandb_resume_prv_runid: Previous run ID if training should be resumed from a previous
             ckpt. Default: None
         wandb_group_name: Group name for the wandb run. Default: None.
+        wandb_delete_local_logs: If True, delete local wandb logs folder after training.
+            If False, keep the folder. If None (default), automatically delete if logging
+            online (wandb_mode != "offline") and keep if logging offline. Default: None.
         optimizer: Optimizer to be used. One of ["Adam", "AdamW"]. Default: "Adam".
         learning_rate: Learning rate of type float. Default: 1e-3.
         amsgrad: Enable AMSGrad with the optimizer. Default: False.
@@ -846,6 +850,7 @@ def get_trainer_config(
             save_viz_imgs_wandb=wandb_save_viz_imgs_wandb,
             prv_runid=wandb_resume_prv_runid,
             group=wandb_group_name,
+            delete_local_logs=wandb_delete_local_logs,
         ),
         save_ckpt=save_ckpt,
         ckpt_dir=ckpt_dir,

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/config/trainer_config.py

@@ -90,6 +90,10 @@ class WandBConfig:
         viz_box_size: (float) Size of keypoint boxes in pixels (for viz_boxes). *Default*: `5.0`.
         viz_confmap_threshold: (float) Threshold for confidence map masks (for viz_masks). *Default*: `0.1`.
         log_viz_table: (bool) If True, also log images to a wandb.Table for backwards compatibility. *Default*: `False`.
+        delete_local_logs: (bool, optional) If True, delete local wandb logs folder after
+            training. If False, keep the folder. If None (default), automatically delete
+            if logging online (wandb_mode != "offline") and keep if logging offline.
+            *Default*: `None`.
     """
 
     entity: Optional[str] = None
@@ -107,6 +111,7 @@ class WandBConfig:
     viz_box_size: float = 5.0
     viz_confmap_threshold: float = 0.1
     log_viz_table: bool = False
+    delete_local_logs: Optional[bool] = None
 
 
 @define
@@ -203,6 +208,23 @@ class EarlyStoppingConfig:
     stop_training_on_plateau: bool = True
 
 
+@define
+class EvalConfig:
+    """Configuration for epoch-end evaluation.
+
+    Attributes:
+        enabled: (bool) Enable epoch-end evaluation metrics. *Default*: `False`.
+        frequency: (int) Evaluate every N epochs. *Default*: `1`.
+        oks_stddev: (float) OKS standard deviation for evaluation. *Default*: `0.025`.
+        oks_scale: (float) OKS scale override. If None, uses default. *Default*: `None`.
+    """
+
+    enabled: bool = False
+    frequency: int = field(default=1, validator=validators.ge(1))
+    oks_stddev: float = field(default=0.025, validator=validators.gt(0))
+    oks_scale: Optional[float] = None
+
+
 @define
 class HardKeypointMiningConfig:
     """Configuration for online hard keypoint mining.
@@ -305,6 +327,7 @@ class TrainerConfig:
         factory=HardKeypointMiningConfig
     )
     zmq: Optional[ZMQConfig] = field(factory=ZMQConfig)  # Required for SLEAP GUI
+    eval: EvalConfig = field(factory=EvalConfig)  # Epoch-end evaluation config
 
     @staticmethod
     def validate_optimizer_name(value):

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/data/custom_datasets.py

@@ -13,6 +13,14 @@ from omegaconf import DictConfig, OmegaConf
 import numpy as np
 from PIL import Image
 from loguru import logger
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    BarColumn,
+    TimeElapsedColumn,
+)
+from rich.console import Console
 import torch
 import torchvision.transforms as T
 from torch.utils.data import Dataset, DataLoader, DistributedSampler
@@ -215,17 +223,51 @@ class BaseDataset(Dataset):
     def _fill_cache(self, labels: List[sio.Labels]):
         """Load all samples to cache."""
         # TODO: Implement parallel processing (using threads might cause error with MediaVideo backend)
-        for sample in self.lf_idx_list:
-            labels_idx = sample["labels_idx"]
-            lf_idx = sample["lf_idx"]
-            img = labels[labels_idx][lf_idx].image
-            if img.shape[-1] == 1:
-                img = np.squeeze(img)
-            if self.cache_img == "disk":
-                f_name = f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
-                Image.fromarray(img).save(f_name, format="JPEG")
-            if self.cache_img == "memory":
-                self.cache[(labels_idx, lf_idx)] = img
+        import os
+        import sys
+
+        total_samples = len(self.lf_idx_list)
+        cache_type = "disk" if self.cache_img == "disk" else "memory"
+
+        # Check for NO_COLOR env var or non-interactive terminal
+        no_color = (
+            os.environ.get("NO_COLOR") is not None
+            or os.environ.get("FORCE_COLOR") == "0"
+        )
+        use_progress = sys.stdout.isatty() and not no_color
+
+        def process_samples(progress=None, task=None):
+            for sample in self.lf_idx_list:
+                labels_idx = sample["labels_idx"]
+                lf_idx = sample["lf_idx"]
+                img = labels[labels_idx][lf_idx].image
+                if img.shape[-1] == 1:
+                    img = np.squeeze(img)
+                if self.cache_img == "disk":
+                    f_name = f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
+                    Image.fromarray(img).save(f_name, format="JPEG")
+                if self.cache_img == "memory":
+                    self.cache[(labels_idx, lf_idx)] = img
+                if progress is not None:
+                    progress.update(task, advance=1)
+
+        if use_progress:
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                BarColumn(),
+                TextColumn("{task.completed}/{task.total}"),
+                TimeElapsedColumn(),
+                console=Console(force_terminal=True),
+                transient=True,
+            ) as progress:
+                task = progress.add_task(
+                    f"Caching images to {cache_type}", total=total_samples
+                )
+                process_samples(progress, task)
+        else:
+            logger.info(f"Caching {total_samples} images to {cache_type}...")
+            process_samples()
 
     def __len__(self) -> int:
         """Return the number of samples in the dataset."""

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/evaluation.py

@@ -29,11 +29,27 @@ def get_instances(labeled_frame: sio.LabeledFrame) -> List[MatchInstance]:
     """
     instance_list = []
     frame_idx = labeled_frame.frame_idx
-    video_path = (
-        labeled_frame.video.backend.source_filename
-        if hasattr(labeled_frame.video.backend, "source_filename")
-        else labeled_frame.video.backend.filename
-    )
+
+    # Extract video path with fallbacks for embedded videos
+    video = labeled_frame.video
+    video_path = None
+    if video is not None:
+        backend = getattr(video, "backend", None)
+        if backend is not None:
+            # Try source_filename first (for embedded videos with provenance)
+            video_path = getattr(backend, "source_filename", None)
+            if video_path is None:
+                video_path = getattr(backend, "filename", None)
+        # Fallback to video.filename if backend doesn't have it
+        if video_path is None:
+            video_path = getattr(video, "filename", None)
+            # Handle list filenames (image sequences)
+            if isinstance(video_path, list) and video_path:
+                video_path = video_path[0]
+    # Final fallback: use a unique identifier
+    if video_path is None:
+        video_path = f"video_{id(video)}" if video is not None else "unknown"
+
     for instance in labeled_frame.instances:
         match_instance = MatchInstance(
             instance=instance, frame_idx=frame_idx, video_path=video_path
@@ -47,6 +63,10 @@ def find_frame_pairs(
 ) -> List[Tuple[sio.LabeledFrame, sio.LabeledFrame]]:
     """Find corresponding frames across two sets of labels.
 
+    This function uses sleap-io's robust video matching API to handle various
+    scenarios including embedded videos, cross-platform paths, and videos with
+    different metadata.
+
     Args:
         labels_gt: A `sio.Labels` instance with ground truth instances.
         labels_pr: A `sio.Labels` instance with predicted instances.
@@ -56,16 +76,15 @@ def find_frame_pairs(
     Returns:
         A list of pairs of `sio.LabeledFrame`s in the form `(frame_gt, frame_pr)`.
     """
+    # Use sleap-io's robust video matching API (added in 0.6.2)
+    # The match() method returns a MatchResult with video_map: {pred_video: gt_video}
+    match_result = labels_gt.match(labels_pr)
+
     frame_pairs = []
-    for video_gt in labels_gt.videos:
-        # Find matching video instance in predictions.
-        video_pr = None
-        for video in labels_pr.videos:
-            if video_gt.matches_content(video) and video_gt.matches_path(video):
-                video_pr = video
-                break
-
-        if video_pr is None:
+    # Iterate over matched video pairs (pred_video -> gt_video mapping)
+    for video_pr, video_gt in match_result.video_map.items():
+        if video_gt is None:
+            # No match found for this prediction video
             continue
 
         # Find labeled frames in this video.
@@ -786,11 +805,26 @@ def run_evaluation(
     """Evaluate SLEAP-NN model predictions against ground truth labels."""
     logger.info("Loading ground truth labels...")
     ground_truth_instances = sio.load_slp(ground_truth_path)
+    logger.info(
+        f"  Ground truth: {len(ground_truth_instances.videos)} videos, "
+        f"{len(ground_truth_instances.labeled_frames)} frames"
+    )
 
     logger.info("Loading predicted labels...")
     predicted_instances = sio.load_slp(predicted_path)
+    logger.info(
+        f"  Predictions: {len(predicted_instances.videos)} videos, "
+        f"{len(predicted_instances.labeled_frames)} frames"
+    )
+
+    logger.info("Matching videos and frames...")
+    # Get match stats before creating evaluator
+    match_result = ground_truth_instances.match(predicted_instances)
+    logger.info(
+        f"  Videos matched: {match_result.n_videos_matched}/{len(match_result.video_map)}"
+    )
 
-    logger.info("Creating evaluator...")
+    logger.info("Matching instances...")
     evaluator = Evaluator(
         ground_truth_instances=ground_truth_instances,
         predicted_instances=predicted_instances,
@@ -799,21 +833,38 @@ def run_evaluation(
         match_threshold=match_threshold,
         user_labels_only=user_labels_only,
     )
+    logger.info(
+        f"  Frame pairs: {len(evaluator.frame_pairs)}, "
+        f"Matched instances: {len(evaluator.positive_pairs)}, "
+        f"Unmatched GT: {len(evaluator.false_negatives)}"
+    )
 
     logger.info("Computing evaluation metrics...")
     metrics = evaluator.evaluate()
 
+    # Compute PCK at specific thresholds (5 and 10 pixels)
+    dists = metrics["distance_metrics"]["dists"]
+    dists_clean = np.copy(dists)
+    dists_clean[np.isnan(dists_clean)] = np.inf
+    pck_5 = (dists_clean < 5).mean()
+    pck_10 = (dists_clean < 10).mean()
+
     # Print key metrics
     logger.info("Evaluation Results:")
-    logger.info(f"mOKS: {metrics['mOKS']['mOKS']:.4f}")
-    logger.info(f"mAP (OKS VOC): {metrics['voc_metrics']['oks_voc.mAP']:.4f}")
-    logger.info(f"mAR (OKS VOC): {metrics['voc_metrics']['oks_voc.mAR']:.4f}")
-    logger.info(f"Average Distance: {metrics['distance_metrics']['avg']:.4f}")
-    logger.info(f"mPCK: {metrics['pck_metrics']['mPCK']:.4f}")
+    logger.info(f"  mOKS: {metrics['mOKS']['mOKS']:.4f}")
+    logger.info(f"  mAP (OKS VOC): {metrics['voc_metrics']['oks_voc.mAP']:.4f}")
+    logger.info(f"  mAR (OKS VOC): {metrics['voc_metrics']['oks_voc.mAR']:.4f}")
+    logger.info(f"  Average Distance: {metrics['distance_metrics']['avg']:.2f} px")
+    logger.info(f"  dist.p50: {metrics['distance_metrics']['p50']:.2f} px")
+    logger.info(f"  dist.p95: {metrics['distance_metrics']['p95']:.2f} px")
+    logger.info(f"  dist.p99: {metrics['distance_metrics']['p99']:.2f} px")
+    logger.info(f"  mPCK: {metrics['pck_metrics']['mPCK']:.4f}")
+    logger.info(f"  PCK@5px: {pck_5:.4f}")
+    logger.info(f"  PCK@10px: {pck_10:.4f}")
     logger.info(
-        f"Visibility Precision: {metrics['visibility_metrics']['precision']:.4f}"
+        f"  Visibility Precision: {metrics['visibility_metrics']['precision']:.4f}"
     )
-    logger.info(f"Visibility Recall: {metrics['visibility_metrics']['recall']:.4f}")
+    logger.info(f"  Visibility Recall: {metrics['visibility_metrics']['recall']:.4f}")
 
     # Save metrics if path provided
     if save_metrics:

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/inference/bottomup.py

@@ -1,5 +1,6 @@
 """Inference modules for BottomUp models."""
 
+import logging
 from typing import Dict, Optional
 import torch
 import lightning as L
@@ -7,6 +8,8 @@ from sleap_nn.inference.peak_finding import find_local_peaks
 from sleap_nn.inference.paf_grouping import PAFScorer
 from sleap_nn.inference.identity import classify_peaks_from_maps
 
+logger = logging.getLogger(__name__)
+
 
 class BottomUpInferenceModel(L.LightningModule):
     """BottomUp Inference model.
@@ -63,8 +66,28 @@ class BottomUpInferenceModel(L.LightningModule):
         return_pafs: Optional[bool] = False,
         return_paf_graph: Optional[bool] = False,
         input_scale: float = 1.0,
+        max_peaks_per_node: Optional[int] = None,
     ):
-        """Initialise the model attributes."""
+        """Initialise the model attributes.
+
+        Args:
+            torch_model: A `nn.Module` that accepts images and predicts confidence maps.
+            paf_scorer: A `PAFScorer` instance for grouping instances.
+            cms_output_stride: Output stride of confidence maps relative to images.
+            pafs_output_stride: Output stride of PAFs relative to images.
+            peak_threshold: Minimum confidence map value for valid peaks.
+            refinement: Peak refinement method: None, "integral", or "local".
+            integral_patch_size: Size of patches for integral refinement.
+            return_confmaps: If True, return confidence maps in output.
+            return_pafs: If True, return PAFs in output.
+            return_paf_graph: If True, return intermediate PAF graph in output.
+            input_scale: Scale factor applied to input images.
+            max_peaks_per_node: Maximum number of peaks allowed per node before
+                skipping PAF scoring. If any node has more peaks than this limit,
+                empty predictions are returned. This prevents combinatorial explosion
+                during early training when confidence maps are noisy. Set to None to
+                disable this check (default). Recommended value: 100.
+        """
         super().__init__()
         self.torch_model = torch_model
         self.paf_scorer = paf_scorer
@@ -77,6 +100,7 @@ class BottomUpInferenceModel(L.LightningModule):
         self.return_pafs = return_pafs
         self.return_paf_graph = return_paf_graph
         self.input_scale = input_scale
+        self.max_peaks_per_node = max_peaks_per_node
 
     def _generate_cms_peaks(self, cms):
         # TODO: append nans to batch them -> tensor (vectorize the initial paf grouping steps)
@@ -124,26 +148,68 @@ class BottomUpInferenceModel(L.LightningModule):
         )  # (batch, h, w, 2*edges)
         cms_peaks, cms_peak_vals, cms_peak_channel_inds = self._generate_cms_peaks(cms)
 
-        (
-            predicted_instances,
-            predicted_peak_scores,
-            predicted_instance_scores,
-            edge_inds,
-            edge_peak_inds,
-            line_scores,
-        ) = self.paf_scorer.predict(
-            pafs=pafs,
-            peaks=cms_peaks,
-            peak_vals=cms_peak_vals,
-            peak_channel_inds=cms_peak_channel_inds,
-        )
-
-        predicted_instances = [p / self.input_scale for p in predicted_instances]
-        predicted_instances_adjusted = []
-        for idx, p in enumerate(predicted_instances):
-            predicted_instances_adjusted.append(
-                p / inputs["eff_scale"][idx].to(p.device)
+        # Check if too many peaks per node (prevents combinatorial explosion)
+        skip_paf_scoring = False
+        if self.max_peaks_per_node is not None:
+            n_nodes = cms.shape[1]
+            for b in range(self.batch_size):
+                for node_idx in range(n_nodes):
+                    n_peaks = int((cms_peak_channel_inds[b] == node_idx).sum().item())
+                    if n_peaks > self.max_peaks_per_node:
+                        logger.warning(
+                            f"Skipping PAF scoring: node {node_idx} has {n_peaks} peaks "
+                            f"(max_peaks_per_node={self.max_peaks_per_node}). "
+                            f"Model may need more training."
+                        )
+                        skip_paf_scoring = True
+                        break
+                if skip_paf_scoring:
+                    break
+
+        if skip_paf_scoring:
+            # Return empty predictions for each sample
+            device = cms.device
+            n_nodes = cms.shape[1]
+            predicted_instances_adjusted = []
+            predicted_peak_scores = []
+            predicted_instance_scores = []
+            for _ in range(self.batch_size):
+                predicted_instances_adjusted.append(
+                    torch.full((0, n_nodes, 2), float("nan"), device=device)
+                )
+                predicted_peak_scores.append(
+                    torch.full((0, n_nodes), float("nan"), device=device)
+                )
+                predicted_instance_scores.append(torch.tensor([], device=device))
+            edge_inds = [
+                torch.tensor([], dtype=torch.int32, device=device)
+            ] * self.batch_size
+            edge_peak_inds = [
+                torch.tensor([], dtype=torch.int32, device=device).reshape(0, 2)
+            ] * self.batch_size
+            line_scores = [torch.tensor([], device=device)] * self.batch_size
+        else:
+            (
+                predicted_instances,
+                predicted_peak_scores,
+                predicted_instance_scores,
+                edge_inds,
+                edge_peak_inds,
+                line_scores,
+            ) = self.paf_scorer.predict(
+                pafs=pafs,
+                peaks=cms_peaks,
+                peak_vals=cms_peak_vals,
+                peak_channel_inds=cms_peak_channel_inds,
             )
+
+            predicted_instances = [p / self.input_scale for p in predicted_instances]
+            predicted_instances_adjusted = []
+            for idx, p in enumerate(predicted_instances):
+                predicted_instances_adjusted.append(
+                    p / inputs["eff_scale"][idx].to(p.device)
+                )
+
         out = {
             "pred_instance_peaks": predicted_instances_adjusted,
             "pred_peak_values": predicted_peak_scores,

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a2}/sleap_nn/train.py

@@ -175,6 +175,7 @@ def train(
     wandb_save_viz_imgs_wandb: bool = False,
     wandb_resume_prv_runid: Optional[str] = None,
     wandb_group_name: Optional[str] = None,
+    wandb_delete_local_logs: Optional[bool] = None,
     optimizer: str = "Adam",
     learning_rate: float = 1e-3,
     amsgrad: bool = False,
@@ -353,6 +354,9 @@ def train(
         wandb_resume_prv_runid: Previous run ID if training should be resumed from a previous
             ckpt. Default: None
         wandb_group_name: Group name for the wandb run. Default: None.
+        wandb_delete_local_logs: If True, delete local wandb logs folder after training.
+            If False, keep the folder. If None (default), automatically delete if logging
+            online (wandb_mode != "offline") and keep if logging offline. Default: None.
         optimizer: Optimizer to be used. One of ["Adam", "AdamW"]. Default: "Adam".
         learning_rate: Learning rate of type float. Default: 1e-3.
         amsgrad: Enable AMSGrad with the optimizer. Default: False.
@@ -456,6 +460,7 @@ def train(
         wandb_save_viz_imgs_wandb=wandb_save_viz_imgs_wandb,
         wandb_resume_prv_runid=wandb_resume_prv_runid,
         wandb_group_name=wandb_group_name,
+        wandb_delete_local_logs=wandb_delete_local_logs,
         optimizer=optimizer,
         learning_rate=learning_rate,
         amsgrad=amsgrad,