Base Classes#

These are the base classes for creating a new transform on top of the predefined routine of kornia.augmentation. Any given augmentation can be classified as either rigid (e.g. affine transformations that manipulate images with a standard transformation matrix) or non-rigid (e.g. cutting out a random area). At the image level, Kornia provides GeometricAugmentationBase2D for rigid transformations that modify the geometric location of image pixels, IntensityAugmentationBase2D for transformations that preserve pixel locations, and the generic AugmentationBase2D, which allows more freedom for customized augmentation design.

The Base-Class Hierarchy#

The bases form a chain, each layer adding one concern and shared by several subclasses — so pick the shallowest base that already does what you need:

nn.Module
└─ _BasicAugmentationBase          parameter sampling + the forward skeleton
   ├─ _AugmentationBase            dispatch to image / mask / box / keypoint / class data keys
   │  ├─ AugmentationBase2D        2D tensor validation  (subclass for a fully custom 2D op)
   │  │  └─ RigidAffineAugmentationBase2D     transform-matrix machinery
   │  │     ├─ IntensityAugmentationBase2D    intensity ops — override apply_transform
   │  │     └─ GeometricAugmentationBase2D    warp ops — also override compute_transformation
   │  └─ AugmentationBase3D … (the 3D mirror of the 2D chain)
   └─ MixAugmentationBaseV2        mix ops (MixUp / CutMix) — bypass the per-key dispatch

Each level is a distinct, reused axis (sampling, data-key dispatch, 2D vs 3D, rigid-matrix vs free-form, intensity vs geometric); the four *Base2D classes are public API that external code subclasses. For a custom augmentation, subclass IntensityAugmentationBase2D or GeometricAugmentationBase2D.

The Predefined Augmentation Routine#

Kornia augmentations generally follow a sample-apply routine.

sample: Kornia aims at flexible tensor-level augmentations that augment every image in a batch with different parameters and probabilities. The sampling step first draws a set of random parameters. The sampled augmentation state is then stored in the _params attribute of the augmentation, for replay. Application-time draws can require additional RNG control; see the reproducibility notes below.
apply: with the generated (or user-provided) parameters, the augmentation is performed accordingly. Apart from transforming image tensors, Kornia also supports inverse operations that revert the transform, and transforms of other data modalities (data keys in Kornia) such as masks, keypoints, and bounding boxes. These features depend on the concrete operation and its data-key handlers. AugmentationSequential passes the matrix of the call’s parameters to the mask, box and keypoint handlers of geometric children and of custom RigidAffineAugmentationBase2D subclasses, recomputing it from replayed params that the child did not last apply to an image or that were edited since. An in-place edit made inside torch.inference_mode is not seen, because inference tensors have no version counter; replace the value instead. Non-rigid coordinate transforms are not supplied automatically (#4420).

Custom Augmentation Classes#

IntensityAugmentationBase2D supplies an identity matrix and default passthrough handlers for most annotations. Subclasses can override these defaults: RandomErasing zero-fills the erased mask region. Direct and container calls pass masks, boxes, and keypoints through unchanged unless a subclass overrides apply_transform_mask, apply_transform_box, or apply_transform_keypoint. GeometricAugmentationBase2D supplies the dispatch used for geometric coordinate transformations.

For a geometric operation, implement compute_transformation and apply_transform. Supporting image inversion also requires inverse_transform; the base provides matrix inversion. Some configurations, such as slice-mode crops, reject inversion.

class kornia.augmentation.GeometricAugmentationBase2D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

GeometricAugmentationBase2D base class for customized geometric augmentation implementations.

See the Convention block on AugmentationBase2D.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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:

pixel coordinates are (x, y) at integer pixel centres, with corners (0, 0) and (W - 1, H - 1) (see Conventions & Pitfalls); rotations, shears and affine maps are centred at ((W - 1) / 2, (H - 1) / 2), and transform_matrix maps input pixel coordinates to output pixel coordinates.
the resampling classes, including the non-rigid RandomElasticTransform and RandomThinPlateSpline, default to bilinear interpolation with zero sampler padding. The align_corners default is True for RandomRotation, RandomRotation90 and the crop and resize classes, and False for RandomAffine, RandomShear, RandomTranslate, RandomPerspective, RandomElasticTransform and RandomThinPlateSpline (#4412).
this base provides a matrix-based inverse interface. Whether a concrete augmentation can invert a call depends on its implementation and configuration: slice-mode crops, for example, do not support it. Inverse resampling cannot recover image or mask information lost through cropping, padding, or interpolation. Tensor-form boxes may lose rotated corners through axis-aligned enclosure. For an empty batch, image and mask inverse restore the spatial dimensions recorded in the parameters.
containers process integer masks in at least float32 and reject labels that cannot be represented exactly in the working dtype. Direct transform_masks calls use the image dtype guard and therefore reject bool masks.
a subclass supplies its own inverse_transform() – the base raises NotImplementedError – while compute_inverse_transformation() defaults to inverting the sampled matrix.
a scalar magnitude x means center ± x, as it does for the intensity classes: its lower end is floored at the parameter’s lower bound, and an upper end past the upper bound raises the same error the explicit range does. The bounds are (-360, 360) for RandomRotation’s and RandomAffine’s degrees and for RandomShear’s and RandomAffine’s shear, (-1, 1) for RandomTranslate’s translate_x and translate_y, (0, 1) for RandomAffine’s translate, and (-3, 3) for RandomRotation90’s times. The 3D angle ranges – RandomRotation3D’s and RandomAffine3D’s degrees, RandomAffine3D’s shears and RandomMotionBlur3D’s angle – raise ValueError past (-360, 360), in scalar and explicit form.

Note

Masks are resampled with nearest neighbour whatever resample the augmentation uses for images, unless resample is passed explicitly to transform_masks() or inverse_masks() – which is what extra_args[DataKey.MASK] does in a container. An explicit value is honoured in both directions. Antialiasing is disabled for masks by default, even when it is enabled for images. Pass antialias=True explicitly to transform_masks() or through extra_args[DataKey.MASK] to filter soft masks.

compute_transformation(input, params, flags)[source]#

Return type:: Tensor

apply_transform(input, params, flags, transform=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Return type:: Tensor

compute_inverse_transformation(transform)[source]#

Compute the inverse transform of given transformation matrices.

Return type:: Tensor

inverse_transform(input, flags, transform=None, size=None)[source]#

By default, the exact transformation as apply_transform will be used.

Return type:: Tensor

For IntensityAugmentationBase2D, the user only needs to override apply_transform.

class kornia.augmentation.IntensityAugmentationBase2D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

IntensityAugmentationBase2D base class for customized intensity augmentation implementations.

See the Convention block on AugmentationBase2D.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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:

an intensity augmentation moves no image pixel, so this base supplies compute_transformation itself and reports the identity matrix, built lazily on the first read of transform_matrix. Its mask, box and keypoint handlers pass their inputs through unless a subclass overrides them (RandomErasing also erases masks). It adds no inverse: inverse skips 2D intensity children.
these classes assume the [0, 1] float image range stated under “Image tensors” in Conventions & Pitfalls. The base does not validate it; each class page states what the class does outside that range, and a class whose page is silent does not clamp.
what this block and the class pages say about an output describes the samples the p gate transforms; every other sample comes back with its input values and, where they are finite, an identity gradient. Below p=1 RandomEqualize, RandomClahe and RandomGamma run their transform on 0.5 for the skipped samples, so an out-of-range skipped image does not make them raise.
the scalar factors a concrete class draws are per sample – one value, or one per channel per sample where the class’s own docstring says so – and same_on_batch=True shares one draw. RandomMotionBlur draws one kernel size for the whole batch; Normalize, Denormalize and RandomDissolving hard-code same_on_batch=True. A drawn whole-image field (gaussian_noise, gradient, plasma, and RandomSaltAndPepperNoise’s mask_salt and mask_pepper) is stored with the batched (B, C, H, W) input shape, even for a (C, H, W) input with keepdim=True; RandomPlasmaShadow’s single-channel map is (B, 1, H, W). With same_on_batch=True, gaussian_noise is stored as (1, C, H, W) and the plasma classes store an expanded view that shares one map across the batch.
where a class documents bounds for a parameter, a range outside them usually raises at construction; the class page states a bound it does not enforce, or checks only on the forward pass. A scalar magnitude x usually means center ± x around a centre the class fixes – the neutral value, for a factor such as contrast – with its lower end floored at the bound, so contrast=1.5 reads as [0, 2.5], while an upper end past the bound raises. RandomSharpness reads a scalar as [0, x] and RandomPosterize as [x, 8].

Warning

Several of these classes return an all-zero image, with no warning, for an input whose values are all negative, and RandomSolarize does for one whose values are all at least 1.5. Tracked in #4430.

apply_transform(input, params, flags, transform=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Return type:: Tensor

The most common case is a pixel-wise augmentation with a random per-sample parameter. Subclass IntensityAugmentationBase2D, declare a parameter generator in __init__, and read the sampled value in apply_transform:

from typing import Any, Dict, Optional

import torch
from torch import Tensor

from kornia.augmentation import IntensityAugmentationBase2D
from kornia.augmentation import random_generator as rg

class RandomAddValue(IntensityAugmentationBase2D):
    """Add a per-sample value drawn uniformly from ``add_range``."""

    def __init__(self, add_range=(0.0, 0.2), same_on_batch=False, p=1.0, keepdim=False):
        super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim)
        # A PlainUniformGenerator sampler is a 4-tuple ``(range, name, center, bound)``:
        # sample a value inside ``range`` and expose it as ``params[name]``. ``center`` and
        # ``bound`` (``None`` here) are optional constraints for centred/bounded ranges.
        self._param_generator = rg.PlainUniformGenerator((add_range, "add", None, None))

    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any],
        transform: Optional[Tensor] = None,
    ) -> Tensor:
        # ``params["add"]`` has shape ``(B,)`` — reshape to broadcast over C, H, W.
        add = params["add"].to(input).view(-1, 1, 1, 1)
        return input + add

aug = RandomAddValue((0.0, 0.2), p=1.0)
out = aug(torch.rand(4, 3, 32, 32))                       # a different value per sample
again = aug(torch.rand(4, 3, 32, 32), params=aug._params)  # reuse the recorded draw

Static (non-random) configuration goes in self.flags (a plain dict), read from the flags argument of apply_transform. A custom augmentation works standalone and inside AugmentationSequential with no extra wiring.

For a rigid geometric augmentation, implement compute_transformation to return the (B, 3, 3) matrix and apply_transform to implement the corresponding image operation. The geometric base propagates the matrix to boxes and keypoints; mask processing uses its mask handler. The following example only illustrates the method signatures, not an invertible geometric warp:

from typing import Any, Dict, Optional

import torch
from torch import Tensor

from kornia.augmentation import GeometricAugmentationBase2D
from kornia.augmentation import random_generator as rg
from kornia.core.ops import eye_like

class MyRandomTransform(GeometricAugmentationBase2D):

   def __init__(
      self,
      factor=(0., 1.),
      same_on_batch: bool = False,
      p: float = 1.0,
      keepdim: bool = False,
   ) -> None:
      super().__init__(p=p, same_on_batch=same_on_batch, keepdim=keepdim)
      self._param_generator = rg.PlainUniformGenerator((factor, "factor", None, None))

   def compute_transformation(self, input, params, flags):
      # return the (B, 3, 3) transform matrix for this augmentation
      # identity shown only to illustrate the required matrix shape
      return eye_like(3, input)

   def apply_transform(
      self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
   ) -> Tensor:
      factor = params["factor"].to(input).view(-1, 1, 1, 1)
      return input * factor

For non-rigid augmentations, the user may implement the apply_transform* and apply_non_transform* APIs as needed. Specifically, apply_transform* applies to the elements of a batch that are selected for augmentation, while apply_non_transform* applies to the elements that are skipped. For example, a crop operation changes the size of the selected elements, so the skipped elements must be resized as well to keep the whole batch tensor at one size.

class kornia.augmentation.AugmentationBase2D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

AugmentationBase2D base class for customized augmentation implementations.

AugmentationBase2D aims at offering a generic base class for a greater level of customization. If the subclass contains routined matrix-based transformations, RigidAffineAugmentationBase2D might be a better fit.

This class is the anchor for the contract the 2D augmentations of kornia.augmentation inherit. Mix and 3D augmentations have their own bases and contracts.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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:

the working layout is (B, C, H, W) float. A (C, H, W) input is promoted to (1, C, H, W) and an (H, W) input to (1, 1, H, W); keepdim=True restores the input rank on the way out and never drops a real batch dimension. The dtype guard accepts float16, bfloat16, float32 and float64 and raises TypeError on an integer tensor. The output keeps the input’s device.
a wrong input rank raises different exception types depending on the entry point that sees it (#4424).
this base samples p per sample and uses p_batch as a call-wide gate. same_on_batch=True shares applicable sampled values across the batch. Concrete constructors need not expose p_batch; see #4425. A sample the gate skips comes back with its input values, its values cannot make the call raise, and where they are finite its gradient is the identity; a class whose transform checks input values or has an unbounded derivative runs it on 0.5 for skipped samples. An augmentation that changes the sample shape still raises ValueError when the gate mixes applied and skipped samples (#4497).
parameter sampling normally uses CPU defaults independently of the image device. Moving RNG state or samplers is incomplete for some generators, and returned parameter placement is a separate concern; see #4426.
reproducibility goes through torch’s global generators; there is no per-instance generator. generator= raises at construction, and forward silently drops it like any other unknown keyword (#4427). Conventions & Pitfalls states the seeding, DataLoader-worker, parameter-placement and serialization rules.
the last draw is kept in _params; forward(x, params=...) replaces that dict wholesale, stores the caller’s dict by reference and, if batch_prob is absent, inserts an all-true gate into it. Replaying parameters reproduces the output, except for RandomDissolving, whose VAE latents are drawn during application and are not stored.
with static p=1.0, p_batch=1.0 every sample is transformed: a batch_prob put into params is not consulted by the image, the transformation matrix, the labels or inverse, in eager and compiled runs alike (#5585). Pass p < 1 to skip samples by hand, which every one of those paths honors.
some _param_generator.* range buffers in state_dict() do not update the samplers when loaded (#4428).
an empty batch gives an empty output, except on a minority of classes that raise on B = 0 (#4429).
rotation-like parameters are in degrees, and a positive angle turns the image counter-clockwise as displayed (top-left origin, y pointing down), as rotate() documents; RandomAffine turns clockwise (#4408).
torch.jit.script is not supported; torch.compile(fullgraph=True) works for some augmentations and flags only.

apply_transform(input, params, flags, transform=None)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Return type:: Tensor

apply_non_transform(input, params, flags, transform=None)[source]#

Return type:: Tensor

apply_transform_mask(input, params, flags, transform=None)[source]#

Process masks corresponding to the inputs that are transformed.

Return type:: Tensor

apply_non_transform_mask(input, params, flags, transform=None)[source]#

Process masks corresponding to the inputs that are no transformation applied.

Return type:: Tensor

apply_transform_box(input, params, flags, transform=None)[source]#

Process boxes corresponding to the inputs that are transformed.

Return type:: Boxes

apply_non_transform_box(input, params, flags, transform=None)[source]#

Process boxes corresponding to the inputs that are no transformation applied.

Return type:: Boxes

apply_transform_keypoint(input, params, flags, transform=None)[source]#

Process keypoints corresponding to the inputs that are transformed.

Return type:: Keypoints

apply_non_transform_keypoint(input, params, flags, transform=None)[source]#

Process keypoints corresponding to the inputs that are no transformation applied.

Return type:: Keypoints

apply_transform_class(input, params, flags, transform=None)[source]#

Process class tags corresponding to the inputs that are transformed.

Return type:: Tensor

apply_non_transform_class(input, params, flags, transform=None)[source]#

Process class tags corresponding to the inputs that are no transformation applied.

Return type:: Tensor

RigidAffineAugmentationBase2D sits between AugmentationBase2D and the two rigid bases. It adds the transform-matrix machinery — compute_transformation and the transform_matrix attribute — but no inverse.

class kornia.augmentation.RigidAffineAugmentationBase2D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

AugmentationBase2D base class for rigid/affine augmentation implementations.

See the Convention block on AugmentationBase2D.

RigidAffineAugmentationBase2D manages transformation matrices. Further subclasses supply the handlers for different data types like masks, boxes, and keypoints.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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:

a subclass implements compute_transformation(), which returns the (B, 3, 3) matrix of the sampled transform, and apply_transform for images. Direct mask, box, and keypoint handlers are incomplete on this base; inherit the geometric base for its matrix-based handlers. AugmentationSequential calls the subclass’s mask, box and keypoint handlers with transform_matrix as transform, as it does for a geometric child, and raises NotImplementedError for a registered data key whose handler is missing.
the matrix of the last call is readable as transform_matrix. Subclasses opt into lazy construction with _compute_matrix_lazily; a lazy matrix keeps only the input’s shape, dtype and device, unless the subclass overrides transform_tensor, generate_transformation_matrix, compute_transformation or identity_matrix, in which case the input is kept until the matrix is read or replaced.
this base does not implement an inverse operation.

compute_transformation(input, params, flags)[source]#

Return type:: Tensor

The 3D bases provide analogous shape, matrix and data-key machinery; see AugmentationBase3D for their contract, including the missing inverse.

class kornia.augmentation.AugmentationBase3D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

AugmentationBase3D base class for customized augmentation implementations.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
same_on_batch (bool, optional) – apply the same transformation across the batch. Default: False

Convention:

the working layout is (B, C, D, H, W) float. Bare 3D augmentations accept inputs of rank three and four and promote them by prepending batch and, for rank three, channel dimensions; keepdim=True restores that original rank. Inside AugmentationSequential, rank-4 input is rejected because (C, D, H, W) is ambiguous with the container’s (B, C, H, W) image layout. The dtype guard accepts only float16, bfloat16, float32, and float64. RandomTransplantation3D overrides forward and takes batched inputs only.
p gates samples and p_batch gates a whole call; same_on_batch=True shares the generated values. Of the concrete 3D constructors only RandomTransplantation3D exposes p_batch (#4425); CenterCrop3D and RandomCrop3D instead map their own p onto the call-wide gate. Parameters use the common augmentation RNG and forward(x, params=...) replays a complete generated dictionary. See Conventions & Pitfalls for the canonical sampling, seeding, and serialization contract.
rigid subclasses expose the last sampled (B, 4, 4) transform_matrix. The 3D bases do not implement inverse: their subclasses have no inverse method, and a geometric 3D child makes an AugmentationSequential inverse raise. RandomTransplantation3D is the exception: it also derives from MixAugmentationBaseV2 and so carries that class’s inverse, which raises RuntimeError.

class kornia.augmentation.RigidAffineAugmentationBase3D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

AugmentationBase3D base class for rigid/affine augmentation implementations.

RigidAffineAugmentationBase3D enables routined transformation with given transformation matrices for different data types like masks, boxes, and keypoints.

See the Convention block on AugmentationBase3D.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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

class kornia.augmentation.GeometricAugmentationBase3D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

Base class for 3D geometric augmentations.

See the Convention block on AugmentationBase3D.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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:

subclasses record a source-to-destination (B, 4, 4) voxel-coordinate matrix. Coordinates are ordered (x, y, z), while volume tensor axes are (D, H, W). The source frame is the original input, including for crops with pre-padding. Resampling follows the matrix at align_corners=True; at align_corners=False normalization can make the applied warp differ from the recorded matrix (#4503). Matrices use the input dtype, so integer voxel translations are rounded when that dtype cannot represent them exactly.
a zero-parameter RandomAffine3D or RandomRotation3D warp reproduces its input up to roundoff in float32 and float64, and in CPU float16/bfloat16, where the warp runs in float32 (#5751). On other devices the half-precision sampling grid is itself rounded, so the error grows with the volume size.

class kornia.augmentation.IntensityAugmentationBase3D(p=0.5, p_batch=1.0, same_on_batch=False, keepdim=False)[source]#

Base class for 3D intensity augmentations.

See the Convention block on AugmentationBase3D.

Parameters:

p (float, optional) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch. 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
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:

these augmentations leave voxel coordinates in place: they record an identity (B, 4, 4) matrix, and a container skips them on its inverse path.

Mix augmentations derive from MixAugmentationBaseV2 and have their own forward and data-key contracts. Some combine different samples; RandomJigsaw rearranges patches within each image. Label and box handling are class-specific: RandomMixUpV2 accepts class labels, while RandomMosaic accepts boxes but does not implement class-label transforms. RandomTransplantation and RandomTransplantation3D require a segmentation mask; a mask-only call needs data_keys=["mask"]. The 3D transplantation class also inherits AugmentationBase3D.

class kornia.augmentation.MixAugmentationBaseV2(p, p_batch, same_on_batch=False, keepdim=False, data_keys=None)[source]#

MixAugmentationBase base class for customized mix augmentation implementations.

For any augmentation, the implementation of “generate_parameters” and “apply_transform” are required. “apply_transform” will need to handle the probabilities internally.

Parameters:

p (float) – probability for applying an augmentation. This param controls the augmentation probabilities element-wise for a batch.
p_batch (float) – probability for applying an augmentation to a batch. This param controls the augmentation probabilities batch-wise.
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
data_keys (Optional[List[Union[str, int, DataKey]]], optional) – the input type sequential for applying augmentations. Accepts “input”, “image”, “mask”, “bbox”, “bbox_xyxy”, “bbox_xywh”, “keypoints”, “class”, “label”. Default: None

Convention:

Sampling, replay, and serialization follow the augmentation-wide contracts in Conventions & Pitfalls.
Image inputs are floating tensors of shape (H, W), (C, H, W), or (B, C, H, W) in float16, bfloat16, float32 or float64, promoted to (B, C, H, W); keepdim=True restores the rank of an unbatched input. RandomTransplantation and RandomTransplantation3D override forward and take batched inputs only.
A mix operation is not geometric: transform_matrix and inverse() raise RuntimeError.
data_keys chooses which positional inputs are dispatched. Called directly, a key the concrete class does not implement raises NotImplementedError (or a class-specific error) before anything is sampled; the two transplantation classes, and any mix child of AugmentationSequential, raise it after the parameters are drawn. The concrete class blocks state the supported non-image keys.
batch_prob gates the result: a selected image is mixed and an unselected image keeps its input values, except where a concrete class changes the output size (see RandomMosaic); the transplantation classes read acceptor_indices instead. The box handlers use the same gate. The class-label handlers of RandomMixUpV2 and RandomCutMixV2 ignore batch_prob; their sampled gate is batch-wide, but a params whose batch_prob selects only some rows labels the unselected, unchanged rows as mixed (#4775).

Some Further Notes#

Probabilities#

AugmentationBase2D states the p / p_batch model. A concrete constructor can map its public p to either gate: RandomMixUpV2 gates the batch and RandomJigsaw individual samples (#4425). When 0 < p_batch < 1, the base draws the batch Bernoulli before the per-sample gate; endpoints skip that draw, and p=1.0, p_batch=0.0 selects no sample.

With static p=1.0, p_batch=1.0 every sample is transformed: a batch_prob put into params is not consulted by the image, the transformation matrix, the labels or inverse, in eager and compiled runs alike (#5585). Pass p < 1 to skip samples by hand.

Random Generators#

To get an automatically generated __repr__ that lists all custom parameters, implement _param_generator by inheriting from RandomGeneratorBase to generate the random parameters, and put all static parameters inside self.flags. You can take advantage of PlainUniformGenerator to generate simple uniform parameters with less boilerplate code.

Random Reproducibility#

See AugmentationBase2D for where parameters are sampled and what params= replays, and Conventions & Pitfalls for seeding, worker seeds and consumption order.

Serialization#

Several constructors that accept nn.Parameter ranges propagate gradients to them.
A kornia.augmentation.auto policy pickles when every operation wrapper in it does, and a wrapper does when its magnitude mapping is picklable; the built-in mappings are.
AugmentationBase2D covers the range buffers in state_dict(), and RigidAffineAugmentationBase2D what a lazily built matrix keeps.