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

sleap_nn/config/data_config.py

@@ -109,10 +109,10 @@ 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`.
+        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*: `None`.
         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`.
+        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*: `None`.
         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`.
         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`.
@@ -129,10 +129,10 @@ 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)
+    rotation_p: Optional[float] = field(default=None)
     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)
+    scale_p: Optional[float] = field(default=None)
     translate_width: float = 0.0
     translate_height: float = 0.0
     translate_p: Optional[float] = field(default=None)
@@ -198,8 +198,6 @@ class DataConfig:
         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`)
@@ -219,13 +217,9 @@ class DataConfig:
     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] = field(
-        factory=lambda: AugmentationConfig(geometric=GeometricConfig())
-    )
+    augmentation_config: Optional[AugmentationConfig] = None
     skeletons: Optional[list] = None
 
 

sleap_nn/config/get_config.py

@@ -463,9 +463,9 @@ def get_data_config(
     crop_size: Optional[int] = None,
     min_crop_size: Optional[int] = 100,
     crop_padding: Optional[int] = None,
-    use_augmentations_train: bool = True,
+    use_augmentations_train: bool = False,
     intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
-    geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = "rotation",
+    geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
 ):
     """Train a pose-estimation model with SLEAP-NN framework.
 
@@ -517,7 +517,7 @@ def get_data_config(
             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: True.
+            training data, else False. Default: False.
         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,8 +529,7 @@ 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_min": -45,
-                        "rotation_max": 45,
+                        "rotation": 45,
                         "affine_p": 1.0
                     }
     """

sleap_nn/config/trainer_config.py

@@ -178,69 +178,19 @@ 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
@@ -258,26 +208,6 @@ 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.
@@ -380,7 +310,6 @@ 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/data/augmentation.py

@@ -1,15 +1,12 @@
-"""This module implements data pipeline blocks for augmentation operations.
+"""This module implements data pipeline blocks for augmentation operations."""
 
-Uses Skia (skia-python) for ~1.5x faster augmentation compared to Kornia.
-"""
-
-from typing import Optional, Tuple
+from typing import Any, Dict, Optional, Tuple, Union
+import kornia as K
 import torch
-
-from sleap_nn.data.skia_augmentation import (
-    apply_intensity_augmentation_skia,
-    apply_geometric_augmentation_skia,
-)
+from kornia.augmentation._2d.intensity.base import IntensityAugmentationBase2D
+from kornia.augmentation.container import AugmentationSequential
+from kornia.augmentation.utils.param_validation import _range_bound
+from kornia.core import Tensor
 
 
 def apply_intensity_augmentation(
@@ -27,8 +24,8 @@ def apply_intensity_augmentation(
     brightness_min: Optional[float] = 1.0,
     brightness_max: Optional[float] = 1.0,
     brightness_p: float = 0.0,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Apply intensity augmentation on image and instances.
+) -> Tuple[torch.Tensor]:
+    """Apply kornia intensity augmentation on image and instances.
 
     Args:
         image: Input image. Shape: (n_samples, C, H, W)
@@ -49,23 +46,66 @@ def apply_intensity_augmentation(
     Returns:
         Returns tuple: (image, instances) with augmentation applied.
     """
-    return apply_intensity_augmentation_skia(
-        image=image,
-        instances=instances,
-        uniform_noise_min=uniform_noise_min,
-        uniform_noise_max=uniform_noise_max,
-        uniform_noise_p=uniform_noise_p,
-        gaussian_noise_mean=gaussian_noise_mean,
-        gaussian_noise_std=gaussian_noise_std,
-        gaussian_noise_p=gaussian_noise_p,
-        contrast_min=contrast_min,
-        contrast_max=contrast_max,
-        contrast_p=contrast_p,
-        brightness_min=brightness_min,
-        brightness_max=brightness_max,
-        brightness_p=brightness_p,
+    aug_stack = []
+    if uniform_noise_p > 0:
+        aug_stack.append(
+            RandomUniformNoise(
+                noise=(uniform_noise_min, uniform_noise_max),
+                p=uniform_noise_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if gaussian_noise_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomGaussianNoise(
+                mean=gaussian_noise_mean,
+                std=gaussian_noise_std,
+                p=gaussian_noise_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if contrast_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomContrast(
+                contrast=(contrast_min, contrast_max),
+                p=contrast_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if brightness_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomBrightness(
+                brightness=(brightness_min, brightness_max),
+                p=brightness_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+
+    augmenter = AugmentationSequential(
+        *aug_stack,
+        data_keys=["input", "keypoints"],
+        keepdim=True,
+        same_on_batch=True,
     )
 
+    inst_shape = instances.shape
+    # Before (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # Before (cropped image): (B=1, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    instances = instances.reshape(inst_shape[0], -1, 2)
+    # (n_samples, C, H, W), (n_samples, n_instances * n_nodes, 2) OR (n_samples, n_nodes, 2)
+
+    aug_image, aug_instances = augmenter(image, instances)
+
+    # After (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # After (cropped image): (n_samples, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    return aug_image, aug_instances.reshape(*inst_shape)
+
 
 def apply_geometric_augmentation(
     image: torch.Tensor,
@@ -88,8 +128,8 @@ def apply_geometric_augmentation(
     mixup_lambda_min: Optional[float] = 0.01,
     mixup_lambda_max: Optional[float] = 0.05,
     mixup_p: float = 0.0,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Apply geometric augmentation on image and instances.
+) -> Tuple[torch.Tensor]:
+    """Apply kornia geometric augmentation on image and instances.
 
     Args:
         image: Input image. Shape: (n_samples, C, H, W)
@@ -120,25 +160,176 @@ def apply_geometric_augmentation(
     Returns:
         Returns tuple: (image, instances) with augmentation applied.
     """
-    return apply_geometric_augmentation_skia(
-        image=image,
-        instances=instances,
-        rotation_min=rotation_min,
-        rotation_max=rotation_max,
-        rotation_p=rotation_p,
-        scale_min=scale_min,
-        scale_max=scale_max,
-        scale_p=scale_p,
-        translate_width=translate_width,
-        translate_height=translate_height,
-        translate_p=translate_p,
-        affine_p=affine_p,
-        erase_scale_min=erase_scale_min,
-        erase_scale_max=erase_scale_max,
-        erase_ratio_min=erase_ratio_min,
-        erase_ratio_max=erase_ratio_max,
-        erase_p=erase_p,
-        mixup_lambda_min=mixup_lambda_min,
-        mixup_lambda_max=mixup_lambda_max,
-        mixup_p=mixup_p,
+    aug_stack = []
+
+    # Check if any individual probability is set
+    use_independent = (
+        rotation_p is not None or scale_p is not None or translate_p is not None
     )
+
+    if use_independent:
+        # New behavior: Apply augmentations independently with separate probabilities
+        if rotation_p is not None and rotation_p > 0:
+            aug_stack.append(
+                K.augmentation.RandomRotation(
+                    degrees=(rotation_min, rotation_max),
+                    p=rotation_p,
+                    keepdim=True,
+                    same_on_batch=True,
+                )
+            )
+
+        if scale_p is not None and scale_p > 0:
+            aug_stack.append(
+                K.augmentation.RandomAffine(
+                    degrees=0,  # No rotation
+                    translate=None,  # No translation
+                    scale=(scale_min, scale_max),
+                    p=scale_p,
+                    keepdim=True,
+                    same_on_batch=True,
+                )
+            )
+
+        if translate_p is not None and translate_p > 0:
+            aug_stack.append(
+                K.augmentation.RandomAffine(
+                    degrees=0,  # No rotation
+                    translate=(translate_width, translate_height),
+                    scale=None,  # No scaling
+                    p=translate_p,
+                    keepdim=True,
+                    same_on_batch=True,
+                )
+            )
+    elif affine_p > 0:
+        # Legacy behavior: Bundled affine transformation
+        aug_stack.append(
+            K.augmentation.RandomAffine(
+                degrees=(rotation_min, rotation_max),
+                translate=(translate_width, translate_height),
+                scale=(scale_min, scale_max),
+                p=affine_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+
+    if erase_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomErasing(
+                scale=(erase_scale_min, erase_scale_max),
+                ratio=(erase_ratio_min, erase_ratio_max),
+                p=erase_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+    if mixup_p > 0:
+        aug_stack.append(
+            K.augmentation.RandomMixUpV2(
+                lambda_val=(mixup_lambda_min, mixup_lambda_max),
+                p=mixup_p,
+                keepdim=True,
+                same_on_batch=True,
+            )
+        )
+
+    augmenter = AugmentationSequential(
+        *aug_stack,
+        data_keys=["input", "keypoints"],
+        keepdim=True,
+        same_on_batch=True,
+    )
+
+    inst_shape = instances.shape
+    # Before (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # Before (cropped image): (B=1, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    instances = instances.reshape(inst_shape[0], -1, 2)
+    # (n_samples, C, H, W), (n_samples, n_instances * n_nodes, 2) OR (n_samples, n_nodes, 2)
+
+    aug_image, aug_instances = augmenter(image, instances)
+
+    # After (full image): (n_samples, C, H, W), (n_samples, n_instances, n_nodes, 2)
+    # or
+    # After (cropped image): (n_samples, C, crop_H, crop_W), (n_samples, n_nodes, 2)
+    return aug_image, aug_instances.reshape(*inst_shape)
+
+
+class RandomUniformNoise(IntensityAugmentationBase2D):
+    """Data transformer for applying random uniform noise to input images.
+
+    This is a custom Kornia augmentation inheriting from `IntensityAugmentationBase2D`.
+    Uniform noise within (min_val, max_val) is applied to the entire input image.
+
+    Note: Inverse transform is not implemented and re-applying the same transformation
+    in the example below does not work when included in an AugmentationSequential class.
+
+    Args:
+        noise: 2-tuple (min_val, max_val); 0.0 <= min_val <= max_val <= 1.0.
+        p: probability for applying an augmentation. This param controls the augmentation probabilities
+          element-wise for a batch.
+        p_batch: probability for applying an augmentation to a batch. This param controls the augmentation
+          probabilities batch-wise.
+        same_on_batch: apply the same transformation across the batch.
+        keepdim: whether to keep the output shape the same as input `True` or broadcast it
+          to the batch form `False`.
+
+    Examples:
+        >>> rng = torch.manual_seed(0)
+        >>> img = torch.rand(1, 1, 2, 2)
+        >>> RandomUniformNoise(min_val=0., max_val=0.1, p=1.)(img)
+        tensor([[[[0.9607, 0.5865],
+                  [0.2705, 0.5920]]]])
+
+    To apply the exact augmentation again, you may take the advantage of the previous parameter state:
+        >>> input = torch.rand(1, 3, 32, 32)
+        >>> aug = RandomUniformNoise(min_val=0., max_val=0.1, p=1.)
+        >>> (aug(input) == aug(input, params=aug._params)).all()
+        tensor(True)
+
+    Ref: `kornia.augmentation._2d.intensity.gaussian_noise
+    <https://kornia.readthedocs.io/en/latest/_modules/kornia/augmentation/_2d/intensity/gaussian_noise.html#RandomGaussianNoise>`_.
+    """
+
+    def __init__(
+        self,
+        noise: Tuple[float, float],
+        p: float = 0.5,
+        p_batch: float = 1.0,
+        clip_output: bool = True,
+        same_on_batch: bool = False,
+        keepdim: bool = False,
+    ) -> None:
+        """Initialize the class."""
+        super().__init__(
+            p=p, p_batch=p_batch, same_on_batch=same_on_batch, keepdim=keepdim
+        )
+        self.flags = {
+            "uniform_noise": _range_bound(noise, "uniform_noise", bounds=(0.0, 1.0))
+        }
+        self.clip_output = clip_output
+
+    def apply_transform(
+        self,
+        input: Tensor,
+        params: Dict[str, Tensor],
+        flags: Dict[str, Any],
+        transform: Optional[Tensor] = None,
+    ) -> Tensor:
+        """Compute the uniform noise, add, and clamp output."""
+        if "uniform_noise" in params:
+            uniform_noise = params["uniform_noise"]
+        else:
+            uniform_noise = (
+                torch.FloatTensor(input.shape)
+                .uniform_(flags["uniform_noise"][0], flags["uniform_noise"][1])
+                .to(input.device)
+            )
+            self._params["uniform_noise"] = uniform_noise
+        if self.clip_output:
+            return torch.clamp(
+                input + uniform_noise, 0.0, 1.0
+            )  # RandomGaussianNoise doesn't clamp.
+        return input + uniform_noise