Add Skia-based augmentation backend for faster data pipeline (#431)

## Summary
- **Replaces Kornia entirely** with Skia-based augmentation + pure
PyTorch operations
- Provides **~1.58x faster** CPU augmentation throughput
- Reduces GPU transfer bandwidth by **4x** (images stay uint8 until GPU
normalization)
- **Removes kornia as a dependency** (one fewer external dependency)

## Changes

### New Files
- **`sleap_nn/data/skia_augmentation.py`**: Skia implementations of
intensity/geometric augmentations and crop_and_resize

### Modified Files
- **`sleap_nn/data/augmentation.py`**: Simplified wrapper delegating to
Skia
- **`sleap_nn/data/custom_datasets.py`**: Use Skia crop_and_resize
- **`sleap_nn/data/instance_cropping.py`**: Use Skia crop_and_resize
- **`sleap_nn/inference/peak_finding.py`**: Pure PyTorch
morphological_dilation (replaces kornia)
- **`sleap_nn/training/lightning_modules.py`**: Add `normalize_on_gpu()`
to all training/validation steps
- **`pyproject.toml`**: Remove kornia, add skia-python>=87.0

## Architecture

```
Data Pipeline (CPU):
  Images stay uint8 → Skia augmentation → uint8 output

Training/Inference (GPU):
  uint8 tensor → normalize_on_gpu() → float32 [0,1] → model
```

This deferred normalization reduces PCIe bandwidth by 4x.

## Benchmark Results (from investigation)

| Dataset | Skia | Kornia | Speedup |
|---------|------|--------|---------|
| SingleInstance | 417.5/s | 359.8/s | **1.16x** |
| Centroid | 455.5/s | 365.3/s | **1.25x** |
| CenteredInstance | 998.0/s | 364.4/s | **2.74x** |
| BottomUp | 297.9/s | 253.6/s | **1.18x** |
| **Average** | | | **1.58x** |

## Test Results
- [x] All data tests pass (52/52)
- [x] All training tests pass (9/9)
- [x] Peak finding tests pass (9/9)
- [ ] End-to-end training verification

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.5 <[email protected]>

Author: gitttt-1234

.github/workflows/ci.yml

@@ -82,6 +82,11 @@ jobs:
         if: matrix.os != 'self-hosted-gpu'
         run: uv sync --extra torch-cpu
 
+      - name: Install graphics dependencies (Ubuntu)
+        if: matrix.os == 'ubuntu'
+        run: |
+          sudo apt-get update && sudo apt-get install -y libglapi-mesa libegl-mesa0 libegl1 libopengl0 libgl1 libglx-mesa0
+
       - name: Print environment info
         run: |
           echo "=== UV Environment ==="
@@ -103,7 +108,7 @@ jobs:
               print('CUDA is not available')
           " || echo "CUDA check failed"
           echo "=== Import Test ==="
-          uv run --frozen --extra torch-cpu python -c "import torch; import lightning; import kornia; print('All imports successful')" || echo "Import test failed"
+          uv run --frozen --extra torch-cpu python -c "import torch; import lightning; import skia; print('All imports successful')" || echo "Import test failed"
 
       - name: Check MPS backend (macOS only)
         if: runner.os == 'macOS'
@@ -126,7 +131,7 @@ jobs:
       - name: Run pytest
         run: |
           echo "=== Final environment check before tests ==="
-          uv run --frozen --extra torch-cpu python -c "import numpy, torch, lightning, kornia; print(f'All packages available: numpy={numpy.__version__}, torch={torch.__version__}')"
+          uv run --frozen --extra torch-cpu python -c "import numpy, torch, lightning, skia; print(f'All packages available: numpy={numpy.__version__}, torch={torch.__version__}')"
           echo "=== Running pytest ==="
           uv run --frozen --extra torch-cpu pytest --cov=sleap_nn --cov-report=xml --durations=-1 tests/
 

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "sleap-io>=0.6.2,<0.7.0",
     "numpy",
     "lightning",
-    "kornia",
+    "skia-python>=87.0",
     "jsonpickle",
     "scipy",
     "attrs",

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,66 +49,23 @@ 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,
@@ -128,8 +88,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]:
-    """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)
@@ -160,176 +120,25 @@ def apply_geometric_augmentation(
     Returns:
         Returns tuple: (image, instances) with augmentation applied.
     """
-    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
+    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,
     )
-
-    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

sleap_nn/data/custom_datasets.py

@@ -1,6 +1,6 @@
 """Custom `torch.utils.data.Dataset`s for different model types."""
 
-from kornia.geometry.transform import crop_and_resize
+from sleap_nn.data.skia_augmentation import crop_and_resize_skia as crop_and_resize
 
 # from concurrent.futures import ThreadPoolExecutor # TODO: implement parallel processing
 # import concurrent.futures

sleap_nn/data/instance_cropping.py

@@ -5,7 +5,7 @@
 import numpy as np
 import sleap_io as sio
 import torch
-from kornia.geometry.transform import crop_and_resize
+from sleap_nn.data.skia_augmentation import crop_and_resize_skia as crop_and_resize
 
 
 def compute_augmentation_padding(