Skip to content

skia_augmentation

sleap_nn.data.skia_augmentation

Skia-based augmentation functions that operate on uint8 tensors.

This module provides augmentation functions using skia-python that: 1. Match the exact API of sleap_nn.data.augmentation 2. Operate on uint8 tensors throughout (avoiding float32 conversions) 3. Provide ~1.5x faster augmentation compared to Kornia

The implementation uses array-backed Skia surfaces (following the sleap-io pattern) which avoids platform-specific BGR/RGBA pixel format issues. By creating surfaces with an existing numpy array as the backing store, Skia writes directly to that array in the specified color format, eliminating the need for channel swapping.

Usage

from sleap_nn.data.skia_augmentation import ( apply_intensity_augmentation_skia, apply_geometric_augmentation_skia, )

Apply augmentations (uint8 in, uint8 out)

image, instances = apply_intensity_augmentation_skia(image, instances, **config) image, instances = apply_geometric_augmentation_skia(image, instances, **config)

Functions:

Name Description
apply_flip_augmentation_skia

Randomly mirror an image and its keypoints left/right, swapping symmetries.

apply_geometric_augmentation_skia

Apply geometric augmentations using Skia.

apply_intensity_augmentation_skia

Apply intensity augmentations on uint8 image tensor.

crop_and_resize_skia

Crop and resize image regions using Skia.

apply_flip_augmentation_skia(image, instances, symmetric_inds=None, flip_p=0.0, masks=None)

Randomly mirror an image and its keypoints left/right, swapping symmetries.

When an image is mirrored left/right, left/right symmetric body parts physically exchange sides, so their slots in the instance array must be swapped to keep semantic labels correct (e.g. left_paw must remain left_paw). This mirrors the behavior of SLEAP v1.4's RandomFlipper (horizontal flip).

The flip is applied to the whole sample with probability flip_p (sampled once per call). Mirroring is exact and lossless (a tensor reverse, no interpolation).

Parameters:

Name Type Description Default
image Tensor

Input tensor of shape (1, C, H, W) with dtype uint8 or float32.

required
instances Tensor

Keypoints tensor of shape (1, n_instances, n_nodes, 2) or (1, n_nodes, 2). The node axis is the second-to-last (-2).

required
symmetric_inds Optional[Sequence[Tuple[int, int]]]

Iterable of (i, j) node-index pairs to swap after mirroring. None or empty means no swap (correct only if the skeleton is truly left/right symmetric in labeling, e.g. centroids).

None
flip_p float

Probability of applying the flip. 0 disables (no-op).

0.0
masks Optional[Tensor]

Optional segmentation masks of shape (1, K, H, W) to co-transform with the image under the SAME flip draw. Mirroring a binary mask is exact and lossless (a tensor reverse). Symmetric-node swapping does not apply to masks (a mask just mirrors). When provided, the return value gains a third element holding the mirrored masks.

None

Returns:

Type Description
Tuple[Tensor, ...]

(image, instances) when masks is None, else (image, instances, masks). When the flip is not applied, the inputs are returned unchanged. NaN keypoints are preserved ((W - 1) - NaN is NaN).

Source code in sleap_nn/data/skia_augmentation.py
def apply_flip_augmentation_skia(
    image: torch.Tensor,
    instances: torch.Tensor,
    symmetric_inds: Optional[Sequence[Tuple[int, int]]] = None,
    flip_p: float = 0.0,
    masks: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, ...]:
    """Randomly mirror an image and its keypoints left/right, swapping symmetries.

    When an image is mirrored left/right, left/right symmetric body parts physically
    exchange sides, so their slots in the instance array must be swapped to keep
    semantic labels correct (e.g. ``left_paw`` must remain ``left_paw``). This
    mirrors the behavior of SLEAP v1.4's ``RandomFlipper`` (horizontal flip).

    The flip is applied to the whole sample with probability ``flip_p`` (sampled once
    per call). Mirroring is exact and lossless (a tensor reverse, no interpolation).

    Args:
        image: Input tensor of shape ``(1, C, H, W)`` with dtype uint8 or float32.
        instances: Keypoints tensor of shape ``(1, n_instances, n_nodes, 2)`` or
            ``(1, n_nodes, 2)``. The node axis is the second-to-last (``-2``).
        symmetric_inds: Iterable of ``(i, j)`` node-index pairs to swap after
            mirroring. ``None`` or empty means no swap (correct only if the skeleton
            is truly left/right symmetric in labeling, e.g. centroids).
        flip_p: Probability of applying the flip. ``0`` disables (no-op).
        masks: Optional segmentation masks of shape ``(1, K, H, W)`` to co-transform
            with the image under the SAME flip draw. Mirroring a binary mask is exact
            and lossless (a tensor reverse). Symmetric-node swapping does not apply to
            masks (a mask just mirrors). When provided, the return value gains a third
            element holding the mirrored masks.

    Returns:
        ``(image, instances)`` when ``masks`` is ``None``, else
        ``(image, instances, masks)``. When the flip is not applied, the inputs are
        returned unchanged. NaN keypoints are preserved (``(W - 1) - NaN`` is ``NaN``).
    """
    if flip_p <= 0 or np.random.random() >= flip_p:
        if masks is not None:
            return image, instances, masks
        return image, instances

    instances = instances.clone()
    width = image.shape[-1]
    image = torch.flip(image, dims=[-1])
    instances[..., 0] = (width - 1) - instances[..., 0]

    # Swap symmetric node pairs on the node axis (-2) so labels stay correct.
    if symmetric_inds is not None:
        for a, b in symmetric_inds:
            swap = instances[..., a, :].clone()
            instances[..., a, :] = instances[..., b, :]
            instances[..., b, :] = swap

    if masks is not None:
        masks = torch.flip(masks, dims=[-1])
        return image, instances, masks

    return image, instances

apply_geometric_augmentation_skia(image, instances, rotation_min=-15.0, rotation_max=15.0, rotation_p=None, scale_min=0.9, scale_max=1.1, scale_p=None, translate_width=0.02, translate_height=0.02, translate_p=None, affine_p=0.0, erase_scale_min=0.0001, erase_scale_max=0.01, erase_ratio_min=1.0, erase_ratio_max=1.0, erase_p=0.0, mixup_lambda_min=0.01, mixup_lambda_max=0.05, mixup_p=0.0, flip_p=0.0, symmetric_inds=None, masks=None)

Apply geometric augmentations using Skia.

Matches API of sleap_nn.data.augmentation.apply_geometric_augmentation.

Parameters:

Name Type Description Default
image Tensor

Input tensor of shape (1, C, H, W) with dtype uint8 or float32.

required
instances Tensor

Keypoints tensor of shape (1, n_instances, n_nodes, 2) or (1, n_nodes, 2).

required
rotation_min float

Minimum rotation angle in degrees.

-15.0
rotation_max float

Maximum rotation angle in degrees.

15.0
rotation_p Optional[float]

Probability of rotation (independent). None = use affine_p.

None
scale_min float

Minimum scale factor.

0.9
scale_max float

Maximum scale factor.

1.1
scale_p Optional[float]

Probability of scaling (independent). None = use affine_p.

None
translate_width float

Max horizontal translation as fraction of width.

0.02
translate_height float

Max vertical translation as fraction of height.

0.02
translate_p Optional[float]

Probability of translation (independent). None = use affine_p.

None
affine_p float

Probability of bundled affine transform.

0.0
erase_scale_min float

Min proportion of image to erase.

0.0001
erase_scale_max float

Max proportion of image to erase.

0.01
erase_ratio_min float

Min aspect ratio of erased area.

1.0
erase_ratio_max float

Max aspect ratio of erased area.

1.0
erase_p float

Probability of random erasing.

0.0
mixup_lambda_min float

Min mixup strength (not implemented).

0.01
mixup_lambda_max float

Max mixup strength (not implemented).

0.05
mixup_p float

Probability of mixup (not implemented).

0.0
flip_p float

Probability of mirroring the sample left/right (with symmetric-node swap).

0.0
symmetric_inds Optional[Sequence[Tuple[int, int]]]

Node-index pairs to swap after mirroring (see apply_flip_augmentation_skia). None/empty means no swap.

None
masks Optional[Tensor]

Optional segmentation masks of shape (1, K, H, W) (float in {0, 1}) to co-transform with the image under the SAME sampled flip + affine matrix. Masks are warped with nearest-neighbor sampling (crisp binary, no interpolation fringe) and re-binarized. Erase/mixup are image-only and never touch masks (they simulate occlusion/blending while the object — and thus its ground-truth mask — remains). When provided, the return value gains a third element holding the co-transformed masks.

None

Returns:

Type Description
Tuple[Tensor, ...]

(augmented_image, augmented_instances) when masks is None, else (augmented_image, augmented_instances, augmented_masks). Image dtype matches input.

Source code in sleap_nn/data/skia_augmentation.py
def apply_geometric_augmentation_skia(
    image: torch.Tensor,
    instances: torch.Tensor,
    rotation_min: float = -15.0,
    rotation_max: float = 15.0,
    rotation_p: Optional[float] = None,
    scale_min: float = 0.9,
    scale_max: float = 1.1,
    scale_p: Optional[float] = None,
    translate_width: float = 0.02,
    translate_height: float = 0.02,
    translate_p: Optional[float] = None,
    affine_p: float = 0.0,
    erase_scale_min: float = 0.0001,
    erase_scale_max: float = 0.01,
    erase_ratio_min: float = 1.0,
    erase_ratio_max: float = 1.0,
    erase_p: float = 0.0,
    mixup_lambda_min: float = 0.01,
    mixup_lambda_max: float = 0.05,
    mixup_p: float = 0.0,
    flip_p: float = 0.0,
    symmetric_inds: Optional[Sequence[Tuple[int, int]]] = None,
    masks: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, ...]:
    """Apply geometric augmentations using Skia.

    Matches API of sleap_nn.data.augmentation.apply_geometric_augmentation.

    Args:
        image: Input tensor of shape (1, C, H, W) with dtype uint8 or float32.
        instances: Keypoints tensor of shape (1, n_instances, n_nodes, 2) or (1, n_nodes, 2).
        rotation_min: Minimum rotation angle in degrees.
        rotation_max: Maximum rotation angle in degrees.
        rotation_p: Probability of rotation (independent). None = use affine_p.
        scale_min: Minimum scale factor.
        scale_max: Maximum scale factor.
        scale_p: Probability of scaling (independent). None = use affine_p.
        translate_width: Max horizontal translation as fraction of width.
        translate_height: Max vertical translation as fraction of height.
        translate_p: Probability of translation (independent). None = use affine_p.
        affine_p: Probability of bundled affine transform.
        erase_scale_min: Min proportion of image to erase.
        erase_scale_max: Max proportion of image to erase.
        erase_ratio_min: Min aspect ratio of erased area.
        erase_ratio_max: Max aspect ratio of erased area.
        erase_p: Probability of random erasing.
        mixup_lambda_min: Min mixup strength (not implemented).
        mixup_lambda_max: Max mixup strength (not implemented).
        mixup_p: Probability of mixup (not implemented).
        flip_p: Probability of mirroring the sample left/right (with symmetric-node
            swap).
        symmetric_inds: Node-index pairs to swap after mirroring (see
            ``apply_flip_augmentation_skia``). None/empty means no swap.
        masks: Optional segmentation masks of shape ``(1, K, H, W)`` (float in
            ``{0, 1}``) to co-transform with the image under the SAME sampled flip +
            affine ``matrix``. Masks are warped with nearest-neighbor sampling (crisp
            binary, no interpolation fringe) and re-binarized. Erase/mixup are
            image-only and never touch masks (they simulate occlusion/blending while
            the object — and thus its ground-truth mask — remains). When provided, the
            return value gains a third element holding the co-transformed masks.

    Returns:
        ``(augmented_image, augmented_instances)`` when ``masks`` is ``None``, else
        ``(augmented_image, augmented_instances, augmented_masks)``. Image dtype
        matches input.
    """
    # Apply flip first (matches SLEAP v1.4 ordering: flip before affine).
    if flip_p > 0:
        if masks is not None:
            image, instances, masks = apply_flip_augmentation_skia(
                image,
                instances,
                symmetric_inds=symmetric_inds,
                flip_p=flip_p,
                masks=masks,
            )
        else:
            image, instances = apply_flip_augmentation_skia(
                image,
                instances,
                symmetric_inds=symmetric_inds,
                flip_p=flip_p,
            )

    # Convert to numpy for Skia processing
    is_float = image.dtype == torch.float32
    if is_float:
        img_np = (image[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    else:
        img_np = image[0].permute(1, 2, 0).numpy().copy()

    h, w = img_np.shape[:2]
    cx, cy = w / 2, h / 2

    # Build transformation matrix
    matrix = skia.Matrix()
    has_transform = False

    use_independent = (
        rotation_p is not None or scale_p is not None or translate_p is not None
    )

    if use_independent:
        if (
            rotation_p is not None
            and rotation_p > 0
            and np.random.random() < rotation_p
        ):
            angle = np.random.uniform(rotation_min, rotation_max)
            rot_matrix = skia.Matrix()
            rot_matrix.setRotate(angle, cx, cy)
            matrix = matrix.preConcat(rot_matrix)
            has_transform = True

        if scale_p is not None and scale_p > 0 and np.random.random() < scale_p:
            scale = np.random.uniform(scale_min, scale_max)
            scale_matrix = skia.Matrix()
            scale_matrix.setScale(scale, scale, cx, cy)
            matrix = matrix.preConcat(scale_matrix)
            has_transform = True

        if (
            translate_p is not None
            and translate_p > 0
            and np.random.random() < translate_p
        ):
            tx = np.random.uniform(-translate_width, translate_width) * w
            ty = np.random.uniform(-translate_height, translate_height) * h
            trans_matrix = skia.Matrix()
            trans_matrix.setTranslate(tx, ty)
            matrix = matrix.preConcat(trans_matrix)
            has_transform = True

    elif affine_p > 0 and np.random.random() < affine_p:
        angle = np.random.uniform(rotation_min, rotation_max)
        scale = np.random.uniform(scale_min, scale_max)
        tx = np.random.uniform(-translate_width, translate_width) * w
        ty = np.random.uniform(-translate_height, translate_height) * h

        matrix.setRotate(angle, cx, cy)
        matrix.preScale(scale, scale, cx, cy)
        matrix.preTranslate(tx, ty)
        has_transform = True

    # Apply geometric transform
    if has_transform:
        img_np = _transform_image_skia(img_np, matrix)
        instances = _transform_keypoints_tensor(instances, matrix)
        if masks is not None:
            masks = _transform_masks_skia(masks, matrix)

    # Apply random erasing (image-only; masks are intentionally untouched because
    # erase simulates occlusion of an object that is still present in the GT).
    if erase_p > 0 and np.random.random() < erase_p:
        img_np = _apply_random_erase(
            img_np, erase_scale_min, erase_scale_max, erase_ratio_min, erase_ratio_max
        )

    # Convert back to tensor
    result_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0)
    if is_float:
        result_tensor = result_tensor.float() / 255.0

    if masks is not None:
        return result_tensor, instances, masks
    return result_tensor, instances

apply_intensity_augmentation_skia(image, instances, uniform_noise_min=0.0, uniform_noise_max=0.04, uniform_noise_p=0.0, gaussian_noise_mean=0.0, gaussian_noise_std=0.02, gaussian_noise_p=0.0, contrast_min=0.9, contrast_max=1.1, contrast_p=0.0, brightness_min=0.9, brightness_max=1.1, brightness_p=0.0)

Apply intensity augmentations on uint8 image tensor.

Matches API of sleap_nn.data.augmentation.apply_intensity_augmentation.

Parameters:

Name Type Description Default
image Tensor

Input tensor of shape (1, C, H, W) with dtype uint8 or float32.

required
instances Tensor

Keypoints tensor (not modified, just passed through).

required
uniform_noise_min float

Minimum uniform noise (0-1 scale, maps to 0-255).

0.0
uniform_noise_max float

Maximum uniform noise (0-1 scale).

0.04
uniform_noise_p float

Probability of uniform noise.

0.0
gaussian_noise_mean float

Gaussian noise mean (0-1 scale).

0.0
gaussian_noise_std float

Gaussian noise std (0-1 scale).

0.02
gaussian_noise_p float

Probability of Gaussian noise.

0.0
contrast_min float

Minimum contrast factor.

0.9
contrast_max float

Maximum contrast factor.

1.1
contrast_p float

Probability of contrast adjustment.

0.0
brightness_min float

Minimum brightness factor.

0.9
brightness_max float

Maximum brightness factor.

1.1
brightness_p float

Probability of brightness adjustment.

0.0

Returns:

Type Description
Tuple[Tensor, Tensor]

Tuple of (augmented_image, instances). Image dtype matches input.

Source code in sleap_nn/data/skia_augmentation.py
def apply_intensity_augmentation_skia(
    image: torch.Tensor,
    instances: torch.Tensor,
    uniform_noise_min: float = 0.0,
    uniform_noise_max: float = 0.04,
    uniform_noise_p: float = 0.0,
    gaussian_noise_mean: float = 0.0,
    gaussian_noise_std: float = 0.02,
    gaussian_noise_p: float = 0.0,
    contrast_min: float = 0.9,
    contrast_max: float = 1.1,
    contrast_p: float = 0.0,
    brightness_min: float = 0.9,
    brightness_max: float = 1.1,
    brightness_p: float = 0.0,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Apply intensity augmentations on uint8 image tensor.

    Matches API of sleap_nn.data.augmentation.apply_intensity_augmentation.

    Args:
        image: Input tensor of shape (1, C, H, W) with dtype uint8 or float32.
        instances: Keypoints tensor (not modified, just passed through).
        uniform_noise_min: Minimum uniform noise (0-1 scale, maps to 0-255).
        uniform_noise_max: Maximum uniform noise (0-1 scale).
        uniform_noise_p: Probability of uniform noise.
        gaussian_noise_mean: Gaussian noise mean (0-1 scale).
        gaussian_noise_std: Gaussian noise std (0-1 scale).
        gaussian_noise_p: Probability of Gaussian noise.
        contrast_min: Minimum contrast factor.
        contrast_max: Maximum contrast factor.
        contrast_p: Probability of contrast adjustment.
        brightness_min: Minimum brightness factor.
        brightness_max: Maximum brightness factor.
        brightness_p: Probability of brightness adjustment.

    Returns:
        Tuple of (augmented_image, instances). Image dtype matches input.
    """
    # Convert to numpy for Skia processing
    is_float = image.dtype == torch.float32
    if is_float:
        img_np = (image[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    else:
        img_np = image[0].permute(1, 2, 0).numpy()

    result = img_np.copy()

    # Apply uniform noise (in uint8 space)
    if uniform_noise_p > 0 and np.random.random() < uniform_noise_p:
        noise = np.random.randint(
            int(uniform_noise_min * 255),
            int(uniform_noise_max * 255) + 1,
            img_np.shape,
            dtype=np.int16,
        )
        result = np.clip(result.astype(np.int16) + noise, 0, 255).astype(np.uint8)

    # Apply Gaussian noise (in uint8 space)
    if gaussian_noise_p > 0 and np.random.random() < gaussian_noise_p:
        noise = np.random.normal(
            gaussian_noise_mean * 255, gaussian_noise_std * 255, img_np.shape
        ).astype(np.int16)
        result = np.clip(result.astype(np.int16) + noise, 0, 255).astype(np.uint8)

    # Apply contrast using lookup table (pure uint8)
    if contrast_p > 0 and np.random.random() < contrast_p:
        factor = np.random.uniform(contrast_min, contrast_max)
        lut = np.arange(256, dtype=np.float32)
        lut = np.clip((lut - 127.5) * factor + 127.5, 0, 255).astype(np.uint8)
        result = lut[result]

    # Apply brightness using lookup table (pure uint8)
    if brightness_p > 0 and np.random.random() < brightness_p:
        factor = np.random.uniform(brightness_min, brightness_max)
        lut = np.arange(256, dtype=np.float32)
        lut = np.clip(lut * factor, 0, 255).astype(np.uint8)
        result = lut[result]

    # Convert back to tensor
    result_tensor = torch.from_numpy(result).permute(2, 0, 1).unsqueeze(0)
    if is_float:
        result_tensor = result_tensor.float() / 255.0

    return result_tensor, instances

crop_and_resize_skia(image, boxes, size)

Crop and resize image regions using Skia.

Replacement for kornia.geometry.transform.crop_and_resize.

Uses array-backed Skia surface pattern from sleap-io to avoid platform-specific BGR/RGBA issues. By creating the surface with an existing numpy array as backing store, Skia writes directly to that array in RGBA format.

Parameters:

Name Type Description Default
image Tensor

Input tensor of shape (1, C, H, W).

required
boxes Tensor

Bounding boxes tensor of shape (1, 4, 2) with corners: [top-left, top-right, bottom-right, bottom-left].

required
size Tuple[int, int]

Output size (height, width).

required

Returns:

Type Description
Tensor

Cropped and resized tensor of shape (1, C, out_h, out_w).

Source code in sleap_nn/data/skia_augmentation.py
def crop_and_resize_skia(
    image: torch.Tensor,
    boxes: torch.Tensor,
    size: Tuple[int, int],
) -> torch.Tensor:
    """Crop and resize image regions using Skia.

    Replacement for kornia.geometry.transform.crop_and_resize.

    Uses array-backed Skia surface pattern from sleap-io to avoid platform-specific
    BGR/RGBA issues. By creating the surface with an existing numpy array as backing
    store, Skia writes directly to that array in RGBA format.

    Args:
        image: Input tensor of shape (1, C, H, W).
        boxes: Bounding boxes tensor of shape (1, 4, 2) with corners:
            [top-left, top-right, bottom-right, bottom-left].
        size: Output size (height, width).

    Returns:
        Cropped and resized tensor of shape (1, C, out_h, out_w).
    """
    is_float = image.dtype == torch.float32
    if is_float:
        img_np = (image[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    else:
        img_np = image[0].permute(1, 2, 0).numpy()

    h, w = img_np.shape[:2]
    out_h, out_w = size
    channels = img_np.shape[2] if img_np.ndim == 3 else 1

    # Get box coordinates (top-left and bottom-right)
    box = boxes[0].numpy()  # (4, 2)
    x1, y1 = box[0]  # top-left
    x2, y2 = box[2]  # bottom-right

    crop_w = x2 - x1
    crop_h = y2 - y1

    # Create transformation matrix
    matrix = skia.Matrix()
    scale_x = out_w / crop_w
    scale_y = out_h / crop_h
    matrix.setScale(scale_x, scale_y)
    matrix.preTranslate(-x1, -y1)

    # Prepare source image as RGBA for Skia
    if channels == 1:
        image_rgba = np.stack(
            [img_np.squeeze()] * 3 + [np.full((h, w), 255, dtype=np.uint8)], axis=-1
        )
    elif channels == 3:
        alpha = np.full((h, w, 1), 255, dtype=np.uint8)
        image_rgba = np.concatenate([img_np, alpha], axis=-1)
    else:
        raise ValueError(f"Unsupported channels: {channels}")

    image_rgba = np.ascontiguousarray(image_rgba, dtype=np.uint8)
    skia_image = skia.Image.fromarray(
        image_rgba, colorType=skia.ColorType.kRGBA_8888_ColorType
    )

    # Create output array and array-backed surface
    # This avoids the BGR/RGBA issue by writing directly to the numpy array
    output_rgba = np.zeros((out_h, out_w, 4), dtype=np.uint8)
    output_rgba[:, :, 3] = 255  # Set alpha to opaque
    surface = skia.Surface(output_rgba, colorType=skia.ColorType.kRGBA_8888_ColorType)
    canvas = surface.getCanvas()
    canvas.clear(skia.Color4f(0, 0, 0, 1))
    canvas.setMatrix(matrix)

    paint = skia.Paint()
    paint.setAntiAlias(True)
    sampling = skia.SamplingOptions(skia.FilterMode.kLinear)
    canvas.drawImage(skia_image, 0, 0, sampling, paint)

    # Flush to ensure drawing is complete
    surface.flushAndSubmit()

    # Extract appropriate channels from output array
    if channels == 1:
        result = output_rgba[:, :, 0:1]
    else:
        result = output_rgba[:, :, :3]

    result_tensor = torch.from_numpy(result).permute(2, 0, 1).unsqueeze(0)
    if is_float:
        result_tensor = result_tensor.float() / 255.0

    return result_tensor