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

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

@@ -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 → sleap_nn-0.1.0a1}/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.0a1}/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-0.1.0a0 → sleap_nn-0.1.0a1}/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.0a1}/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-0.1.0a0 → sleap_nn-0.1.0a1}/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.0a1}/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-0.1.0a0 → sleap_nn-0.1.0a1}/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 → sleap_nn-0.1.0a1}/sleap_nn.egg-info/PKG-INFO

@@ -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 → sleap_nn-0.1.0a1}/tests/config/test_trainer_config.py

@@ -123,6 +123,52 @@ def test_wandb_config():
     assert custom_conf.project == "test_project"
 
 
+def test_wandb_config_delete_local_logs():
+    """Test delete_local_logs field in WandBConfig.
+
+    Verifies default value and auto-detection logic.
+    """
+    # Check default value is None
+    conf = OmegaConf.structured(WandBConfig)
+    assert conf.delete_local_logs is None
+
+    # Test explicit True
+    custom_conf = OmegaConf.structured(WandBConfig(delete_local_logs=True))
+    assert custom_conf.delete_local_logs is True
+
+    # Test explicit False
+    custom_conf = OmegaConf.structured(WandBConfig(delete_local_logs=False))
+    assert custom_conf.delete_local_logs is False
+
+    # Test auto-detection logic: online mode (wandb_mode=None) should delete
+    config = WandBConfig(wandb_mode=None)
+    should_delete = config.delete_local_logs is True or (
+        config.delete_local_logs is None and config.wandb_mode != "offline"
+    )
+    assert should_delete is True
+
+    # Test auto-detection logic: offline mode should keep
+    config = WandBConfig(wandb_mode="offline")
+    should_delete = config.delete_local_logs is True or (
+        config.delete_local_logs is None and config.wandb_mode != "offline"
+    )
+    assert should_delete is False
+
+    # Test explicit True overrides offline mode
+    config = WandBConfig(wandb_mode="offline", delete_local_logs=True)
+    should_delete = config.delete_local_logs is True or (
+        config.delete_local_logs is None and config.wandb_mode != "offline"
+    )
+    assert should_delete is True
+
+    # Test explicit False overrides online mode
+    config = WandBConfig(wandb_mode=None, delete_local_logs=False)
+    should_delete = config.delete_local_logs is True or (
+        config.delete_local_logs is None and config.wandb_mode != "offline"
+    )
+    assert should_delete is False
+
+
 def test_optimizer_config(caplog):
     """optimizer_config tests.
 

{sleap_nn-0.1.0a0 → sleap_nn-0.1.0a1}/tests/test_cli.py

@@ -323,6 +323,22 @@ def sample_config(tmp_path, minimal_instance):
     return config_data
 
 
+def _strip_ansi_codes(text: str) -> str:
+    """Remove ANSI escape codes from text.
+
+    This is necessary because some libraries (PyTorch Lightning, Rich, tqdm)
+    emit ANSI codes even when NO_COLOR is set, which breaks YAML parsing.
+    """
+    import re
+
+    # Pattern matches ANSI escape sequences including:
+    # - CSI sequences: \x1b[...X (e.g., colors, cursor movement)
+    # - Private mode sequences: \x1b[?...X (e.g., cursor show/hide)
+    # - OSC sequences: \x1b]...X (e.g., window title)
+    ansi_pattern = re.compile(r"\x1b\[[?0-9;]*[a-zA-Z]|\x1b\][^\x07]*\x07")
+    return ansi_pattern.sub("", text)
+
+
 def test_main_cli(sample_config, tmp_path):
     cmd = [
         "uv",
@@ -351,6 +367,12 @@ def test_main_cli(sample_config, tmp_path):
     )
     OmegaConf.save(sample_config, (Path(tmp_path) / "test_config.yaml").as_posix())
 
+    import os
+    import re
+
+    # Environment to disable color output from rich/loguru/lightning
+    no_color_env = {**os.environ, "NO_COLOR": "1", "FORCE_COLOR": "0"}
+
     cmd = [
         "uv",
         "run",
@@ -368,13 +390,19 @@ def test_main_cli(sample_config, tmp_path):
         cmd,
         capture_output=True,
         text=True,
+        env=no_color_env,
     )
     # Exit code should be 0
     assert result.returncode == 0
     # Try to parse the output back into the yaml, truncate the beginning (starts with "data_config")
-    # Only keep stdout starting from "data_config"
-    stripped_out = result.stdout[result.stdout.find("data_config") :].strip()
-    stripped_out = stripped_out[: stripped_out.find(" | INFO") - 19]
+    # Only keep stdout starting from "data_config" and ending at the next log timestamp
+    # Strip ANSI codes first as some libraries emit them even with NO_COLOR
+    stdout_clean = _strip_ansi_codes(result.stdout)
+    stripped_out = stdout_clean[stdout_clean.find("data_config") :].strip()
+    # Find the next log line (timestamp pattern: YYYY-MM-DD HH:MM:SS |)
+    match = re.search(r"\n\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} \|", stripped_out)
+    if match:
+        stripped_out = stripped_out[: match.start()]
     output = OmegaConf.create(stripped_out)
     assert output == sample_config
 
@@ -400,11 +428,15 @@ def test_main_cli(sample_config, tmp_path):
         cmd,
         capture_output=True,
         text=True,
+        env=no_color_env,
     )
     # Exit code should be 0
     assert result.returncode == 0
-    stripped_out = result.stdout[result.stdout.find("data_config") :].strip()
-    stripped_out = stripped_out[: stripped_out.find(" | INFO") - 19]
+    stdout_clean = _strip_ansi_codes(result.stdout)
+    stripped_out = stdout_clean[stdout_clean.find("data_config") :].strip()
+    match = re.search(r"\n\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} \|", stripped_out)
+    if match:
+        stripped_out = stripped_out[: match.start()]
     output = OmegaConf.create(stripped_out)
     assert output == sample_config
 
@@ -429,6 +461,7 @@ def test_main_cli(sample_config, tmp_path):
         cmd,
         capture_output=True,
         text=True,
+        env=no_color_env,
     )
     assert (
         Path(f"{sample_config.trainer_config.ckpt_dir}")
@@ -440,8 +473,12 @@ def test_main_cli(sample_config, tmp_path):
     # Exit code should be 0
     assert result.returncode == 0
     # Check that overrides are applied
-    stripped_out = result.stdout[result.stdout.find("data_config") :].strip()
-    stripped_out = stripped_out[: stripped_out.find(" | INFO") - 19]
+    stdout_clean = _strip_ansi_codes(result.stdout)
+    stripped_out = stdout_clean[stdout_clean.find("data_config") :].strip()
+    # Find the next log line (timestamp pattern: YYYY-MM-DD HH:MM:SS |)
+    match = re.search(r"\n\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2} \|", stripped_out)
+    if match:
+        stripped_out = stripped_out[: match.start()]
     output = OmegaConf.create(stripped_out)
     assert output.trainer_config.max_epochs == 3
     assert output.data_config.preprocessing.scale == 1.5