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

sleap_nn/data/augmentation.py

@@ -1,12 +1,15 @@
-"""This module implements data pipeline blocks for augmentation operations."""
+"""This module implements data pipeline blocks for augmentation operations.
 
-from typing import Any, Dict, Optional, Tuple, Union
-import kornia as K
+Uses Skia (skia-python) for ~1.5x faster augmentation compared to Kornia.
+"""
+
+from typing import Optional, Tuple
 import torch
-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
+
+from sleap_nn.data.skia_augmentation import (
+    apply_intensity_augmentation_skia,
+    apply_geometric_augmentation_skia,
+)
 
 
 def apply_intensity_augmentation(
@@ -24,8 +27,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]:
-    """Apply kornia intensity augmentation on image and instances.
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply intensity augmentation on image and instances.
 
     Args:
         image: Input image. Shape: (n_samples, C, H, W)
@@ -46,76 +49,36 @@ def apply_intensity_augmentation(
     Returns:
         Returns tuple: (image, instances) with augmentation applied.
     """
-    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,
+    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,
     )
 
-    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,
     instances: torch.Tensor,
     rotation_min: Optional[float] = -15.0,
     rotation_max: Optional[float] = 15.0,
+    rotation_p: Optional[float] = None,
     scale_min: Optional[float] = 0.9,
     scale_max: Optional[float] = 1.1,
+    scale_p: Optional[float] = None,
     translate_width: Optional[float] = 0.02,
     translate_height: Optional[float] = 0.02,
+    translate_p: Optional[float] = None,
     affine_p: float = 0.0,
     erase_scale_min: Optional[float] = 0.0001,
     erase_scale_max: Optional[float] = 0.01,
@@ -125,19 +88,26 @@ 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]:
-    """Apply kornia geometric augmentation on image and instances.
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Apply geometric augmentation on image and instances.
 
     Args:
         image: Input image. Shape: (n_samples, C, H, W)
         instances: Input keypoints. (n_samples, n_instances, n_nodes, 2) or (n_samples, n_nodes, 2)
         rotation_min: Minimum rotation angle in degrees. Default: -15.0.
         rotation_max: Maximum rotation angle in degrees. Default: 15.0.
+        rotation_p: Probability of applying random rotation independently. If None,
+            falls back to affine_p for bundled behavior. Default: None.
         scale_min: Minimum scaling factor for isotropic scaling. Default: 0.9.
         scale_max: Maximum scaling factor for isotropic scaling. Default: 1.1.
+        scale_p: Probability of applying random scaling independently. If None,
+            falls back to affine_p for bundled behavior. Default: None.
         translate_width: Maximum absolute fraction for horizontal translation. Default: 0.02.
         translate_height: Maximum absolute fraction for vertical translation. Default: 0.02.
-        affine_p: Probability of applying random affine transformations. Default: 0.0.
+        translate_p: Probability of applying random translation independently. If None,
+            falls back to affine_p for bundled behavior. Default: None.
+        affine_p: Probability of applying random affine transformations (rotation, scale,
+            translate bundled). Used when individual *_p params are None. Default: 0.0.
         erase_scale_min: Minimum value of range of proportion of erased area against input image. Default: 0.0001.
         erase_scale_max: Maximum value of range of proportion of erased area against input image. Default: 0.01.
         erase_ratio_min: Minimum value of range of aspect ratio of erased area. Default: 1.
@@ -150,134 +120,25 @@ def apply_geometric_augmentation(
     Returns:
         Returns tuple: (image, instances) with augmentation applied.
     """
-    aug_stack = []
-    if affine_p > 0:
-        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,
+    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,
     )
-
-    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