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.
crop_and_resize_skia_window	:func:crop_and_resize_skia of a frame, given only a window of it.
crop_window_bounds	The frame window (y0, y1, x0, x1) (half-open) a crop box samples from.

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

crop_and_resize_skia_window(window, origin, frame_hw, boxes, size)

:func:crop_and_resize_skia of a frame, given only a window of it.

Returns the same bytes :func:crop_and_resize_skia returns for the whole frame, when window covers :func:crop_window_bounds of the box, in time that grows with the crop instead of the frame.

Skia maps each output pixel to the source through the box's transform, in float32. Cropping a window and shifting the box (or drawing the window at an offset) changes that arithmetic, and crops can then differ by a grey level. So the transform is left as it is: the window is written into a frame-sized RGBA buffer (reused, one per thread) that Skia wraps without a copy. The rest of that buffer is never sampled, so what it holds does not matter, and nothing frame-sized is built or copied. Keep the transform and the drawing below in step with :func:crop_and_resize_skia.

Parameters:

Name	Type	Description	Default
window	Tensor	The frame window, (1, C, h, w) (C is 1 or 3), uint8 or float32 in [0, 1] like :func:crop_and_resize_skia's input.	required
origin	Tuple[int, int]	The window's top-left (x0, y0) in the frame.	required
frame_hw	Tuple[int, int]	The frame's (height, width).	required
boxes	Tensor	Box corners (1, 4, 2) in frame coordinates.	required
size	Tuple[int, int]	Output size (height, width).	required

Returns:

Type	Description
Tensor	The crop, (1, C, out_h, out_w), of the window's dtype.

Source code in sleap_nn/data/skia_augmentation.py

def crop_and_resize_skia_window(
    window: torch.Tensor,
    origin: Tuple[int, int],
    frame_hw: Tuple[int, int],
    boxes: torch.Tensor,
    size: Tuple[int, int],
) -> torch.Tensor:
    """:func:`crop_and_resize_skia` of a frame, given only a window of it.

    Returns the same bytes :func:`crop_and_resize_skia` returns for the whole frame,
    when ``window`` covers :func:`crop_window_bounds` of the box, in time that grows
    with the crop instead of the frame.

    Skia maps each output pixel to the source through the box's transform, in float32.
    Cropping a window and shifting the box (or drawing the window at an offset) changes
    that arithmetic, and crops can then differ by a grey level. So the
    transform is left as it is: the window is written into a frame-sized RGBA buffer
    (reused, one per thread) that Skia wraps without a copy. The rest of that buffer
    is never sampled, so what it holds does not matter, and nothing frame-sized is
    built or copied. Keep the transform and the drawing below in step with
    :func:`crop_and_resize_skia`.

    Args:
        window: The frame window, ``(1, C, h, w)`` (``C`` is 1 or 3), uint8 or float32
            in ``[0, 1]`` like :func:`crop_and_resize_skia`'s input.
        origin: The window's top-left ``(x0, y0)`` in the frame.
        frame_hw: The frame's ``(height, width)``.
        boxes: Box corners ``(1, 4, 2)`` in frame coordinates.
        size: Output size ``(height, width)``.

    Returns:
        The crop, ``(1, C, out_h, out_w)``, of the window's dtype.
    """
    is_float = window.dtype == torch.float32
    if is_float:
        win_np = (window[0].permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    else:
        win_np = window[0].permute(1, 2, 0).numpy()

    frame_h, frame_w = int(frame_hw[0]), int(frame_hw[1])
    out_h, out_w = size
    channels = win_np.shape[2] if win_np.ndim == 3 else 1
    if channels not in (1, 3):
        raise ValueError(f"Unsupported channels: {channels}")

    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

    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)

    # Frame-sized RGBA source; only the window is written.
    image_rgba = _frame_rgba_scratch(frame_h, frame_w)
    ox, oy = int(origin[0]), int(origin[1])
    h, w = win_np.shape[:2]
    if h > 0 and w > 0:
        target = image_rgba[oy : oy + h, ox : ox + w]
        target[..., :3] = win_np
        target[..., 3] = 255
    skia_image = skia.Image.fromarray(
        image_rgba, colorType=skia.ColorType.kRGBA_8888_ColorType, copy=False
    )

    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)
    surface.flushAndSubmit()
    # `skia_image` borrows `image_rgba`; release it before the buffer goes.
    del skia_image

    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

crop_window_bounds(boxes, frame_hw)

The frame window (y0, y1, x0, x1) (half-open) a crop box samples from.

An output pixel center of :func:crop_and_resize_skia maps strictly inside the box, and bilinear sampling reads the two source pixels around it, so every tap lies within one pixel of the box. The window adds a margin of _CROP_WINDOW_MARGIN pixels and is clipped to the frame; it can be empty when the box lies entirely outside the frame.

Parameters:

Name	Type	Description	Default
boxes	Tensor	Box corners of shape (1, 4, 2), as :func:crop_and_resize_skia takes them.	required
frame_hw	Tuple[int, int]	The (height, width) of the frame the box is in.	required

Source code in sleap_nn/data/skia_augmentation.py

def crop_window_bounds(
    boxes: torch.Tensor, frame_hw: Tuple[int, int]
) -> Tuple[int, int, int, int]:
    """The frame window ``(y0, y1, x0, x1)`` (half-open) a crop box samples from.

    An output pixel center of :func:`crop_and_resize_skia` maps strictly inside the
    box, and bilinear sampling reads the two source pixels around it, so every tap
    lies within one pixel of the box. The window adds a margin of
    ``_CROP_WINDOW_MARGIN`` pixels and is clipped to the frame; it can be empty
    when the box lies entirely outside the frame.

    Args:
        boxes: Box corners of shape ``(1, 4, 2)``, as :func:`crop_and_resize_skia`
            takes them.
        frame_hw: The ``(height, width)`` of the frame the box is in.
    """
    box = boxes[0].numpy()
    height, width = int(frame_hw[0]), int(frame_hw[1])
    x0 = max(0, int(math.floor(float(box[:, 0].min()))) - _CROP_WINDOW_MARGIN)
    y0 = max(0, int(math.floor(float(box[:, 1].min()))) - _CROP_WINDOW_MARGIN)
    x1 = min(width, int(math.floor(float(box[:, 0].max()))) + _CROP_WINDOW_MARGIN + 1)
    y1 = min(height, int(math.floor(float(box[:, 1].max()))) + _CROP_WINDOW_MARGIN + 1)
    return y0, max(y0, y1), x0, max(x0, x1)