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

sleap_nn/config/data_config.py

@@ -6,7 +6,7 @@ the parameters required to initialize the data config.
 
 from attrs import define, field, validators
 from omegaconf import MISSING
-from typing import Optional, Tuple, Any, List
+from typing import Optional, Tuple, Any, List, Union
 from loguru import logger
 import sleap_io as sio
 import yaml
@@ -20,11 +20,15 @@ class PreprocessingConfig:
     Attributes:
         ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one channel when this is set to `True`, then the images from single-channel is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. *Default*: `False`.
         ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this 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`.
-        max_height: (int) Maximum height the image should be padded to. If not provided, the original image size will be retained. *Default*: `None`.
-        max_width: (int) Maximum width the image should be padded to. If not provided, the original image size will be retained. *Default*: `None`.
+        max_height: (int) 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: (int) Maximum width the original image should be resized and padded to. If not provided, the original image size will be retained. *Default*: `None`.
         scale: (float) Factor to resize the image dimensions by, specified as a float. *Default*: `1.0`.
-        crop_size: (int) 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`.
+        crop_size: (int) 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.
+            If `scale` is provided, then the cropped image will be resized according to `scale`.*Default*: `None`.
         min_crop_size: (int) Minimum crop size to be used if `crop_size` is `None`. *Default*: `100`.
+        crop_padding: (int) Padding in pixels to add around the instance bounding box when computing crop size.
+            If `None`, padding is auto-computed based on augmentation settings (rotation/scale).
+            Only used when `crop_size` is `None`. *Default*: `None`.
     """
 
     ensure_rgb: bool = False
@@ -36,6 +40,7 @@ class PreprocessingConfig:
     )
     crop_size: Optional[int] = None
     min_crop_size: Optional[int] = 100  # to help app work in case of error
+    crop_padding: Optional[int] = None
 
     def validate_scale(self):
         """Scale Validation.
@@ -104,11 +109,14 @@ class GeometricConfig:
     Attributes:
         rotation_min: (float) Minimum rotation angle in degrees. A random angle in (rotation_min, rotation_max) will be sampled and applied to both images and keypoints. Set to 0 to disable rotation augmentation. *Default*: `-15.0`.
         rotation_max: (float) Maximum rotation angle in degrees. A random angle in (rotation_min, rotation_max) will be sampled and applied to both images and keypoints. Set to 0 to disable rotation augmentation. *Default*: `15.0`.
+        rotation_p: (float, optional) Probability of applying random rotation independently. If set, rotation is applied separately from scale/translate. If `None`, falls back to `affine_p` for bundled behavior. *Default*: `1.0`.
         scale_min: (float) Minimum scaling factor. If scale_min and scale_max are provided, the scale is randomly sampled from the range scale_min <= scale <= scale_max for isotropic scaling. *Default*: `0.9`.
         scale_max: (float) Maximum scaling factor. If scale_min and scale_max are provided, the scale is randomly sampled from the range scale_min <= scale <= scale_max for isotropic scaling. *Default*: `1.1`.
+        scale_p: (float, optional) Probability of applying random scaling independently. If set, scaling is applied separately from rotation/translate. If `None`, falls back to `affine_p` for bundled behavior. *Default*: `1.0`.
         translate_width: (float) Maximum absolute fraction for horizontal translation. For example, if translate_width=a, then horizontal shift is randomly sampled in the range -img_width * a < dx < img_width * a. Will not translate by default. *Default*: `0.0`.
         translate_height: (float) Maximum absolute fraction for vertical translation. For example, if translate_height=a, then vertical shift is randomly sampled in the range -img_height * a < dy < img_height * a. Will not translate by default. *Default*: `0.0`.
-        affine_p: (float) Probability of applying random affine transformations. *Default*: `0.0`.
+        translate_p: (float, optional) Probability of applying random translation independently. If set, translation is applied separately from rotation/scale. If `None`, falls back to `affine_p` for bundled behavior. *Default*: `None`.
+        affine_p: (float) Probability of applying random affine transformations (rotation, scale, translate bundled together). Used for backwards compatibility when individual `*_p` params are not set. *Default*: `0.0`.
         erase_scale_min: (float) Minimum value of range of proportion of erased area against input image. *Default*: `0.0001`.
         erase_scale_max: (float) Maximum value of range of proportion of erased area against input image. *Default*: `0.01`.
         erase_ratio_min: (float) Minimum value of range of aspect ratio of erased area. *Default*: `1.0`.
@@ -121,10 +129,13 @@ class GeometricConfig:
 
     rotation_min: float = field(default=-15.0, validator=validators.ge(-180))
     rotation_max: float = field(default=15.0, validator=validators.le(180))
+    rotation_p: Optional[float] = field(default=1.0)
     scale_min: float = field(default=0.9, validator=validators.ge(0))
     scale_max: float = field(default=1.1, validator=validators.ge(0))
+    scale_p: Optional[float] = field(default=1.0)
     translate_width: float = 0.0
     translate_height: float = 0.0
+    translate_p: Optional[float] = field(default=None)
     affine_p: float = field(default=0.0, validator=validate_proportion)
     erase_scale_min: float = 0.0001
     erase_scale_max: float = 0.01
@@ -149,6 +160,28 @@ class AugmentationConfig:
     geometric: Optional[GeometricConfig] = None
 
 
+def validate_test_file_path(instance, attribute, value):
+    """Validate test_file_path to accept str or List[str].
+
+    Args:
+        instance: The instance being validated.
+        attribute: The attribute being validated.
+        value: The value to validate.
+
+    Raises:
+        ValueError: If value is not None, str, or list of strings.
+    """
+    if value is None:
+        return
+    if isinstance(value, str):
+        return
+    if isinstance(value, (list, tuple)) and all(isinstance(p, str) for p in value):
+        return
+    message = f"{attribute.name} must be a string or list of strings, got {type(value).__name__}"
+    logger.error(message)
+    raise ValueError(message)
+
+
 @define
 class DataConfig:
     """Data configuration.
@@ -157,13 +190,16 @@ class DataConfig:
         train_labels_path: (List[str]) List of paths to training data (`.slp` file(s)). *Default*: `None`.
         val_labels_path: (List[str]) List of paths to validation data (`.slp` file(s)). *Default*: `None`.
         validation_fraction: (float) Float between 0 and 1 specifying the fraction of the 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: (str) Path to test dataset (`.slp` file or `.mp4` file). *Note*: This is used only with CLI to get evaluation on test set after training is completed. *Default*: `None`.
+        use_same_data_for_val: (bool) 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: (str or List[str]) Path or list of paths to test dataset(s) (`.slp` file(s) or `.mp4` file(s)). *Note*: This is used only with CLI to get evaluation on test set after training is completed. *Default*: `None`.
         provider: (str) Provider class to read the input sleap files. Only "LabelsReader" is currently supported for the training pipeline. *Default*: `"LabelsReader"`.
         user_instances_only: (bool) `True` if only user labeled instances should be used for training. If `False`, both user labeled and predicted instances would be used. *Default*: `True`.
         data_pipeline_fw: (str) Framework to create the data loaders. One of [`torch_dataset`, `torch_dataset_cache_img_memory`, `torch_dataset_cache_img_disk`]. *Default*: `"torch_dataset"`. (Note: When using `torch_dataset`, `num_workers` in `trainer_config` should be set to 0 as multiprocessing doesn't work with pickling video backends.)
         cache_img_path: (str) Path to save `.jpg` images created with `torch_dataset_cache_img_disk` data pipeline framework. If `None`, the path provided in `trainer_config.save_ckpt` is used. The `train_imgs` and `val_imgs` dirs are created inside this path. *Default*: `None`.
         use_existing_imgs: (bool) Use existing train and val images/ chunks in the `cache_img_path` for `torch_dataset_cache_img_disk` frameworks. If `True`, the `cache_img_path` should have `train_imgs` and `val_imgs` dirs. *Default*: `False`.
         delete_cache_imgs_after_training: (bool) If `False`, the images (torch_dataset_cache_img_disk) are retained after training. Else, the files are deleted. *Default*: `True`.
+        parallel_caching: (bool) If `True`, use parallel processing to cache images (significantly faster for large datasets). *Default*: `True`.
+        cache_workers: (int) Number of worker threads for parallel caching. If 0, uses min(4, cpu_count). *Default*: `0`.
         preprocessing: Configuration options related to data preprocessing.
         use_augmentations_train: (bool) True if the data augmentation should be applied to the training data, else False. *Default*: `True`.
         augmentation_config: Configurations related to augmentation. (only if `use_augmentations_train` is `True`)
@@ -173,16 +209,23 @@ class DataConfig:
     train_labels_path: Optional[List[str]] = None
     val_labels_path: Optional[List[str]] = None  # TODO : revisit MISSING!
     validation_fraction: float = 0.1
-    test_file_path: Optional[str] = None
+    use_same_data_for_val: bool = False
+    test_file_path: Optional[Any] = field(
+        default=None, validator=validate_test_file_path
+    )
     provider: str = "LabelsReader"
     user_instances_only: bool = True
     data_pipeline_fw: str = "torch_dataset"
     cache_img_path: Optional[str] = None
     use_existing_imgs: bool = False
     delete_cache_imgs_after_training: bool = True
+    parallel_caching: bool = True
+    cache_workers: int = 0
     preprocessing: PreprocessingConfig = field(factory=PreprocessingConfig)
     use_augmentations_train: bool = True
-    augmentation_config: Optional[AugmentationConfig] = None
+    augmentation_config: Optional[AugmentationConfig] = field(
+        factory=lambda: AugmentationConfig(geometric=GeometricConfig())
+    )
     skeletons: Optional[list] = None
 
 

sleap_nn/config/get_config.py

@@ -131,27 +131,18 @@ def get_aug_config(
 
         for g in geometric_aug:
             if g == "rotation":
-                aug_config.geometric.affine_p = 1.0
-                aug_config.geometric.scale_min = 1.0
-                aug_config.geometric.scale_max = 1.0
-                aug_config.geometric.translate_height = 0
-                aug_config.geometric.translate_width = 0
+                # Use new independent rotation probability
+                aug_config.geometric.rotation_p = 1.0
             elif g == "scale":
+                # Use new independent scale probability
                 aug_config.geometric.scale_min = 0.9
                 aug_config.geometric.scale_max = 1.1
-                aug_config.geometric.affine_p = 1.0
-                aug_config.geometric.rotation_min = 0
-                aug_config.geometric.rotation_max = 0
-                aug_config.geometric.translate_height = 0
-                aug_config.geometric.translate_width = 0
+                aug_config.geometric.scale_p = 1.0
             elif g == "translate":
+                # Use new independent translate probability
                 aug_config.geometric.translate_height = 0.2
                 aug_config.geometric.translate_width = 0.2
-                aug_config.geometric.affine_p = 1.0
-                aug_config.geometric.rotation_min = 0
-                aug_config.geometric.rotation_max = 0
-                aug_config.geometric.scale_min = 1.0
-                aug_config.geometric.scale_max = 1.0
+                aug_config.geometric.translate_p = 1.0
             elif g == "erase_scale":
                 aug_config.geometric.erase_p = 1.0
             elif g == "mixup":
@@ -456,7 +447,8 @@ def get_data_config(
     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",
@@ -470,9 +462,10 @@ def get_data_config(
     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",
 ):
     """Train a pose-estimation model with SLEAP-NN framework.
 
@@ -486,7 +479,11 @@ def get_data_config(
             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".
@@ -508,16 +505,19 @@ def get_data_config(
         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`.
@@ -529,7 +529,8 @@ def get_data_config(
             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
                     }
     """
@@ -541,6 +542,7 @@ def get_data_config(
         scale=scale,
         crop_size=crop_size,
         min_crop_size=min_crop_size,
+        crop_padding=crop_padding,
     )
     augmentation_config = None
     if use_augmentations_train:
@@ -553,6 +555,7 @@ def get_data_config(
         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,
@@ -675,6 +678,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,
@@ -744,6 +748,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.
@@ -844,6 +851,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

@@ -84,6 +84,16 @@ class WandBConfig:
         prv_runid: (str) Previous run ID if training should be resumed from a previous ckpt. *Default*: `None`.
         group: (str) Group for wandb logging. *Default*: `None`.
         current_run_id: (str) Run ID for the current model training. (stored once the training starts). *Default*: `None`.
+        viz_enabled: (bool) If True, log pre-rendered matplotlib images to wandb. *Default*: `True`.
+        viz_boxes: (bool) If True, log interactive keypoint boxes. *Default*: `False`.
+        viz_masks: (bool) If True, log confidence map overlay masks. *Default*: `False`.
+        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
@@ -95,6 +105,13 @@ class WandBConfig:
     prv_runid: Optional[str] = None
     group: Optional[str] = None
     current_run_id: Optional[str] = None
+    viz_enabled: bool = True
+    viz_boxes: bool = False
+    viz_masks: bool = False
+    viz_box_size: float = 5.0
+    viz_confmap_threshold: float = 0.1
+    log_viz_table: bool = False
+    delete_local_logs: Optional[bool] = None
 
 
 @define
@@ -161,19 +178,69 @@ class ReduceLROnPlateauConfig:
         raise ValueError(message)
 
 
+@define
+class CosineAnnealingWarmupConfig:
+    """Configuration for Cosine Annealing with Linear Warmup scheduler.
+
+    The learning rate increases linearly during warmup, then decreases following
+    a cosine curve to the minimum value.
+
+    Attributes:
+        warmup_epochs: (int) Number of epochs for linear warmup phase. *Default*: `5`.
+        max_epochs: (int) Total number of training epochs. Will be overridden by
+            trainer's max_epochs if not specified. *Default*: `None`.
+        warmup_start_lr: (float) Learning rate at start of warmup. *Default*: `0.0`.
+        eta_min: (float) Minimum learning rate at end of cosine decay. *Default*: `0.0`.
+    """
+
+    warmup_epochs: int = field(default=5, validator=validators.ge(0))
+    max_epochs: Optional[int] = None
+    warmup_start_lr: float = field(default=0.0, validator=validators.ge(0))
+    eta_min: float = field(default=0.0, validator=validators.ge(0))
+
+
+@define
+class LinearWarmupLinearDecayConfig:
+    """Configuration for Linear Warmup + Linear Decay scheduler.
+
+    The learning rate increases linearly during warmup, then decreases linearly
+    to the end learning rate.
+
+    Attributes:
+        warmup_epochs: (int) Number of epochs for linear warmup phase. *Default*: `5`.
+        max_epochs: (int) Total number of training epochs. Will be overridden by
+            trainer's max_epochs if not specified. *Default*: `None`.
+        warmup_start_lr: (float) Learning rate at start of warmup. *Default*: `0.0`.
+        end_lr: (float) Learning rate at end of training. *Default*: `0.0`.
+    """
+
+    warmup_epochs: int = field(default=5, validator=validators.ge(0))
+    max_epochs: Optional[int] = None
+    warmup_start_lr: float = field(default=0.0, validator=validators.ge(0))
+    end_lr: float = field(default=0.0, validator=validators.ge(0))
+
+
 @define
 class LRSchedulerConfig:
     """Configuration for lr_scheduler.
 
+    Only one scheduler should be configured at a time. If multiple are set,
+    priority order is: cosine_annealing_warmup > linear_warmup_linear_decay >
+    step_lr > reduce_lr_on_plateau.
+
     Attributes:
         step_lr: Configuration for StepLR scheduler.
         reduce_lr_on_plateau: Configuration for ReduceLROnPlateau scheduler.
+        cosine_annealing_warmup: Configuration for Cosine Annealing with Linear Warmup scheduler.
+        linear_warmup_linear_decay: Configuration for Linear Warmup + Linear Decay scheduler.
     """
 
     step_lr: Optional[StepLRConfig] = None
     reduce_lr_on_plateau: Optional[ReduceLROnPlateauConfig] = field(
         factory=ReduceLROnPlateauConfig
     )
+    cosine_annealing_warmup: Optional[CosineAnnealingWarmupConfig] = None
+    linear_warmup_linear_decay: Optional[LinearWarmupLinearDecayConfig] = None
 
 
 @define
@@ -191,6 +258,26 @@ 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`.
+        match_threshold: (float) Maximum distance in pixels for centroid matching.
+            Only used for centroid model evaluation. *Default*: `50.0`.
+    """
+
+    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
+    match_threshold: float = field(default=50.0, validator=validators.gt(0))
+
+
 @define
 class HardKeypointMiningConfig:
     """Configuration for online hard keypoint mining.
@@ -293,6 +380,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):