sleap-nn 0.1.0a0__py3-none-any.whl → 0.1.0a1__py3-none-any.whl

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.0a1"
 
 # Public API
 from sleap_nn.evaluation import load_metrics

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/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

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/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,

sleap_nn/training/model_trainer.py

@@ -898,6 +898,17 @@ class ModelTrainer:
             )
             loggers.append(wandb_logger)
 
+            # Log message about wandb local logs cleanup
+            should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
+                wandb_config.delete_local_logs is None
+                and wandb_config.wandb_mode != "offline"
+            )
+            if should_delete_wandb_logs:
+                logger.info(
+                    "WandB local logs will be deleted after training completes. "
+                    "To keep logs, set trainer_config.wandb.delete_local_logs=false"
+                )
+
             # Learning rate monitor callback - logs LR at each step for dynamic schedulers
             # Only added when wandb is enabled since it requires a logger
             callbacks.append(LearningRateMonitor(logging_interval="step"))
@@ -1314,6 +1325,25 @@ class ModelTrainer:
             if self.trainer.global_rank == 0 and self.config.trainer_config.use_wandb:
                 wandb.finish()
 
+                # Delete local wandb logs if configured
+                wandb_config = self.config.trainer_config.wandb
+                should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
+                    wandb_config.delete_local_logs is None
+                    and wandb_config.wandb_mode != "offline"
+                )
+                if should_delete_wandb_logs:
+                    wandb_dir = (
+                        Path(self.config.trainer_config.ckpt_dir)
+                        / self.config.trainer_config.run_name
+                        / "wandb"
+                    )
+                    if wandb_dir.exists():
+                        logger.info(
+                            f"Deleting local wandb logs at {wandb_dir}... "
+                            "(set trainer_config.wandb.delete_local_logs=false to disable)"
+                        )
+                        shutil.rmtree(wandb_dir, ignore_errors=True)
+
             # delete image disk caching
             if (
                 self.config.data_config.data_pipeline_fw

{sleap_nn-0.1.0a0.dist-info → sleap_nn-0.1.0a1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.1.0a0
+Version: 0.1.0a1
 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

{sleap_nn-0.1.0a0.dist-info → sleap_nn-0.1.0a1.dist-info}/RECORD

@@ -1,11 +1,11 @@
 sleap_nn/.DS_Store,sha256=HY8amA79eHkt7o5VUiNsMxkc9YwW6WIPyZbYRj_JdSU,6148
-sleap_nn/__init__.py,sha256=DzQeiZIFUmfhpf6mk4j1AKAY2bofVMyIa31xbiSu-ls,1317
+sleap_nn/__init__.py,sha256=l5Lwiad8GOurqkAhMwWw8-UcpH6af2TnMURf-oKj_U8,1362
 sleap_nn/cli.py,sha256=U4hpEcOxK7a92GeItY95E2DRm5P1ME1GqU__mxaDcW0,21167
 sleap_nn/evaluation.py,sha256=3u7y85wFoBgCwOB2xOGTJIDrd2dUPWOo4m0s0oW3da4,31095
 sleap_nn/legacy_models.py,sha256=8aGK30DZv3pW2IKDBEWH1G2mrytjaxPQD4miPUehj0M,20258
 sleap_nn/predict.py,sha256=8QKjRbS-L-6HF1NFJWioBPv3HSzUpFr2oGEB5hRJzQA,35523
 sleap_nn/system_info.py,sha256=7tWe3y6s872nDbrZoHIdSs-w4w46Z4dEV2qCV-Fe7No,14711
-sleap_nn/train.py,sha256=fWx_b1HqkadQ-GM_VEM1frCd8WkzJLqRARBNn8UoUbo,27181
+sleap_nn/train.py,sha256=XvVhzMXL9rNQLx1-6jIcp5BAO1pR7AZjdphMn5ZX-_I,27558
 sleap_nn/architectures/__init__.py,sha256=w0XxQcx-CYyooszzvxRkKWiJkUg-26IlwQoGna8gn40,46
 sleap_nn/architectures/common.py,sha256=MLv-zdHsWL5Q2ct_Wv6SQbRS-5hrFtjK_pvBEfwx-vU,3660
 sleap_nn/architectures/convnext.py,sha256=l9lMJDxIMb-9MI3ShOtVwbOUMuwOLtSQlxiVyYHqjvE,13953
@@ -17,15 +17,15 @@ sleap_nn/architectures/unet.py,sha256=rAy2Omi6tv1MNW2nBn0Tw-94Nw_-1wFfCT3-IUyPcg
 sleap_nn/architectures/utils.py,sha256=L0KVs0gbtG8U75Sl40oH_r_w2ySawh3oQPqIGi54HGo,2171
 sleap_nn/config/__init__.py,sha256=l0xV1uJsGJfMPfWAqlUR7Ivu4cSCWsP-3Y9ueyPESuk,42
 sleap_nn/config/data_config.py,sha256=5a5YlXm4V9qGvkqgFNy6o0XJ_Q06UFjpYJXmNHfvXEI,24021
-sleap_nn/config/get_config.py,sha256=vN_aOPTj9F-QBqGGfVSv8_aFSAYl-RfXY0pdbdcqjcM,42021
+sleap_nn/config/get_config.py,sha256=rjNUffKU9z-ohLwrOVmJNGCqwUM93eh68h4KJfrSy8Y,42396
 sleap_nn/config/model_config.py,sha256=XFIbqFno7IkX0Se5WF_2_7aUalAlC2SvpDe-uP2TttM,57582
-sleap_nn/config/trainer_config.py,sha256=PaoNtRSNc2xgzwN955aR9kTZL8IxCWdevGljLxS6jOk,28073
+sleap_nn/config/trainer_config.py,sha256=ZMXxns6VYakgYHRhkM541Eje76DdaTdDi4FFPNjJtP4,28413
 sleap_nn/config/training_job_config.py,sha256=v12_ME_tBUg8JFwOxJNW4sDQn-SedDhiJOGz-TlRwT0,5861
 sleap_nn/config/utils.py,sha256=GgWgVs7_N7ifsJ5OQG3_EyOagNyN3Dx7wS2BAlkaRkg,5553
 sleap_nn/data/__init__.py,sha256=eMNvFJFa3gv5Rq8oK5wzo6zt1pOlwUGYf8EQii6bq7c,54
 sleap_nn/data/augmentation.py,sha256=Kqw_DayPth_DBsmaO1G8Voou_-cYZuSPOjSQWSajgRI,13618
 sleap_nn/data/confidence_maps.py,sha256=PTRqZWSAz1S7viJhxu7QgIC1aHiek97c_dCUsKUwG1o,6217
-sleap_nn/data/custom_datasets.py,sha256=2qAyLeiCPI9uudFFP7zlj6d_tbxc5OVzpnzT23mRkVw,98472
+sleap_nn/data/custom_datasets.py,sha256=SO-aNB1-bB9DL5Zw-oGYDsliBxwI4iKX_FmwgZjKOgQ,99975
 sleap_nn/data/edge_maps.py,sha256=75qG_7zHRw7fC8JUCVI2tzYakIoxxneWWmcrTwjcHPo,12519
 sleap_nn/data/identity.py,sha256=7vNup6PudST4yDLyDT9wDO-cunRirTEvx4sP77xrlfk,5193
 sleap_nn/data/instance_centroids.py,sha256=SF-3zJt_VMTbZI5ssbrvmZQZDd3684bn55EAtvcbQ6o,2172
@@ -55,11 +55,11 @@ sleap_nn/training/__init__.py,sha256=vNTKsIJPZHJwFSKn5PmjiiRJunR_9e7y4_v0S6rdF8U
 sleap_nn/training/callbacks.py,sha256=TVnQ6plNC2MnlTiY2rSCRuw2WRk5cQSziek_VPUcOEg,25994
 sleap_nn/training/lightning_modules.py,sha256=G3c4xJkYWW-iSRawzkgTqkGd4lTsbPiMTcB5Nvq7jes,85512
 sleap_nn/training/losses.py,sha256=gbdinUURh4QUzjmNd2UJpt4FXwecqKy9gHr65JZ1bZk,1632
-sleap_nn/training/model_trainer.py,sha256=InDKHrQxBwbltZKutW4yrBR9NThLdRpWNUGhmB0xAi4,57863
+sleap_nn/training/model_trainer.py,sha256=loCmEX0DfBtdV_pN-W8s31fn2_L-lbpWaq3OQXeSp-0,59337
 sleap_nn/training/utils.py,sha256=ivdkZEI0DkTCm6NPszsaDOh9jSfozkONZdl6TvvQUWI,20398
-sleap_nn-0.1.0a0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-sleap_nn-0.1.0a0.dist-info/METADATA,sha256=lxSmGNTUg9eetqHCvhw8Tv5zJtia-dIM5RzOeoDccj8,5637
-sleap_nn-0.1.0a0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sleap_nn-0.1.0a0.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
-sleap_nn-0.1.0a0.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
-sleap_nn-0.1.0a0.dist-info/RECORD,,
+sleap_nn-0.1.0a1.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+sleap_nn-0.1.0a1.dist-info/METADATA,sha256=h3d4WPIu_JunY32jaRqJ4-fXp4KruTWT57FWb3L6dps,5637
+sleap_nn-0.1.0a1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sleap_nn-0.1.0a1.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
+sleap_nn-0.1.0a1.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
+sleap_nn-0.1.0a1.dist-info/RECORD,,