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.

RandomMixUpV2 blends each image with a paired batch image, and RandomCutMixV2 pastes in a rectangle from it. RandomMosaic replaces an image by a crop of a mosaic of batch images; RandomJigsaw shuffles the tiles of each image. RandomTransplantation copies the region of one mask label, together with its mask, from the previous image in the batch.

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

Replace a square patch with that of a paired sample, applied to a batch with probability p (default 1.0).

Try in browser

output = PatchMix(alpha, patch_size, ...)(input) replaces a square patch in each image with the same region of a randomly paired image of the batch, which can be the image itself.

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 of applying the augmentation to the batch. 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. same_on_batch=True shares the patch corner, not the pairing. Only image input is implemented.

See the Convention block on MixAugmentationBaseV2.

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

See also

• RandomCutMixV2: CutMix with paired labels and a sampled rectangle.

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)

Mix rectangular image regions and labels, applied to a batch with probability p (default 1.0).

Try in browser

images, labels = RandomCutMixV2(data_keys=["input", "class"])(input, label) mixes image patches and returns label pairs with their coefficients. Use RandomMixUpV2 to blend complete images instead.

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 (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: _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 each mix’s cut size and placement across rows; each mix still draws its own pairing and geometry. At p=0 the image is unchanged and each class row contains the original label twice with lambda zero.

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

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

See the Convention block on MixAugmentationBaseV2.

See also

• RandomMixUpV2: blends complete paired images instead of rectangular patches.

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)

Shuffle the tiles of each image like a jigsaw puzzle, applied per sample with probability p (default 0.5).

Try in browser

output = RandomJigsaw(grid, ...)(input) rearranges the grid cells of each selected image.

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.

See the Convention block on MixAugmentationBaseV2.

See also

• RandomMosaic: composes a crop from multiple images in the batch.

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)

Mix pairs of images and labels, applied to a batch with probability p (default 1.0).

Try in browser

images, labels = RandomMixUpV2(data_keys=["input", "class"])(input, label) returns blended images and paired labels with mixing coefficients. Use RandomCutMixV2 to replace a rectangular region instead.

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: _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.

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

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

See the Convention block on MixAugmentationBaseV2.

See also

• RandomCutMixV2: mixes paired images through sampled rectangular regions.

Examples

The loss and accuracy can use the returned original label, paired label and lambda:

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()

>>> 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')

Replace each sample, with probability p (default 0.7), by a crop of a mosaic of batch images.

Try in browser

output = RandomMosaic(output_size, ...)(input) arranges batch images in a mosaic before taking one output crop.

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"

See the Convention block on MixAugmentationBaseV2.

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.

See also

• RandomJigsaw: shuffles tiles within each individual image.

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]))

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

Paste a labelled region from the previous batch image, applied per sample with probability p (default 0.5).

image_out, mask_out = RandomTransplantation(data_keys=["input", "mask"])(input, mask) copies the region of one drawn label, image features and segmentation mask alike, from the previous batch image into each selected image. Use RandomTransplantation3D for volume dispatch in AugmentationSequential.

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 of applying the transformation to each sample. 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.

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

See also

• RandomTransplantation3D: performs transplantation on 3D volumes and masks.

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.]]])