2D mix transforms#

Transforms that combine samples or rearrange patches within an image. Supported label, box and mask keys vary by class; mixing does not imply support for every annotation type.

class kornia.augmentation.PatchMix(alpha=1.0, patch_size=16, p=1.0, same_on_batch=False, keepdim=False, data_keys=None)[source]#

PatchMix augmentation.

Try in browser

Replaces a random patch in each image of a batch with the corresponding region from a randomly chosen batch image, which can be the same image.

See the Convention block on MixAugmentationBaseV2.

Implementation for CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features [YHO+19].

Parameters:

alpha (float, optional) – Hyperparameter for the Beta distribution used to generate the mixing parameter lambda. Default: 1.0
patch_size (int, optional) – The size of the square patch to mix. Default: 16
p (float, optional) – Probability for applying the augmentation. Default: 1.0
same_on_batch (bool, optional) – Apply the same transformation across the batch. Default: False
keepdim (bool, optional) – Whether to keep the output shape the same as input True or broadcast it to the batch form False. Default: False

Convention:

patch_size is the side length in pixels of a square patch and must not exceed min(H, W); a larger value, such as the default 16 on a smaller image, raises ValueError when parameters are drawn, whatever the gate. Each output sample copies that square from its paired input, _params["mix_pairs"], at the same sampled (x, y) top-left corner, _params["patch_coords"]. The sampled lam is stored but blends nothing.
p is a batch-wide gate (#4425). same_on_batch=True shares the patch corner, not the pairing, which may pair a sample with itself. Only image input is implemented.

Examples

>>> aug = PatchMix(alpha=1.0, patch_size=4)
>>> x = torch.rand(2, 3, 32, 32)
>>> out = aug(x)
>>> out.shape
torch.Size([2, 3, 32, 32])

class kornia.augmentation.RandomCutMixV2(num_mix=1, cut_size=None, beta=None, same_on_batch=False, p=1.0, keepdim=False, data_keys=None, use_correct_lambda=False)[source]#

Apply CutMix augmentation to a batch of torch.Tensor images.

Try in browser

Implementation for CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features [YHO+19].

The function returns (inputs, labels), in which the inputs is the torch.Tensor that contains the mixup images while the labels is a \((\text{num_mixes}, B, 3)\) torch.Tensor that contains (label_batch, label_permuted_batch, lambda) for each cutmix.

The implementation referred to the following repository: clovaai/CutMix-PyTorch.

See the Convention block on MixAugmentationBaseV2.

Parameters:

p (float, optional) – probability for applying an augmentation to a batch. This param controls the augmentation probabilities batch-wisely. Default: 1.0
num_mix (int, optional) – cut mix times. Default: 1
beta (Union[Tensor, float, None], optional) – hyperparameter for generating cut size from beta distribution. Beta cannot be set to 0 after torch 1.8.0. If None, it will be set to 1. Default: None
cut_size (Union[Tensor, Tuple[float, float], None], optional) – the [min, max] clamp, within [0, 1], applied to the Beta-sampled mixing coefficient lambda, not the fraction of the image to cut. The cut side is floor(sqrt(1 - lambda) * side), so a larger cut_size gives a smaller cut, and a minimum of 1 is rejected because it would cut nothing. If None, it will be set to [0, 1], which means no restriction. Default: None
same_on_batch (bool, optional) – apply the same transformation across the batch. This flag will not maintain permutation order. Default: False
keepdim (bool, optional) – whether to keep the output shape the same as input (True) or broadcast it to the batch form (False). Default: False
data_keys (Optional[List[Union[str, int, DataKey]]], optional) – list of keys to apply the mix augmentation to. Default: None
use_correct_lambda (bool, optional) – if True, compute lambda according to the CutMix paper (lam = 1 - area_ratio). Defaults to False (lam = area_ratio) for backward compatibility, but will raise a deprecation warning when False. Default: False

Inputs:

Input image tensors, shape of \((B, C, H, W)\).
Raw labels, shape of \((B)\).

Returns:

Adjusted image, shape of \((B, C, H, W)\).
Raw labels, permuted labels and lambdas for each mix, shape of \((num_mix, B, 3)\).

Return type:

Tuple[torch.Tensor, torch.Tensor]

Convention:

With data_keys=["input", "class"], the class output is (num_mix, B, 3). Its last axis holds the original label, the label selected by that mix’s permutation, and the area fraction returned as lambda. Labels must be one-dimensional and are returned as floating values in the image dtype, except that a float16 or bfloat16 image yields float32 labels so integer class ids up to 2 ** 24 stay exact. use_correct_lambda=True returns 1 - cut_area / image_area; the compatibility default returns cut_area / image_area and emits a deprecation warning.
p is a batch-wide gate applied once, unlike the per-sample p of RandomJigsaw (#4425): _params["batch_prob"] is all ones or all zeros, and every row and mix of a selected batch receives a cut, though self-pairing or a zero-sized cut can leave a row unchanged. With same_on_batch=False each row and mix draws its own cut size and placement; same_on_batch=True shares one geometry, and with num_mix > 1 one pairing too, so every mix repeats the first cut while the labels credit the donor once per mix (#4805). At p=0 the image is unchanged and each class row contains the original label twice with lambda zero.

Note

This implementation would randomly cutmix images in a batch. Ideally, the larger batch size would be preferred.

Examples

>>> rng = torch.manual_seed(3)
>>> input = torch.rand(2, 1, 3, 3)
>>> input[0] = torch.ones((1, 3, 3))
>>> label = torch.tensor([0, 1])
>>> cutmix = RandomCutMixV2(data_keys=["input", "class"], use_correct_lambda=True)
>>> cutmix(input, label)
[tensor([[[[1.0000, 1.0000, 1.0000],
          [0.1498, 0.4015, 1.0000],
          [0.4594, 0.1756, 1.0000]]],


        [[[0.8879, 1.0000, 1.0000],
          [0.1498, 1.0000, 1.0000],
          [0.4594, 0.1756, 0.9492]]]]), tensor([[[0.0000, 1.0000, 0.5556],
         [1.0000, 0.0000, 0.5556]]])]

class kornia.augmentation.RandomJigsaw(grid=(4, 4), data_keys=None, p=0.5, same_on_batch=False, keepdim=False, ensure_perm=True)[source]#

RandomJigsaw augmentation.

Try in browser

Make Jigsaw puzzles for each image individually. To mix with different images in a batch, referring to kornia.augmentation.RandomMosaic.

See the Convention block on MixAugmentationBaseV2.

Parameters:

grid (Tuple[int, int], optional) – the Jigsaw puzzle grid. e.g. (2, 2) means each output will mix image patches in a 2x2 grid. Default: (4, 4)
ensure_perm (bool, optional) – reject the image-preserving permutation arange(N).view(rows, columns).T.flatten() when drawing, so a selected sample is never returned unchanged. For a single-row or single-column grid that permutation is [0, ..., N - 1]; see the Convention block. A 1 x 1 grid has no other permutation, so it raises ValueError unless ensure_perm=False. Default: True
data_keys (Optional[List[Union[str, int, DataKey]]], optional) – the input type sequential for applying augmentations. Only “input” and “image” are implemented; see the Convention block. Default: None
p (float, optional) – probability of applying the transformation to each sample. Default: 0.5
same_on_batch (bool, optional) – apply the same transformation across the batch. Default: False
keepdim (bool, optional) – whether to keep the output shape the same as input True or broadcast it to the batch form False. Default: False

Convention:

grid=(rows, columns) partitions each image independently. Both image dimensions must be divisible by their grid entries, otherwise RuntimeError is raised whatever the gate. An entry’s position selects the destination cell in column-major order, while its value indexes the source patch in row-major order, so the identity permutation reproduces the image only on a single-row or single-column grid (on a square grid it transposes the patches). The image-preserving permutation is arange(N).view(rows, columns).T.flatten(), [0, 2, 1, 3] for a 2 x 2 grid.
p is a per-sample gate; same_on_batch=True shares one gate draw and one permutation. Only image input is implemented.

Examples

>>> jigsaw = RandomJigsaw((4, 4))
>>> input = torch.randn(8, 3, 256, 256)
>>> out = jigsaw(input)
>>> out.shape
torch.Size([8, 3, 256, 256])

class kornia.augmentation.RandomMixUpV2(lambda_val=None, same_on_batch=False, p=1.0, keepdim=False, data_keys=None)[source]#

Apply MixUp augmentation to a batch of torch.Tensor images.

Try in browser

Implementation for mixup: BEYOND EMPIRICAL RISK MINIMIZATION [ZnYNDLP18].

The function returns (inputs, labels), in which the inputs is the torch.Tensor that contains the mixup images while the labels is a \((B, 3)\) torch.Tensor that contains (label_batch, label_permuted_batch, lambda) for each image.

The implementation is on top of the following repository: hongyi-zhang/mixup.

See the Convention block on MixAugmentationBaseV2.

The loss and accuracy are computed as:

def loss_mixup(y, logits):
    criterion = F.cross_entropy
    loss_a = criterion(logits, y[:, 0].long(), reduction='none')
    loss_b = criterion(logits, y[:, 1].long(), reduction='none')
    return ((1 - y[:, 2]) * loss_a + y[:, 2] * loss_b).mean()

def acc_mixup(y, logits):
    pred = torch.argmax(logits, dim=1).to(y.device)
    return (1 - y[:, 2]) * pred.eq(y[:, 0]).float() + y[:, 2] * pred.eq(y[:, 1]).float()

Parameters:

p (float, optional) – probability for applying an augmentation to a batch. This param controls the augmentation probabilities batch-wisely. Default: 1.0
lambda_val (Union[Tensor, Tuple[float, float], None], optional) – min-max value of mixup strength. Default is 0-1. Default: None
same_on_batch (bool, optional) – apply the same transformation across the batch. This flag will not maintain permutation order. Default: False
keepdim (bool, optional) – whether to keep the output shape the same as input (True) or broadcast it to the batch form (False). Default: False

Inputs:

Input image tensors, shape of \((B, C, H, W)\).
Label: raw labels, shape of \((B)\).

Returns:

Adjusted image, shape of \((B, C, H, W)\).
Raw labels, permuted labels and lambdas for each mix, shape of \((B, 3)\).

Return type:

Tuple[torch.Tensor, torch.Tensor]

Convention:

With data_keys=["input", "class"], the class output is (B, 3): its columns are the original label, the label selected by _params["mixup_pairs"], and the sampled lambda. The image is input * (1 - lambda) + paired_input * lambda. Labels must be one-dimensional and are returned as floating values in the image dtype, except that a float16 or bfloat16 image yields float32 labels so integer class ids up to 2 ** 24 stay exact.
p is a batch-wide gate applied once, unlike the per-sample p of RandomJigsaw (#4425): _params["batch_prob"] is all ones or all zeros, and every row of a selected batch draws its lambda from lambda_val. At p=0 the image is unchanged and the class output holds the original label twice with a zero lambda. same_on_batch=True shares the lambda draw but not the pairing, which may pair a sample with itself.

Note

This implementation would randomly mixup images in a batch. Ideally, the larger batch size would be preferred.

Examples

>>> rng = torch.manual_seed(1)
>>> input = torch.rand(2, 1, 3, 3)
>>> label = torch.tensor([0, 1])
>>> mixup = RandomMixUpV2(data_keys=["input", "class"])
>>> mixup(input, label)
[tensor([[[[0.7576, 0.2793, 0.4031],
          [0.7347, 0.0293, 0.7999],
          [0.3971, 0.7544, 0.5695]]],


        [[[0.4388, 0.6387, 0.5247],
          [0.6826, 0.3051, 0.4635],
          [0.4550, 0.5725, 0.4980]]]]), tensor([[0.0000, 0.0000, 0.1980],
        [1.0000, 1.0000, 0.4162]])]

class kornia.augmentation.RandomMosaic(output_size=None, mosaic_grid=(2, 2), start_ratio_range=(0.3, 0.7), min_bbox_size=0.0, data_keys=None, p=0.7, keepdim=False, padding_mode='constant', resample=Resample.BILINEAR.name, align_corners=True, cropping_mode='slice')[source]#

Mosaic augmentation.

Try in browser

https://raw.githubusercontent.com/kornia/data/main/random_mosaic.png

Given a certain number of images, mosaic transform combines them into one output image. The output image is composed of the parts from each sub-image. To mess up each image individually, referring to kornia.augmentation.RandomJigsaw.

The mosaic transform steps are as follows:

Concate selected images into a super-image.

Crop out the outcome image according to the top-left corner and crop size.

See the Convention block on MixAugmentationBaseV2.

Parameters:

output_size (Optional[Tuple[int, int]], optional) – the output (height, width) after mosaicing. Default: None
start_ratio_range (Tuple[float, float], optional) – the (low, high) range from which both top-left crop ratios (x / W, y / H) are drawn. Default: (0.3, 0.7)
mosaic_grid (Tuple[int, int], optional) – the number of images and image arrangement. e.g. (2, 2) means each output will mix 4 images in a 2x2 grid. Default: (2, 2)
min_bbox_size (float, optional) – minimum area of bounding boxes. Default to 0. Default: 0.0
data_keys (Optional[List[Union[str, int, DataKey]]], optional) – the input type sequential for applying augmentations. Only “input”, “image”, “bbox”, “bbox_xyxy” and “bbox_xywh” are implemented; see the Convention block. Default: None
p (float, optional) – probability of applying the transformation to each sample. Default: 0.7
keepdim (bool, optional) – whether to keep the output shape the same as input True or broadcast it to the batch form False. Default: False
padding_mode (str, optional) – Type of padding, used only with cropping_mode="resample". Should be: constant, reflect, replicate. Default: "constant"
resample (Union[str, int, Resample], optional) – the interpolation mode. Default: Resample.BILINEAR.name
align_corners (bool, optional) – interpolation flag. Default: True
cropping_mode (str, optional) – The used algorithm to crop. slice will use advanced slicing to extract the torch.Tensor based on the sampled indices. resample will use warp_affine using the affine transformation to extract the window. Use slice for efficiency, or resample for proper differentiability. Default: "slice"

Examples

>>> mosaic = RandomMosaic((300, 300), data_keys=["input", "bbox_xyxy"])
>>> boxes = torch.tensor([[
...     [70, 5, 150, 100],
...     [60, 180, 175, 220],
... ]]).repeat(8, 1, 1)
>>> input = torch.randn(8, 3, 224, 224)
>>> out = mosaic(input, boxes)
>>> out[0].shape, out[1].shape
(torch.Size([8, 3, 300, 300]), torch.Size([8, 8, 4]))

Convention:

output_size and the default output shape are ordered (height, width). With output_size=None the output preserves the input’s (H, W) even when they differ, in both cropping_mode="slice" and cropping_mode="resample". start_ratio_range draws a pair used as (x / W, y / H) for the crop’s top-left corner, which is floored to a whole pixel that both cropping modes and the boxes share. These are the repaired axis conventions from #4438.
p is per sample and this class fixes same_on_batch=False. It composes mosaic_grid[0] tiles along width and mosaic_grid[1] tiles along height, then crops each result. It supports "bbox", "bbox_xyxy", and "bbox_xywh" in addition to image inputs; it does not support masks, keypoints, or class labels.
Both cropping_mode values take an output_size window of the composed canvas at the drawn corner, without rescaling; where the window passes the canvas edge it is zero. With an explicit output_size an unselected sample is zero-padded or cropped to that size rather than returned unchanged. The boxes of a selected sample are translated with its tiles, not rescaled, and clipped to the output window [0, W_out] x [0, H_out].
A tensor box input comes back as one dense tensor. When the gate selects any sample, every row holds N * mosaic_grid[0] * mosaic_grid[1] boxes for N input boxes per sample: a selected row holds its mosaic boxes, and an unselected row holds its own N boxes unchanged followed by all-zero padding rows ([0, 0, 0, 0] in "bbox_xyxy" and "bbox_xywh", four (0, 0) vertices in "bbox"), which have zero area. When the gate selects no sample the boxes are returned unchanged with their N rows.
A list box input comes back as a list with one tensor per sample, from a direct call and from AugmentationSequential alike, whatever the gate selects. An unselected sample’s tensor is its own boxes unchanged, with no padding rows. A selected sample’s tensor holds the boxes of each tile’s source image, with no padding rows.
In either form, a selected sample’s box that falls below min_bbox_size is zeroed by filter_boxes_by_area(). It exports as exact zeros in every box format, from a direct call and from AugmentationSequential.

class kornia.augmentation.RandomTransplantation(excluded_labels=None, p=0.5, p_batch=1.0, data_keys=None)[source]#

RandomTransplantation augmentation.

Randomly transplant (copy and paste) image features and corresponding segmentation masks between images in a batch. The transplantation transform works as follows:

Based on the parameter p, a certain number of images in the batch are selected as acceptor of a transplantation.

For each acceptor, the image below in the batch is selected as donor (via circling: \(i - 1 \mod B\)).

From the donor, a random label is selected and the corresponding image features and segmentation mask are transplanted to the acceptor.

The augmentation is described in Semantic segmentation of surgical hyperspectral images under geometric domain shifts [SSSF+23].

Parameters:

excluded_labels (Union[Sequence[int], Tensor, None], optional) – sequence of labels which should not be transplanted from a donor. This can be useful if only parts of the image are annotated and the non-annotated regions (with a specific label index) should be excluded from the augmentation. If no label is left in a donor image, that acceptor receives nothing while the rest of the batch is transplanted as usual. Default: None
p (float, optional) – probability for applying an augmentation to an image. This parameter controls how many images in a batch receive a transplant. Default: 0.5
p_batch (float, optional) – probability for applying an augmentation to a batch. This param controls the augmentation probabilities batch-wise. Default: 1.0
data_keys (Optional[list[str | int | DataKey]], optional) – the input type sequential for applying augmentations. There must be at least one “mask” torch.Tensor. If no data keys are given, the first torch.Tensor is assumed to be DataKey.INPUT and the second torch.Tensor DataKey.MASK. Accepts “input”, “mask”. Default: None

Note

This augmentation requires that segmentation masks are available for all images in the batch and that at least some objects in the image are annotated.
When using this class directly (RandomTransplantation()(…)), it works for arbitrary spatial dimensions including 2D and 3D images. When wrapping in kornia.augmentation.container.AugmentationSequential, use kornia.augmentation.RandomTransplantation for 2D and kornia.augmentation.RandomTransplantation3D for 3D images.

Inputs:

Segmentation mask torch.Tensor which is used to determine the objects for transplantation: \((B, *)\).
(optional) Additional image or mask tensors where the features are transplanted based on the first segmentation mask: \((B, C, *)\) (DataKey.INPUT) or \((B, *)\) (DataKey.MASK).

Returns:

torch.Tensor:

Augmented mask tensors: \((B, *)\).

list[torch.Tensor]:

Augmented mask tensors: \((B, *)\).
Additional augmented image or mask tensors: \((B, C, *)\) (DataKey.INPUT) or \((B, *)\) (DataKey.MASK).

Return type:

torch.Tensor | list[torch.Tensor]

Convention:

the transplant is driven by the first "mask" input. An image is (B, C, *spatial) and a mask (B, *spatial) for any number of spatial dimensions, including none (a (B,) mask moves whole images). There is no unbatched form, and the inherited keepdim is inert on a direct call. In AugmentationSequential a 2D (B, H, W) mask comes back as (B, 1, H, W), like from any other child, unless the container sets keepdim=True.
p gates samples and p_batch the whole call. The donor of acceptor i is image (i - 1) mod B of the full batch, with its original content, whether or not it is an acceptor itself; at B = 1 the call is an identity.
one label is drawn uniformly from the donor’s distinct labels minus excluded_labels (over labels, not area), on the global CPU generator whatever the mask’s device. An acceptor whose donor has no eligible label is dropped from acceptor_indices and its batch_prob is cleared.
selection in _params is what the transform reads: row d marks the positions moved from donor_indices[d] into acceptor_indices[d]. A params passed back is used as given and completed in place. Labels are drawn only when selected_labels and selection are both missing; a missing selection alone is rebuilt from the given labels as they are (by design, excluded_labels filters only drawn labels), and a label list shorter than the acceptors leaves the trailing ones untouched.
images are float16, bfloat16, float32 or float64; a mask keeps its own dtype (bool, integer or floating). Only "input" / "image" and "mask" keys are implemented; any other key raises NotImplementedError. A mask is needed only to draw the parameters; a call that draws them without one, or whose inputs do not match data_keys, raises a kornia error that names the mask (and the key count, on a mismatch) before any input is read. With a complete params any subset of the inputs is accepted, which is how AugmentationSequential applies the transplant, one input at a time. Outputs come back in input order, a single one as a bare tensor.
like every mix augmentation it is not geometric: transform_matrix and inverse() both raise RuntimeError, and the inverse() of a AugmentationSequential holding one raises the same error.
a mask with a singleton channel axis, (B, 1, *spatial), next to an image of the same rank is read as the (B, *spatial) mask: the labels and selection come from its spatial layout and it comes back in its own (B, 1, *spatial) layout. That is the layout a container hands on once any other augmentation has run, so the transplant works at any position in a pipeline, not only as the first step (#4707). A further mask one rank above the driving mask is read as (B, C, *spatial), and every channel moves through the same selection. With no image in the call, a mask keeps its own rank and a size-1 first spatial axis is not squeezed.

Examples

>>> import torch
>>> rng = torch.manual_seed(0)
>>> aug = RandomTransplantation(p=1.)
>>> image = torch.randn(2, 3, 5, 5)
>>> mask = torch.randint(0, 3, (2, 5, 5))
>>> mask
tensor([[[0, 0, 1, 1, 0],
         [1, 2, 0, 0, 0],
         [1, 2, 1, 1, 0],
         [0, 0, 0, 0, 2],
         [2, 2, 2, 0, 2]],

        [[2, 0, 0, 2, 1],
         [2, 1, 0, 2, 1],
         [2, 0, 1, 0, 2],
         [2, 2, 2, 0, 2],
         [2, 1, 0, 0, 0]]])
>>> image_out, mask_out = aug(image, mask)
>>> image_out.shape
torch.Size([2, 3, 5, 5])
>>> mask_out.shape
torch.Size([2, 5, 5])
>>> mask_out
tensor([[[2, 0, 1, 2, 0],
         [2, 2, 0, 2, 0],
         [2, 2, 1, 1, 2],
         [2, 2, 2, 0, 2],
         [2, 2, 2, 0, 2]],

        [[0, 0, 0, 2, 0],
         [2, 1, 0, 0, 0],
         [2, 0, 1, 0, 0],
         [0, 0, 0, 0, 2],
         [2, 1, 0, 0, 0]]])
>>> aug._params["selected_labels"]  # Image 0 received label 2 from image 1 and image 1 label 0 from image 0
tensor([2, 0])

You can apply the same augmentation again in which case the same objects get transplanted between the images:

>>> aug._params["selection"]  # The pixels (objects) which get transplanted
tensor([[[ True, False, False,  True, False],
         [ True, False, False,  True, False],
         [ True, False, False, False,  True],
         [ True,  True,  True, False,  True],
         [ True, False, False, False, False]],

        [[ True,  True, False, False,  True],
         [False, False,  True,  True,  True],
         [False, False, False, False,  True],
         [ True,  True,  True,  True, False],
         [False, False, False,  True, False]]])
>>> image2 = torch.zeros(2, 3, 5, 5)
>>> image2[1] = 1
>>> image2[:, 0]
tensor([[[0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.],
         [0., 0., 0., 0., 0.]],

        [[1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.]]])
>>> image_out2, mask_out2 = aug(image2, mask, params=aug._params)
>>> image_out2[:, 0]
tensor([[[1., 0., 0., 1., 0.],
         [1., 0., 0., 1., 0.],
         [1., 0., 0., 0., 1.],
         [1., 1., 1., 0., 1.],
         [1., 0., 0., 0., 0.]],

        [[0., 0., 1., 1., 0.],
         [1., 1., 0., 0., 0.],
         [1., 1., 1., 1., 0.],
         [0., 0., 0., 0., 1.],
         [1., 1., 1., 0., 1.]]])