sleap.nn.data.augmentation

Transformers for applying data augmentation.

class sleap.nn.data.augmentation.AlbumentationsAugmenter(augmenter: Compose, image_key: str = 'image', instances_key: str = 'instances')

Data transformer based on the albumentations library.

This class can generate a tf.data.Dataset from an existing one that generates image and instance data. Element of the output dataset will have a set of augmentation transformations applied.

augmenter

An instance of albumentations.Compose that will be applied to each element of the input dataset.

Type:

albumentations.core.composition.Compose

image_key

Name of the example key where the image is stored. Defaults to “image”.

Type:

str

instances_key

Name of the example key where the instance points are stored. Defaults to “instances”.

Type:

str

classmethod from_config(config: AugmentationConfig, image_key: str = 'image', instances_key: str = 'instances') → AlbumentationsAugmenter

Create an augmenter from a set of configuration parameters.

Parameters:

• config – An AugmentationConfig instance with the desired parameters.

• image_key – Name of the example key where the image is stored. Defaults to “image”.

• instances_key – Name of the example key where the instance points are stored. Defaults to “instances”.

Returns:

An instance of AlbumentationsAugmenter with the specified augmentation configuration.

property input_keys: List[str]

Return the keys that incoming elements are expected to have.

property output_keys: List[str]

Return the keys that outgoing elements will have.

transform_dataset(input_ds: DatasetV2) → DatasetV2

Create a tf.data.Dataset with elements containing augmented data.

Parameters:

input_ds – A dataset with elements that contain the keys “image” and “instances”. This is typically raw data from a data provider.

Returns:

A tf.data.Dataset with the same keys as the input, but with images and instance points updated with the applied augmentations.

Notes

The “scale” key in examples are not modified when scaling augmentation is applied.

class sleap.nn.data.augmentation.RandomCropper(crop_height: int = 256, crop_width: int = 256)

Data transformer for applying random crops to input images.

This class can generate a tf.data.Dataset from an existing one that generates image and instance data. Element of the output dataset will have random crops applied.

crop_height

The height of the cropped region in pixels.

Type:

int

crop_width

The width of the cropped region in pixels.

Type:

int

transform_dataset(input_ds: DatasetV2)

Create a tf.data.Dataset with elements containing augmented data.

Parameters:

input_ds – A dataset with elements that contain the keys “image” and “instances”. This is typically raw data from a data provider.

Returns:

A tf.data.Dataset with the same keys as the input, but with images and instance points updated with the applied random crop.

Additionally, the "crop_bbox" key will contain the bounding box of the crop in the form [y1, x1, y2, x2].

class sleap.nn.data.augmentation.RandomFlipper(symmetric_inds: ndarray | None = None, horizontal: bool = True, probability: float = 0.5)

Data transformer for applying random flipping to input images.

This class can generate a tf.data.Dataset from an existing one that generates image and instance data. Elements of the output dataset will have random horizontal flips applied.

symmetric_inds

Indices of symmetric pairs of nodes as a an array of shape (n_symmetries, 2). Each row contains the indices of nodes that are mirror symmetric, e.g., left/right body parts. The ordering of the list or which node comes first (e.g., left/right vs right/left) does not matter. Each pair of nodes will be swapped to account for the reflection if this is not None (the default).

Type:

numpy.ndarray | None

horizontal

If True (the default), flips are applied horizontally instead of vertically.

Type:

bool

probability

The probability that the augmentation should be applied.

Type:

float

classmethod from_skeleton(skeleton: Skeleton, horizontal: bool = True, probability: float = 0.5) → RandomFlipper

Create an instance of RandomFlipper from a skeleton.

Parameters:

• skeleton – A sleap.Skeleton that may define symmetric nodes.

• horizontal – If True (the default), flips are applied horizontally instead of vertically.

• probability – The probability that the augmentation should be applied.

Returns:

An instance of RandomFlipper.

transform_dataset(input_ds: DatasetV2)

Create a tf.data.Dataset with elements containing augmented data.

Parameters:

input_ds – A dataset with elements that contain the keys "image" and "instances". This is typically raw data from a data provider.

Returns:

A tf.data.Dataset with the same keys as the input, but with images and instance points updated with the applied random flip.

sleap.nn.data.augmentation.flip_instances_lr(instances: Tensor, img_width: int, symmetric_inds: Tensor | None = None) → Tensor

Flip a set of instance points horizontally with symmetric node adjustment.

Parameters:

• instances – Instance points as a tf.Tensor of shape (n_instances, n_nodes, 2) and dtype tf.float32.

• img_width – Width of image in the same units as instances.

• symmetric_inds – Indices of symmetric pairs of nodes as a tf.Tensor of shape (n_symmetries, 2) and dtype tf.int32. Each row contains the indices of nodes that are mirror symmetric, e.g., left/right body parts. The ordering of the list or which node comes first (e.g., left/right vs right/left) does not matter. Each pair of nodes will be swapped to account for the reflection if this is not None (the default).

Returns:

The instance points with x-coordinates flipped horizontally.

sleap.nn.data.augmentation.flip_instances_ud(instances: Tensor, img_height: int, symmetric_inds: Tensor | None = None) → Tensor

Flip a set of instance points vertically with symmetric node adjustment.

Parameters:

• instances – Instance points as a tf.Tensor of shape (n_instances, n_nodes, 2) and dtype tf.float32.

• img_height – Height of image in the same units as instances.

• symmetric_inds – Indices of symmetric pairs of nodes as a tf.Tensor of shape (n_symmetries, 2) and dtype tf.int32. Each row contains the indices of nodes that are mirror symmetric, e.g., left/right body parts. The ordering of the list or which node comes first (e.g., left/right vs right/left) does not matter. Each pair of nodes will be swapped to account for the reflection if this is not None (the default).

Returns:

The instance points with y-coordinates flipped horizontally.