sleap-nn 0.0.5__py3-none-any.whl → 0.1.0__py3-none-any.whl

sleap_nn/train.py

@@ -6,6 +6,7 @@ from datetime import datetime
 from time import time
 from omegaconf import DictConfig, OmegaConf
 from typing import Any, Dict, Optional, List, Tuple, Union
+import sleap_io as sio
 from sleap_nn.config.training_job_config import TrainingJobConfig
 from sleap_nn.training.model_trainer import ModelTrainer
 from sleap_nn.predict import run_inference as predict
@@ -15,15 +16,31 @@ from sleap_nn.config.get_config import (
     get_model_config,
     get_data_config,
 )
+from sleap_nn.system_info import get_startup_info_string
 
 
-def run_training(config: DictConfig):
-    """Create ModelTrainer instance and start training."""
+def run_training(
+    config: DictConfig,
+    train_labels: Optional[List[sio.Labels]] = None,
+    val_labels: Optional[List[sio.Labels]] = None,
+):
+    """Create ModelTrainer instance and start training.
+
+    Args:
+        config: Training configuration as a DictConfig.
+        train_labels: List of Labels objects for training.
+        val_labels: List of Labels objects for validation.
+            If not provided, the labels will be loaded from paths in the config.
+    """
     start_train_time = time()
     start_timestamp = str(datetime.now())
     logger.info(f"Started training at: {start_timestamp}")
+    logger.info(get_startup_info_string())
 
-    trainer = ModelTrainer.get_model_trainer_from_config(config)
+    # provide the labels as the train labels, val labels will be split from the train labels
+    trainer = ModelTrainer.get_model_trainer_from_config(
+        config, train_labels=train_labels, val_labels=val_labels
+    )
     trainer.train()
 
     finish_timestamp = str(datetime.now())
@@ -39,48 +56,44 @@ def run_training(config: DictConfig):
         # run inference on val dataset
         if trainer.config.trainer_config.save_ckpt:
             data_paths = {}
-            for index, path in enumerate(trainer.config.data_config.train_labels_path):
-                logger.info(
-                    f"Training labels path for index {index}: {(Path(trainer.config.trainer_config.ckpt_dir) / trainer.config.trainer_config.run_name).as_posix()}"
-                )
-                data_paths[f"train_{index}"] = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"labels_train_gt_{index}.slp"
+            run_path = (
+                Path(trainer.config.trainer_config.ckpt_dir)
+                / trainer.config.trainer_config.run_name
+            )
+            for index, _ in enumerate(trainer.train_labels):
+                logger.info(f"Run path for index {index}: {run_path.as_posix()}")
+                data_paths[f"train.{index}"] = (
+                    run_path / f"labels_gt.train.{index}.slp"
                 ).as_posix()
-                data_paths[f"val_{index}"] = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"labels_val_gt_{index}.slp"
+                data_paths[f"val.{index}"] = (
+                    run_path / f"labels_gt.val.{index}.slp"
                 ).as_posix()
 
-            if (
-                OmegaConf.select(config, "data_config.test_file_path", default=None)
-                is not None
-            ):
-                data_paths["test"] = config.data_config.test_file_path
+            # Handle test_file_path as either a string or list of strings
+            test_file_path = OmegaConf.select(
+                config, "data_config.test_file_path", default=None
+            )
+            if test_file_path is not None:
+                # Normalize to list of strings
+                if isinstance(test_file_path, str):
+                    test_paths = [test_file_path]
+                else:
+                    test_paths = list(test_file_path)
+                # Add each test path to data_paths (always use index for consistency)
+                for idx, test_path in enumerate(test_paths):
+                    data_paths[f"test.{idx}"] = test_path
 
             for d_name, path in data_paths.items():
-                pred_path = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"pred_{d_name}.slp"
-                )
-                metrics_path = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"{d_name}_pred_metrics.npz"
-                )
+                # d_name is now in format: "train.0", "val.0", "test.0", etc.
+                pred_path = run_path / f"labels_pr.{d_name}.slp"
+                metrics_path = run_path / f"metrics.{d_name}.npz"
 
                 pred_labels = predict(
                     data_path=path,
-                    model_paths=[
-                        Path(trainer.config.trainer_config.ckpt_dir)
-                        / trainer.config.trainer_config.run_name
-                    ],
+                    model_paths=[run_path],
                     peak_threshold=0.2,
                     make_labels=True,
-                    device=trainer.trainer.strategy.root_device,
+                    device=str(trainer.trainer.strategy.root_device),
                     output_path=pred_path,
                     ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                     ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
@@ -105,12 +118,77 @@ def run_training(config: DictConfig):
                 logger.info(f"p90 dist: {metrics['distance_metrics']['p90']}")
                 logger.info(f"p50 dist: {metrics['distance_metrics']['p50']}")
 
+                # Log test metrics to wandb summary
+                if (
+                    d_name.startswith("test")
+                    and trainer.config.trainer_config.use_wandb
+                ):
+                    import wandb
+
+                    if wandb.run is not None:
+                        summary_metrics = {
+                            f"eval/{d_name}/mOKS": metrics["mOKS"]["mOKS"],
+                            f"eval/{d_name}/oks_voc_mAP": metrics["voc_metrics"][
+                                "oks_voc.mAP"
+                            ],
+                            f"eval/{d_name}/oks_voc_mAR": metrics["voc_metrics"][
+                                "oks_voc.mAR"
+                            ],
+                            f"eval/{d_name}/mPCK": metrics["pck_metrics"]["mPCK"],
+                            f"eval/{d_name}/PCK_5": metrics["pck_metrics"]["PCK@5"],
+                            f"eval/{d_name}/PCK_10": metrics["pck_metrics"]["PCK@10"],
+                            f"eval/{d_name}/distance_avg": metrics["distance_metrics"][
+                                "avg"
+                            ],
+                            f"eval/{d_name}/distance_p50": metrics["distance_metrics"][
+                                "p50"
+                            ],
+                            f"eval/{d_name}/distance_p95": metrics["distance_metrics"][
+                                "p95"
+                            ],
+                            f"eval/{d_name}/distance_p99": metrics["distance_metrics"][
+                                "p99"
+                            ],
+                            f"eval/{d_name}/visibility_precision": metrics[
+                                "visibility_metrics"
+                            ]["precision"],
+                            f"eval/{d_name}/visibility_recall": metrics[
+                                "visibility_metrics"
+                            ]["recall"],
+                        }
+                        for key, value in summary_metrics.items():
+                            wandb.run.summary[key] = value
+
+            # Finish wandb run and cleanup after all evaluation is complete
+            if trainer.config.trainer_config.use_wandb:
+                import wandb
+                import shutil
+
+                if wandb.run is not None:
+                    wandb.finish()
+
+                # Delete local wandb logs if configured
+                wandb_config = trainer.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 = run_path / "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)
+
 
 def train(
     train_labels_path: Optional[List[str]] = None,
     val_labels_path: Optional[List[str]] = None,
     validation_fraction: float = 0.1,
-    test_file_path: Optional[str] = None,
+    use_same_data_for_val: bool = False,
+    test_file_path: Optional[Union[str, List[str]]] = None,
     provider: str = "LabelsReader",
     user_instances_only: bool = True,
     data_pipeline_fw: str = "torch_dataset",
@@ -124,9 +202,10 @@ def train(
     max_width: Optional[int] = None,
     crop_size: Optional[int] = None,
     min_crop_size: Optional[int] = 100,
-    use_augmentations_train: bool = False,
+    crop_padding: Optional[int] = None,
+    use_augmentations_train: bool = True,
     intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
-    geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
+    geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = "rotation",
     init_weight: str = "default",
     pretrained_backbone_weights: Optional[str] = None,
     pretrained_head_weights: Optional[str] = None,
@@ -160,6 +239,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,
@@ -188,7 +268,11 @@ def train(
             training set to sample for generating the validation set. The remaining
             labeled frames will be left in the training set. If the `validation_labels`
             are already specified, this has no effect. Default: 0.1.
-        test_file_path: Path to test dataset (`.slp` file or `.mp4` file).
+        use_same_data_for_val: If `True`, use the same data for both training and
+            validation (train = val). Useful for intentional overfitting on small
+            datasets. When enabled, `val_labels_path` and `validation_fraction` are
+            ignored. Default: False.
+        test_file_path: Path or list of paths to test dataset(s) (`.slp` file(s) or `.mp4` file(s)).
             Note: This is used to get evaluation on test set after training is completed.
         provider: Provider class to read the input sleap files. Only "LabelsReader"
             supported for the training pipeline. Default: "LabelsReader".
@@ -210,16 +294,19 @@ def train(
             is set to True, then we convert the image to grayscale (single-channel)
             image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
         scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
-        max_height: Maximum height the image should be padded to. If not provided, the
+        max_height: Maximum height the original image should be resized and padded to. If not provided, the
             original image size will be retained. Default: None.
-        max_width: Maximum width the image should be padded to. If not provided, the
+        max_width: Maximum width the original image should be resized and padded to. If not provided, the
             original image size will be retained. Default: None.
         crop_size: Crop size of each instance for centered-instance model.
             If `None`, this would be automatically computed based on the largest instance
-            in the `sio.Labels` file. Default: None.
+            in the `sio.Labels` file. If `scale` is provided, then the cropped image will be resized according to `scale`. Default: None.
         min_crop_size: Minimum crop size to be used if `crop_size` is `None`. Default: 100.
+        crop_padding: Padding in pixels to add around instance bounding box when computing
+            crop size. If `None`, padding is auto-computed based on augmentation settings.
+            Only used when `crop_size` is `None`. Default: None.
         use_augmentations_train: True if the data augmentation should be applied to the
-            training data, else False. Default: False.
+            training data, else False. Default: True.
         intensity_aug: One of ["uniform_noise", "gaussian_noise", "contrast", "brightness"]
             or list of strings from the above allowed values. To have custom values, pass
             a dict with the structure in `sleap_nn.config.data_config.IntensityConfig`.
@@ -231,7 +318,8 @@ def train(
             or list of strings from the above allowed values. To have custom values, pass
             a dict with the structure in `sleap_nn.config.data_config.GeometryConfig`.
             For eg: {
-                        "rotation": 45,
+                        "rotation_min": -45,
+                        "rotation_max": 45,
                         "affine_p": 1.0
                     }
         init_weight: model weights initialization method. "default" uses kaiming uniform
@@ -331,6 +419,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.
@@ -376,6 +467,7 @@ def train(
         train_labels_path=train_labels_path,
         val_labels_path=val_labels_path,
         validation_fraction=validation_fraction,
+        use_same_data_for_val=use_same_data_for_val,
         test_file_path=test_file_path,
         provider=provider,
         user_instances_only=user_instances_only,
@@ -390,6 +482,7 @@ def train(
         max_width=max_width,
         crop_size=crop_size,
         min_crop_size=min_crop_size,
+        crop_padding=crop_padding,
         use_augmentations_train=use_augmentations_train,
         intensity_aug=intensity_aug,
         geometry_aug=geometry_aug,
@@ -432,6 +525,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,