kornia.contrib

Augmentation Geometry Audit

kornia.contrib.augmentation_audit.audit executes an AugmentationSequential once and returns (outputs, report). It records the executed operation order, native parameter snapshots, image shapes, effective flags, and coordinate matrices. It also checks keypoint and box round trips against the labels actually returned by the pipeline. The ordinary forward path has no reporting overhead.

For example, a training-data validation step can inspect coordinates and export a report without storing image pixels:

import json
from pathlib import Path
import torch
import kornia.augmentation as K
from kornia.contrib import audit

image = torch.rand(2, 3, 32, 48)
points = torch.tensor([[[4., 6.], [20., 15.]]]).expand(2, -1, -1)
pipeline = K.AugmentationSequential(
    K.RandomHorizontalFlip(p=1.0),
    K.RandomAffine(degrees=15, p=1.0),
    data_keys=["input", "keypoints"],
)
(image_out, points_out), report = audit(
    pipeline, image, points, roundtrip_tolerance=1e-3, out_of_frame_tolerance=0.25,
)
print(report.summary())
point_audit = report.spatial[0]
assert report.geometry_status == "available"
assert torch.equal(point_audit.roundtrip_valid_count, point_audit.count)
assert bool((point_audit.roundtrip_max <= 1e-3).all())
Path("augmentation-audit.json").write_text(report.to_json(), encoding="utf-8")
assert json.loads(report.to_json())["image_reconstruction_evaluated"] is False

# Native parameter snapshots replay the run with this pipeline configuration.
replay_image, replay_points = pipeline(image, points, params=report.params)
torch.testing.assert_close(replay_image, image_out)
torch.testing.assert_close(replay_points, points_out)

Interpreting the report

• matrix maps original pixel coordinates to final image coordinates, including RandomCrop prepadding when it belongs to the returned image mapping. inverse_matrix is its algebraic inverse, with nonfinite entries for batches that cannot be inverted. invertible identifies those batches explicitly.

• geometry_status="available" describes matrix availability rather than spatial alignment. Round trips inverse-map the returned labels and compare them with the source. An invertible matrix can therefore accompany a nonzero label error.

• Keypoint errors are Euclidean distances. Box errors are symmetric Hausdorff distances between corner sets. Tensor box exports take axis-aligned envelopes, so rotation or shear can introduce measurable loss; Boxes objects retain the transformed quadrilateral.

• Per-batch statistics include only finite errors. Inspect roundtrip_valid_count together with the mean, median, or maximum. Empty labels and unavailable inverses do not produce a successful zero-error measurement.

• Out-of-frame counts use inclusive pixel centers from zero to width - 1 and height - 1. A box is outside when any corner is outside. Nonfinite labels are counted separately.

• Crop matrices can invert coordinates even after image content has been discarded. The report records this distinction and never evaluates image reconstruction. Cropping and downsampling emit possible content-loss warnings.

• A RandomCrop skipped for the whole batch leaves images and labels unchanged, including with padding or pad_if_needed (#4473). Manually supplied parameters that apply a shape-changing operation to only some rows raise ValueError, as the ordinary forward call does, because the skipped rows cannot keep their shape (#4497). With p=1.0, p_batch=1.0, an overridden batch_prob is not consulted: every image row is cropped and the labels move with it (#5585). Pass p < 1 to skip rows by hand.

• Non-rigid and unknown operations are explicitly unsupported for matrix composition. A transformation-matrix identity fallback is not treated as proof of correspondence. Supported neighboring operations remain in the provenance.

• Captured operations must match the ordered operations selected by the sampled parameters. Missing or out-of-order captures make the geometry unsupported. RandAugment, AutoAugment, TrivialAugment, and custom sequential subclasses are currently opaque and therefore unsupported when selected.

Diagnostics use float32 for half-precision inputs and disable autocast for matrix composition and inverse mapping. The pipeline forward itself inherits the caller’s autocast context, so returned labels remain identical to an ordinary forward call. Autocast can quantize spatial-label propagation and produce real round-trip errors; the report does not suppress those differences.

The API accepts one nonempty BCHW image first, followed by optional batched masks, keypoints, or boxes in the usual positional data_keys forms. Dictionaries, ragged or unbatched inputs, video, patch, and 3D containers, and modules registered under multiple names are unsupported. Repeated sampled operations and nested built-in 2D sequences are supported. Spatial operations must preserve label shapes and cardinality. This is an eager diagnostic utility and the same stateful pipeline must not be audited concurrently.

params is an independent, detached native replay snapshot. JSON is a portable diagnostic export rather than a replay loader. Tensor exports include dtype, device, and shape metadata; nonfinite values become JSON null.

kornia.contrib.audit(sequence, *args, params=None, data_keys=None, roundtrip_tolerance=1e-3, out_of_frame_tolerance=0.0)

Run an augmentation sequence once and return its outputs plus a geometry audit report.

The sequence must use one BCHW image first, followed by batched masks ((B, C, H, W) or (B, H, W)), boxes or keypoints described by positional data_keys. params and data_keys have the same meaning as in kornia.augmentation.AugmentationSequential.forward(). Unknown or non-rigid operations run normally and make composed geometry unavailable. Keypoints and boxes through RandomElasticTransform, RandomThinPlateSpline or RandomFisheye raise NotImplementedError, as in the sequence (#4420).

Reporting is an eager diagnostic operation. It snapshots executed operations and parameters, checks spatial round trips, and never evaluates image reconstruction. The returned outputs retain the ordinary forward structure, values and gradients. Do not audit the same stateful sequence concurrently.

Convention:

Augmentation Geometry Audit states the direction of report.matrix, the pixel-centre rectangle of the out-of-frame test, the dtypes of the diagnostics and the inputs it accepts.

Parameters:

• sequence (AugmentationSequential) – The 2D augmentation sequence to execute and inspect.

• args (Union[Tensor, List[Tensor], Boxes, Keypoints]) – Image and optional spatial inputs accepted by the sequence.

• params (list[ParamItem] | None, optional) – Optional sampled parameters to replay. Default: None

• data_keys (list[str | int | DataKey] | None, optional) – Optional positional data-key override. Default: None

• roundtrip_tolerance (float, optional) – Nonnegative pixel threshold for round-trip warnings. Default: 1e-3

• out_of_frame_tolerance (float, optional) – Fraction in [0, 1] above which an input gets a warning. Default: 0.0

Return type:

tuple[Union[Tensor, List[Tensor], Boxes, Keypoints, list[Union[Tensor, List[Tensor], Boxes, Keypoints]], dict[str, Union[Tensor, List[Tensor], Boxes, Keypoints]]], AugmentationAuditReport]

Returns:

A pair (outputs, report). outputs matches an ordinary forward call.

class kornia.contrib.AugmentationAuditReport(inputs, output_shape, steps, params, configured_extra_args, geometry_status, matrix, inverse_matrix, invertible, spatial, warnings, roundtrip_tolerance, out_of_frame_tolerance, image_reconstruction_evaluated=False)

Geometry provenance and spatial diagnostics from one augmentation call.

params is a detached snapshot of native ParamItem values for replay with the same pipeline configuration. JSON is a diagnostic export, not a parameter deserializer. Reports retain neither image pixels nor autograd graphs. Tensor fields are independent snapshots, not read-only tensors.

geometry_status is unsupported when any executed operation lacks a supported matrix, singular when any composed matrix cannot be inverted, or available. Inspect the per-batch invertible mask and spatial errors rather than interpreting availability as alignment success. inverse_matrix has NaNs for singular/nonfinite batches. No image inverse is evaluated: coordinate invertibility does not imply pixel reconstruction, particularly after cropping, interpolation or intensity changes.

Convention:

See Augmentation Geometry Audit.

summary()

Describe matrix availability and diagnostic warnings without claiming image invertibility.

Return type:

str

to_dict()

Export JSON-safe metadata; tensor values include dtype/device/shape, and NaNs become None.

Return type:

dict[str, Any]

to_json(*, indent=2)

Serialize the report as strict JSON (nonfinite diagnostic values become null).

Return type:

str

class kornia.contrib.AugmentationAuditStep(name, occurrence, module, input_shape, output_shape, params, flags, matrix, unsupported_reason)

Snapshot of one captured leaf operation, including repeated occurrences.

matrix maps this occurrence’s local input pixel coordinates to its output coordinates, including RandomCrop prepadding. The full pipeline mapping is stored separately on the report. None means that the operation has no supported 2D matrix, not that it is an identity transform. Parameters and flags are detached copies; occurrence is zero-based for each qualified module name. Shapes describe the image path. Flags include keyword overrides observed on that image call.

Convention:

See Augmentation Geometry Audit.

class kornia.contrib.SpatialAudit(input_index, data_key, metric, count, out_of_frame, nonfinite, out_of_frame_fraction, roundtrip_valid_count, roundtrip_mean, roundtrip_median, roundtrip_max)

Per-batch diagnostics for one returned box or keypoint input.

Each point or box is one item. A box is outside if any corner is outside the inclusive pixel-center rectangle [0, width-1] x [0, height-1]. Nonfinite items are counted separately, not counted as outside. Fractions use the total item count; empty inputs have a NaN fraction.

Round trips inverse-map the actual returned coordinates with the recorded composed matrix. Errors are Euclidean point distances or symmetric Hausdorff distances between the four box corners (independent of corner ordering). Tensor box exports can lose orientation by taking an axis-aligned envelope; this loss is included in the measurement. Boxes objects retain corners. Statistics cover only finite errors, in pixels; a batch with no valid errors has NaN statistics and zero roundtrip_valid_count. Diagnostics use float32 for half inputs, otherwise the image dtype.

Convention:

See Augmentation Geometry Audit.

Models

Base

class kornia.models.base.ModelBase(*args, **kwargs)

Abstract model class with some utilities function.

compile(*, fullgraph=False, dynamic=False, backend='inductor', mode=None, options=None, disable=False)

Compile this model with torch.compile().

Parameters:

• fullgraph (bool, optional) – Whether Dynamo should require a single full graph. Default: False

• dynamic (bool, optional) – Whether dynamic shape tracing is enabled. Default: False

• backend (str, optional) – Compilation backend name passed to torch.compile(). Default: "inductor"

• mode (Optional[str], optional) – Optional backend-specific compilation mode. Default: None

• options (Optional[dict[Any, Any]], optional) – Optional backend-specific option dictionary. Default: None

• disable (bool, optional) – If True, return an uncompiled model wrapper according to PyTorch’s compile semantics. Default: False

Return type:

ModelBase[TypeVar(ModelConfig)]

Returns:

Compiled model object with the same high-level interface as this instance.

abstractmethod static from_config(config)

Build/load the model.

Parameters:

config (TypeVar(ModelConfig)) – The specifications for the model be build/loaded

Return type:

ModelBase[TypeVar(ModelConfig)]

load_checkpoint(checkpoint, device=None)

Load checkpoint from a given url or file.

Either way the file is loaded with weights_only=True, so it may hold tensors, primitive types and plain containers only; see kornia.core.download.load_state_dict_from_url().

Parameters:

• checkpoint (str | list[str]) – The url or filepath for the respective checkpoint

• device (Optional[device], optional) – The desired device to load the weights and move the model Default: None

Return type:

None

EfficientViT

class kornia.models.efficient_vit.EfficientViT(backbone)

EfficientViT backbone model.

The pretrained backbones were trained on images normalized with the ImageNet mean and standard deviation, so a \([0, 1]\) image goes in as kornia.enhance.normalize(image, torch.tensor([0.485, 0.456, 0.406]), torch.tensor([0.229, 0.224, 0.225])).

__init__(backbone)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(images)

Extract features from the input images.

Parameters:

images (Tensor) – input images tensor of shape (B, C, H, W).

Return type:

dict[str, Tensor]

Returns:

a dictionary with the input under "input", the output of each stage under "stage0", "stage1", and so on, and the last stage’s output again under "stage_final".

static from_config(config)

Build the EfficientViT model from a configuration object.

Parameters:

config (EfficientViTConfig) – EfficientViT configuration object. See EfficientViTConfig.

Returns:

the EfficientViT model, in evaluation mode, since it always loads config.checkpoint (see Pretrained models).

Return type:

EfficientViT

load_checkpoint(checkpoint, device=None)

Load checkpoint from a given url or file.

Either way the file is loaded with weights_only=True, so it may hold tensors, primitive types and plain containers only; see kornia.core.download.load_state_dict_from_url().

Parameters:

• checkpoint (str | list[str]) – The url or filepath for the respective checkpoint

• device (Optional[device], optional) – The desired device to load the weights and move the model Default: None

Return type:

None

class kornia.models.efficient_vit.EfficientViTConfig(checkpoint=<factory>)

Configuration to construct EfficientViT model.

Model weights can be loaded from a checkpoint URL or local path. The model weights are hosted on HuggingFace’s model hub: https://huggingface.co/kornia.

Parameters:

checkpoint (str, optional) – URL or local path of model weights. EfficientViT.from_config() reads the architecture from the file name, which must start with the model type and a - (b0, b1, b2, b3, l0, l1, l2 or l3), as in b1-r224.pt: my_efficientvit.pt raises ValueError. Default: <factory>

checkpoint: str

classmethod from_pretrained(model_type, resolution)

Return a configuration object from a pre-trained model.

Parameters:

• model_type (Literal['b1', 'b2', 'b3']) – model type, one of "b1", "b2", "b3".

• resolution (Literal[224, 256, 288]) – input resolution, one of 224, 256, 288.

Return type:

EfficientViTConfig

Backbones

class kornia.models.efficient_vit.backbone.EfficientViTBackbone(width_list, depth_list, in_channels=3, dim=32, expand_ratio=4, norm='bn2d', act_func='hswish')

Implement the EfficientViT backbone architecture.

EfficientViT is a high-speed vision transformer designed for efficient inference on mobile and edge devices by optimizing the attention mechanism and structural blocks.

Parameters:

• width_list (list[int]) – List of widths for each stage.

• depth_list (list[int]) – List of depths (number of blocks) for each stage.

• in_channels (int, optional) – Number of input image channels. Default: 3.

• dim (int, optional) – Dimension of the query, key, and value tensors in the attention mechanism. Default: 32.

• expand_ratio (float, optional) – Expansion ratio for the MBConv blocks. Default: 4.

• norm (str, optional) – Normalization layer type. Default: “bn2d”.

• act_func (str, optional) – Activation function type. Default: “hswish”.

static build_local_block(in_channels, out_channels, stride, expand_ratio, norm, act_func, fewer_norm=False)

Build the local convolution block used between EfficientViT stages.

Parameters:

• in_channels (int) – Number of input feature channels.

• out_channels (int) – Number of output feature channels.

• stride (int) – Spatial stride for the block.

• expand_ratio (float) – Expansion ratio used by MBConv-style blocks.

• norm (str) – Normalization layer name.

• act_func (str) – Activation function name.

• fewer_norm (bool, optional) – If True, omit selected normalization layers. Default: False

Return type:

Module

Returns:

Depthwise-separable or inverted-bottleneck convolution block.

forward(x)

Run the EfficientViT backbone and collect stage outputs.

Parameters:

x (Tensor) – Image tensor with shape (B, C, H, W).

Return type:

dict[str, Tensor]

Returns:

Dictionary containing the input, each stage output, and "stage_final" for the final feature map.

kornia.models.efficient_vit.backbone.efficientvit_backbone_b0(**kwargs)

Create EfficientViT B0.

Return type:

EfficientViTBackbone

kornia.models.efficient_vit.backbone.efficientvit_backbone_b1(**kwargs)

Create EfficientViT B1.

Return type:

EfficientViTBackbone

kornia.models.efficient_vit.backbone.efficientvit_backbone_b2(**kwargs)

Create EfficientViT B2.

Return type:

EfficientViTBackbone

kornia.models.efficient_vit.backbone.efficientvit_backbone_b3(**kwargs)

Create EfficientViT B3.

Return type:

EfficientViTBackbone

class kornia.models.efficient_vit.backbone.EfficientViTLargeBackbone(width_list, depth_list, in_channels=3, qkv_dim=32, norm='bn2d', act_func='gelu')

Implement the large-scale variant of the EfficientViT backbone.

This backbone is designed for high-resolution dense prediction tasks. It utilizes multi-scale linear attention to achieve a global receptive field while maintaining linear computational complexity relative to the input resolution.

Parameters:

• width_list (list[int]) – List of channel widths for each stage of the backbone.

• depth_list (list[int]) – List of number of blocks for each stage.

• in_channels (int, optional) – Number of input image channels. Default: 3.

• qkv_dim (int, optional) – The internal dimension for query, key, and value projections in the attention layers. Default: 32.

• norm (str, optional) – Normalization layer type to use (e.g., “bn2d”, “ln”). Default: “bn2d”.

• act_func (str, optional) – Activation function type to use (e.g., “gelu”, “relu”). Default: “gelu”.

static build_local_block(stage_id, in_channels, out_channels, stride, expand_ratio, norm, act_func, fewer_norm=False)

Build a local block for an EfficientViT large stage.

Parameters:

• stage_id (int) – Index of the stage being constructed.

• in_channels (int) – Number of input feature channels.

• out_channels (int) – Number of output feature channels.

• stride (int) – Spatial stride for the block.

• expand_ratio (float) – Expansion ratio controlling intermediate channels.

• norm (str) – Normalization layer name.

• act_func (str) – Activation function name.

• fewer_norm (bool, optional) – If True, use the reduced-normalization variant. Default: False

Return type:

Module

Returns:

Residual, fused-MBConv, or MBConv block chosen for the stage.

forward(x)

Run the EfficientViT large backbone and collect stage outputs.

Parameters:

x (Tensor) – Image tensor with shape (B, C, H, W).

Return type:

dict[str, Tensor]

Returns:

Dictionary with entries for each stage and "stage_final" for the final feature map.

kornia.models.efficient_vit.backbone.efficientvit_backbone_l0(**kwargs)

Create EfficientViT L0.

Return type:

EfficientViTLargeBackbone

kornia.models.efficient_vit.backbone.efficientvit_backbone_l1(**kwargs)

Create EfficientViT L.

Return type:

EfficientViTLargeBackbone

kornia.models.efficient_vit.backbone.efficientvit_backbone_l2(**kwargs)

Create EfficientViT L2.

Return type:

EfficientViTLargeBackbone

kornia.models.efficient_vit.backbone.efficientvit_backbone_l3(**kwargs)

Create EfficientViT L3.

Return type:

EfficientViTLargeBackbone

Structures

class kornia.models.structures.SegmentationResults(logits, scores, mask_threshold=0.0, _original_res_logits=None)

Encapsulate the results obtained by a Segmentation model.

Parameters:

• logits (Tensor) – Results logits with shape (K, C, H, W) for one image: \(K\) is the number of its prompts (the leading size of its point, box or mask prompts; 1 when it has none), \(C\) the number of masks predicted per prompt (3 with multimask_output, else 1), and \(H = W = 256\) for SAM.

• scores (Tensor) – The predicted IoU of each mask. Shape (K, C)

• mask_threshold (float, optional) – The threshold value to generate the binary_masks from the logits Default: 0.0

property binary_masks: Tensor

Binary mask generated from logits considering the mask_threshold.

Shape will be the same of logits (K, C, H, W) where \(C\) is the number masks predicted per prompt.

Note

If you run original_res_logits, this will generate the masks based on the original resolution logits. Otherwise, this will use the low resolution logits (self.logits).

logits: Tensor

mask_threshold: float = 0.0

original_res_logits(input_size, original_size, image_size_encoder)

Remove padding and upscale the logits to the original image size.

Resize to image encoder input -> remove padding (bottom and right) -> Resize to original size

Note

This method set a internal original_res_logits which will be used if available for the binary masks.

Parameters:

• input_size (tuple[int, int]) – The size of the image input to the model, in (H, W) format. Used to remove padding.

• original_size (tuple[int, int]) – The original size of the image before resizing for input to the model, in (H, W) format.

• image_size_encoder (Optional[tuple[int, int]]) – The size of the input image for image encoder, in (H, W) format. Used to resize the logits back to encoder resolution before remove the padding.

Return type:

Tensor

Returns:

Batched logits in (K, C, H, W) format, where (H, W) is given by original_size.

scores: Tensor

squeeze(dim=0)

Realize a squeeze for the dim given for all properties.

Return type:

SegmentationResults

class kornia.models.structures.Prompts(points=None, boxes=None, masks=None)

Encapsulate the prompts inputs for a Model.

Parameters:

• points (Optional[tuple[Tensor, Tensor]], optional) – A tuple with the keypoints (coordinates x, y) and their respective labels. Shape (K, N, 2) for the keypoints, and (K, N) Default: None

• boxes (Optional[Tensor], optional) – Batched box inputs, with shape (K, 4). Expected to be into xyxy format. Default: None

• masks (Optional[Tensor], optional) – Batched mask prompts to the model with shape (K, 1, H, W) Default: None

boxes: Optional[Tensor] = None

property keypoints: Tensor | None

The keypoints from the points.

property keypoints_labels: Tensor | None

The keypoints labels from the points.

masks: Optional[Tensor] = None

points: Optional[tuple[Tensor, Tensor]] = None

VisualPrompter

class kornia.contrib.visual_prompter.VisualPrompter(config=None, device=None, dtype=None)

Allow the user to run multiple queries with multiple prompts on one image.

At the moment, we just support the SAM model. The model is loaded based on the given config.

The prompter holds one image at a time: set_image() takes a (3, H, W) or (1, 3, H, W) image, and predict() then answers prompts on it. To segment a batch of images, each with its own prompts, use kornia.models.sam.Sam.forward().

The image is resized so that its long side equals image_encoder.img_size. The output size \((h, w)\) comes from transforms, an AugmentationSequential that must hold exactly one stage, kornia.augmentation.Resize(img_size, side="long", align_corners=False, antialias=True). The image itself is resized as SAM’s reference resizes it (ResizeLongestSide.apply_image_torch), with torch.nn.functional.interpolate(image, size, mode="bilinear", align_corners=False, antialias=True), computed in float32 (float64 for a float64 image) and cast back. The image is then passed to the method preprocess_image, which normalizes it and pads it to the size the SAM model takes, \((\text{image_encoder.img_size}, \text{image_encoder.img_size})\). For default the image is normalized by the mean and standard deviation of the SAM dataset values.

The prompts follow the per-axis scales \(s_x = w / W\) and \(s_y = h / H\) of an \(H \times W\) image, computed in float32 (float64 for a float64 model), and come out in that dtype. A keypoint marks a pixel centre and follows the half-pixel map \(x \mapsto (x + 0.5)\,s_x - 0.5\) (likewise for \(y\)), the map this resize samples on. A box marks pixel edges, kornia’s xyxy convention (width = x2 - x1), and both its corners follow \(x \mapsto x\,s_x\), as in the reference’s apply_boxes: at \(s = 8\) the box (10, 20, 60, 50) goes to (80, 160, 480, 400), and \((0, 0, W, H)\) to the whole resized image \((0, 0, w, h)\).

The image does not go through kornia’s own antialiased resize(), which blurs with a Gaussian of standard deviation \((f - 1) / 2\) before downscaling by \(f\), a weaker low-pass than the antialiased bilinear filter of SAM’s preprocessing. On a 1380 x 2076 image, against a transcription of Meta’s SamPredictor driving the same MobileSAM weights, the three masks of a point prompt reach IoU 0.962, 0.967 and 0.956 with kornia’s filter and 0.969, 0.969 and 0.968 with SAM’s.

Two details depart on purpose from Meta’s SamPredictor (segment-anything@3518c86):

• The resized short side is truncated, where the reference rounds it (a 1380 x 2076 image becomes 680 x 1024 here and 681 x 1024 there). Truncation is the rule of kornia.augmentation.Resize, and the masks are resized back through the same recorded input size.

• Keypoints follow the half-pixel map \(x \mapsto (x + 0.5)\,s - 0.5\), kornia’s pixel-centre convention and the map the image resize samples on, where the reference’s apply_coords uses \(x \mapsto x\,s\) for points too. With the scale \(s = 8\), the keypoint (40, 30) goes to (323.5, 243.5) here and to (320, 240) in the reference.

The held model follows the mode rule of Pretrained models: a config that loads weights gives a model in evaluation mode, and a randomly initialised one stays in training mode, where MobileSAM’s BatchNorm normalises with the statistics of the batch. Call prompter.model.eval() before inference on random weights.

Parameters:

• config (Optional[SamConfig], optional) – A model config to generate the model. Now just the SAM model is supported. Default: None

• device (Optional[device], optional) – The desired device to use the model. Default: None

• dtype (Optional[dtype], optional) – The desired dtype to use the model. Default: None

Raises:

NotImplementedError – If config is neither a SamConfig nor None.

Example

>>> # prompter = VisualPrompter() # Will load the vit h for default
>>> # You can load a custom SAM type for modifying the config
>>> prompter = VisualPrompter(SamConfig('vit_b'))
>>> image = torch.rand(3, 25, 30)
>>> prompter.set_image(image)
>>> # The K box prompts of the image, as a (1, K, 4) xyxy tensor; here K = 1.
>>> boxes = Boxes.from_tensor(
...     torch.tensor([[[0.0, 0.0, 10.0, 10.0]]], device=prompter.device, dtype=torch.float32), mode='xyxy'
... )
>>> prediction = prompter.predict(boxes=boxes)
>>> prediction.logits.shape
torch.Size([1, 3, 256, 256])

compile(*, fullgraph=False, dynamic=False, backend='inductor', mode=None, options=None, disable=False)

Apply torch.compile(…)/dynamo API into the VisualPrompter API.

Note

For more information about the dynamo API check the official docs https://pytorch.org/docs/stable/generated/torch.compile.html

Parameters:

• fullgraph (bool, optional) – Whether it is ok to break model into several subgraphs Default: False

• dynamic (bool, optional) – Use dynamic shape tracing Default: False

• backend (str, optional) – backend to be used Default: "inductor"

• mode (Optional[str], optional) – Can be either “default”, “reduce-overhead” or “max-autotune” Default: None

• options (Optional[dict[Any, Any]], optional) – A dictionary of options to pass to the backend. Default: None

• disable (bool, optional) – Turn torch.compile() into a no-op for testing Default: False

Return type:

None

Example

>>> # prompter = VisualPrompter()
>>> # prompter.compile()
>>> # Use the prompter methods ...

predict(keypoints=None, keypoints_labels=None, boxes=None, masks=None, multimask_output=True, output_original_size=True)

Predict masks for the given image based on the input prompts.

Parameters:

• keypoints (Union[Keypoints, Tensor, None], optional) – Point prompts to the model. Each point is in (X,Y) in pixels. Shape (K, N, 2). Where N is the number of points and K the number of prompts. Default: None

• keypoints_labels (Optional[Tensor], optional) – Labels for the point prompts. 1 indicates a foreground point and 0 indicates a background point. Shape (K, N). Where N is the number of points, and K the number of prompts. Required with keypoints. Default: None

• boxes (Union[Boxes, Tensor, None], optional) – Box prompts to the model, one per prompt: a torch.Tensor in xyxy mode, shape (K, 4) or (1, K, 4), whose corners are taken as given, or a Boxes of shape (K, 4, 2) or (1, K, 4, 2), read through to_tensor(mode="xyxy"). K is the number of prompts. Default: None

• masks (Optional[Tensor], optional) – A low resolution mask input to the model, typically coming from a previous prediction iteration. Has shape (K, 1, H, W), where for SAM, H=W=256. Default: None

• multimask_output (bool, optional) – If true, the model will return three masks. For ambiguous input prompts (such as a single click), this will often produce better masks than a single prediction. If only a single mask is needed, the model’s predicted quality score can be used to select the best mask. For non-ambiguous prompts, such as multiple input prompts, multimask_output=False can give better results. Default: True

• output_original_size (bool, optional) – If true, the logits of SegmentationResults will be post-process to match the original input image size. Default: True

Return type:

SegmentationResults

Returns:

A prediction with the logits and scores (IoU of each predicted mask)

preprocess_image(x, mean=None, std=None)

Normalize and F.pad a torch.Tensor.

For F.normalize the tensor: will prioritize the mean and std passed as argument, if None will use the default Sam Dataset values.

For F.pad the tensor: Will F.pad the torch.Tensor into the right and bottom to match with the size of self.model.image_encoder.img_size

Parameters:

• x (Tensor) – The image to be preprocessed

• mean (Optional[Tensor], optional) – Mean for each channel. Default: None

• std (Optional[Tensor], optional) – Standard deviations for each channel. Default: None

Return type:

Tensor

Returns:

The image preprocessed (normalized if has mean and str available and padded to encoder size)

preprocess_prompts(keypoints=None, keypoints_labels=None, boxes=None, masks=None)

Validate and preprocess the given prompts to be aligned with the input image.

Keypoints follow the half-pixel map \((x + 0.5)\,s - 0.5\) of a pixel centre, and both corners of each xyxy box the edge map \(x\,s\), as the reference’s apply_boxes does, with the per-axis scales recorded by set_image(). They are computed, and returned, in float32 (float64 for a float64 model). A tensor box is taken as given; a Boxes is read through to_tensor(mode="xyxy").

Return type:

Prompts

reset_image()

Clear cached image state and prompt-transform metadata.

This method invalidates previously computed image embeddings and resets all size/transform bookkeeping so a new call to set_image() starts from a clean state.

In practice, this resets: - transformed-image parameters and the prompt scales, - original/input/encoder spatial sizes, - cached image embeddings, - is_image_set status flag.

Return type:

None

set_image(image, mean=None, std=None)

Set the embeddings from the given image with image_decoder of the model.

Resize the image to the size given by transforms, with SAM’s filter (see the class docstring), and prepare it with the preprocess method.

Parameters:

• image (Tensor) – One RGB image, shape (3, H, W) or (1, 3, H, W), in the dtype of the model. Normally images with range of [0-1], the model preprocess F.normalize the pixel values with the mean and std defined in its initialization. A batch of more than one image raises; use kornia.models.sam.Sam.forward() for batches.

• mean (Optional[Tensor], optional) – mean value of dataset for normalization. Default: None

• std (Optional[Tensor], optional) – standard deviation of dataset for normalization. Default: None

Return type:

None

Edge Detection

class kornia.contrib.EdgeDetector(model, pre_processor, post_processor, name=None)

EdgeDetector is a module that wraps an edge detection model.

This is a high-level API that wraps edge detection models like kornia.models.DexiNed. EdgeDetectorBuilder builds one with the input conversion, the network and the sigmoid in model.

Parameters:

• model (Module) – The edge detection model.

• pre_processor (Module) – Pre-processing module (e.g., ResizePreProcessor).

• post_processor (Module) – Post-processing module (e.g., ResizePostProcessor).

• name (Optional[str], optional) – Optional name for the detector. Default: None

Example

>>> detector = EdgeDetectorBuilder.build("dexined", pretrained=True, image_size=352)
>>> img = torch.rand(1, 3, 320, 320)  # RGB in [0, 1]
>>> out = detector(img)  # a list with one (1, 1, 320, 320) edge probability map per image

Face Detection

class kornia.contrib.FaceDetector(top_k=5000, confidence_threshold=0.3, nms_threshold=0.3, keep_top_k=750)

Detect faces in a given image using YuNet model.

Try in browser

This is a high-level API that wraps the kornia.models.YuNet model for face detection. By default, it uses the method described in [FYP+21].

The input is an RGB image with pixel values in \([0, 255]\), not \([0, 1]\). The detector flips it to BGR before YuNet, whose pretrained weights were trained on BGR images in \([0, 255]\).

Parameters:

• top_k (int, optional) – the maximum number of detections to return before the nms. Default: 5000

• confidence_threshold (float, optional) – the threshold used to discard detections. Default: 0.3

• nms_threshold (float, optional) – the threshold used by the nms for iou. Default: 0.3

• keep_top_k (int, optional) – the maximum number of detections to return after the nms. Default: 750

Returns:

A list of B tensors with shape (N,15) to be used with kornia.contrib.FaceDetectorResult.

Example

>>> img = torch.rand(1, 3, 320, 320) * 255  # RGB in [0, 255]
>>> detect = FaceDetector()
>>> res = detect(img)

class kornia.contrib.FaceKeypoint(value)

Define the keypoints detected in a face.

The left/right convention is based on the screen viewer.

EYE_LEFT = 0

EYE_RIGHT = 1

MOUTH_LEFT = 3

MOUTH_RIGHT = 4

NOSE = 2

class kornia.contrib.FaceDetectorResult(data)

Encapsulate the results obtained by the kornia.contrib.FaceDetector.

Parameters:

data (Tensor) – the detections returned by kornia.contrib.FaceDetector, one row of 15 values per face with shape (15,) or (N, 15): the box xmin, ymin, xmax, ymax, five (x, y) landmarks (see kornia.contrib.FaceKeypoint) and the score.

property bottom_left: Tensor

The [x y] position of the top-left coordinate of the bounding box.

property bottom_right: Tensor

The [x y] position of the bottom-right coordinate of the bounding box.

get_keypoint(keypoint)

Get the [x y] position of a given facial keypoint.

Parameters:

keypoint (FaceKeypoint) – the keypoint type to return the position.

Return type:

Tensor

property height: Tensor

The bounding box height.

property score: Tensor

The detection score.

to(device=None, dtype=None)

Like torch.nn.Module.to() method.

Return type:

FaceDetectorResult

property top_left: Tensor

The [x y] position of the top-left coordinate of the bounding box.

property top_right: Tensor

The [x y] position of the top-left coordinate of the bounding box.

property width: Tensor

The bounding box width.

property xmax: Tensor

The bounding box bottom-right x-coordinate.

property xmin: Tensor

The bounding box top-left x-coordinate.

property ymax: Tensor

The bounding box bottom-right y-coordinate.

property ymin: Tensor

The bounding box top-left y-coordinate.

Object Detection

class kornia.contrib.object_detection.BoundingBoxDataFormat(value)

Enum class that maps bounding box data format.

XYWH = 0

XYXY = 1

CXCYWH = 2

CENTER_XYWH = 2

class kornia.contrib.object_detection.BoundingBox(data, data_format)

Bounding box data class.

Useful for representing bounding boxes in different formats for object detection.

Parameters:

• data (tuple[float, float, float, float]) – tuple of bounding box data. The length of the tuple depends on the data format.

• data_format (BoundingBoxDataFormat) – bounding box data format.

data: tuple[float, float, float, float]

data_format: BoundingBoxDataFormat

class kornia.contrib.object_detection.ObjectDetectorResult(class_id, confidence, bbox)

Object detection result.

Parameters:

• class_id (int) – class id of the detected object.

• confidence (float) – confidence score of the detected object.

• bbox (BoundingBox) – bounding box of the detected object in xywh format.

bbox: BoundingBox

class_id: int

confidence: float

class kornia.contrib.object_detection.ObjectDetector(model, pre_processor, post_processor)

Wrap an object detection model and perform pre-processing and post-processing.

__init__(model, pre_processor, post_processor)

Initialize ObjectDetector.

Parameters:

• model (Module) – The object detection model.

• pre_processor (Module) – Pre-processing module (e.g., ResizePreProcessor).

• post_processor (Module) – Post-processing module (e.g., DETRPostProcessor).

compile(*, fullgraph=False, dynamic=False, backend='inductor', mode=None, options=None, disable=False)

Compile the internal object detection model with torch.compile().

Return type:

None

forward(images)

Detect objects in a given list of images.

Parameters:

images (Union[Tensor, list[Tensor]]) – If list of RGB images. Each image is a torch.Tensor with shape (3, H, W). If torch.Tensor, a torch.Tensor with shape (B, 3, H, W).

Return type:

Union[Tensor, list[Tensor]]

Returns:

list of detections found in each image. For item in a batch, shape is (D, 6), where \(D\) is the number of detections in the given image, \(6\) represents class id, score, and xywh bounding box.

static from_config(config)

Build ObjectDetector from config.

This is a placeholder to satisfy the abstract method requirement. Use kornia.contrib.object_detection.RTDETRDetectorBuilder.build() or instantiate ObjectDetector directly.

Parameters:

config (Any) – Configuration object (not used, kept for interface compatibility).

Return type:

ObjectDetector

Returns:

ObjectDetector instance.

name: str = 'detection'

save(images, detections=None, directory=None)

Save the output image(s) to a directory.

Parameters:

• images (Union[Tensor, list[Tensor]]) – input torch.Tensor.

• detections (Optional[Tensor], optional) – detection torch.Tensor. Default: None

• directory (Optional[str], optional) – directory to save the images. Default: None

Return type:

None

to_onnx(onnx_name=None, image_size=640, include_pre_and_post_processor=True, save=True, additional_metadata=None, **kwargs)

Export an RT-DETR object detection model to ONNX format.

Exports this detector, as built by RTDETRDetectorBuilder.build() or the constructor.

Parameters:

• onnx_name (Optional[str], optional) – The name of the output ONNX file. If not provided, "kornia_<name>_<image_size>.onnx", with the detector’s name, for example "kornia_detection_640.onnx" ("kornia_detection_None.onnx" for image_size=None). Default: None

• image_size (Optional[int], optional) – The height and width of the exported graph’s input. If None, both are dynamic. With the pre- and post-processor, the pre-processor resizes the input to the size it was built with, whatever image_size is; without them, the model runs at the size of its input. For RTDETR, recommended scales include [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]. Default: 640

• include_pre_and_post_processor (bool, optional) – Whether to include the pre-processor and post-processor in the exported model. Default: True

• save (bool, optional) – If True, write the model to onnx_name; it is returned either way. Default: True

• additional_metadata (Optional[list[tuple[str, str]]], optional) – Additional metadata to add to the ONNX model. Default: None

• kwargs (Any) – Additional keyword arguments of torch.onnx.export().

Return type:

ModelProto

Note

With a post-processor that returns a list or tuple of per-image detections, as a kornia.models.rt_detr.DETRPostProcessor that filters by confidence does, the model is exported for a batch of one image: no ONNX output holds a list as long as the batch. Which kind of output the detector returns is read from one forward pass for one image at the export size, in eval mode and without gradients, and a tuple as long as the batch from a second one at another batch size, which tells per-image results from a fixed set of batched outputs such as (boxes, scores). A per-image result requires a fixed batch; a tensor result is checked again at the actual tracing batch before accepting a dynamic graph. Build an RT-DETR detector with confidence_filtering=False to export a batch of any size. A dynamic_axes or dynamic_shapes argument that makes the batch of such an export dynamic raises a ValueError after the export, before anything is saved.

Without the pre- and post-processor, a model that declares ONNX_OUTPUT_NAMES, as RTDETR does (pred_logits and pred_boxes), has its outputs exported under those names, each with a dynamic batch; any other model has one output, output.

visualize(images, detections=None, output_type='torch')

Very simple drawing.

Needs to be more fancy later.

Return type:

Union[Tensor, list[Tensor], list[Image]]

class kornia.contrib.object_detection.ResizePreProcessor(height, width, interpolation_mode='bilinear')

Resize a list of image tensors to the given size.

Additionally, also returns the original image sizes for further post-processing.

Parameters:

• height (int) – Height of the resized image.

• width (int) – Width of the resized image.

• interpolation_mode (str, optional) – Interpolation mode for image resizing. Supported values: nearest, bilinear, bicubic, area, nearest-exact. Default: “bilinear”.

Example

>>> import torch
>>> from kornia.models.processors import ResizePreProcessor
>>> processor = ResizePreProcessor(height=224, width=224)
>>> imgs = torch.randn(2, 3, 480, 640)
>>> resized, sizes = processor(imgs)
>>> print(resized.shape, sizes.shape)
torch.Size([2, 3, 224, 224]) torch.Size([2, 2])

forward(imgs)

Resize input images to the target size.

Parameters:

imgs (Union[Tensor, List[Tensor]]) – Input images, either a tensor of shape (B, C, H, W) or a list of tensors of shape (C, H, W).

Returns:

• resized_imgs: Resized images as a tensor of shape \((B, C, H_{\text{new}}, W_{\text{new}})\).

• original_sizes: Original image sizes of shape (B, 2) containing (height, width), as int64 on the images’ device.

Return type:

Tuple containing

kornia.contrib.object_detection.results_from_detections(detections, format)

Convert a detection torch.Tensor to a list of ObjectDetectorResult.

Parameters:

• detections (Tensor) – torch.Tensor with shape (D, 6), where \(D\) is the number of detections in the given image, \(6\) represents class id, score, and xywh bounding box.

• format (str | BoundingBoxDataFormat) – detection format.

Return type:

list[ObjectDetectorResult]

Returns:

list of ObjectDetectorResult.

Real-Time Detection Transformer (RT-DETR)

class kornia.models.rt_detr.RTDETRModelType(value)

Enum class that maps RT-DETR model type.

resnet18d = 0

resnet34d = 1

resnet50d = 2

resnet101d = 3

hgnetv2_l = 4

hgnetv2_x = 5

resnet50d_m = 6

class kornia.models.rt_detr.RTDETRConfig(model_type, num_classes, input_size=640, checkpoint=None, neck_hidden_dim=None, neck_dim_feedforward=None, neck_expansion=None, head_hidden_dim=256, head_num_queries=300, head_num_decoder_layers=None, confidence_threshold=0.3, eval_idx=None)

Configuration to construct RT-DETR model.

Parameters:

• model_type (RTDETRModelType | str | int) –

  model variant. Available models are

  • ResNetD-18: 0, 'resnet18d' or RTDETRModelType.resnet18d

  • ResNetD-34: 1, 'resnet34d' or RTDETRModelType.resnet34d

  • ResNetD-50: 2, 'resnet50d' or RTDETRModelType.resnet50d

  • ResNetD-101: 3, 'resnet101d' or RTDETRModelType.resnet101d

  • HGNetV2-L: 4, 'hgnetv2_l' or RTDETRModelType.hgnetv2_l

  • HGNetV2-X: 5, 'hgnetv2_x' or RTDETRModelType.hgnetv2_x

  • ResNetD-50-M: 6, 'resnet50d_m' or RTDETRModelType.resnet50d_m

• num_classes (int) – number of classes.

• input_size (int, optional) – size of the square image that build() resizes the input to. Default: 640

• checkpoint (Optional[str], optional) – URL or local path of a state dict with kornia’s parameter names, loaded by RTDETR.from_config(). The published RT-DETR checkpoints keep the reference implementation’s names and load only through RTDETR.from_pretrained(). Default: None

• neck_hidden_dim (Optional[int], optional) – hidden dim for neck. Default: None

• neck_dim_feedforward (Optional[int], optional) – feed-forward network dim for neck. Default: None

• neck_expansion (Optional[float], optional) – expansion ratio for neck. Default: None

• head_hidden_dim (int, optional) – hidden dim for head. Default: 256

• head_num_queries (int, optional) – number of queries for Deformable DETR transformer decoder. Default: 300

• head_num_decoder_layers (Optional[int], optional) – number of decoder layers for Deformable DETR transformer decoder. Default: None

• confidence_threshold (float, optional) – score threshold of the post-processor that build() creates from this config when it is given no confidence_threshold. Default: 0.3

• eval_idx (Optional[int], optional) – index of the decoder layer whose predictions the model returns; a negative value counts back from the last layer (-1), and the layers after the selected one do not run. None selects the variant’s layer: index 2, the third of six, for resnet50d_m, as in the reference configs of that checkpoint (or the last layer of a decoder with fewer than three), and the last layer for the other variants. Default: None

checkpoint: Optional[str] = None

confidence_threshold: float = 0.3

eval_idx: Optional[int] = None

static from_name(model_name, num_classes=80)

Load model without pretrained weights.

Parameters:

• model_name (str) – ‘rtdetr_r18vd’, ‘rtdetr_r34vd’, ‘rtdetr_r50vd_m’, ‘rtdetr_r50vd’, ‘rtdetr_r101vd’.

• num_classes (int, optional) – Number of classes to detect. Default: 80

Return type:

RTDETRConfig

head_hidden_dim: int = 256

head_num_decoder_layers: Optional[int] = None

head_num_queries: int = 300

input_size: int = 640

model_type: RTDETRModelType | str | int

neck_dim_feedforward: Optional[int] = None

neck_expansion: Optional[float] = None

neck_hidden_dim: Optional[int] = None

num_classes: int

class kornia.models.rt_detr.RTDETR(backbone, encoder, decoder)

RT-DETR Object Detection model, as described in https://arxiv.org/abs/2304.08069.

__init__(backbone, encoder, decoder)

Construct RT-DETR Object Detection model.

Parameters:

• backbone (ResNetD | PPHGNetV2) – backbone network for feature extraction.

• encoder (HybridEncoder) – neck network for feature fusion.

• decoder (RTDETRHead) – head network to decode features into detection results.

forward(images)

Detect objects in an image.

Parameters:

images (Tensor) – images to be detected. Shape (N, C, H, W).

Return type:

tuple[Tensor, Tensor]

Returns:

• logits - Tensor of shape (N, Q, K), where \(Q\) is the number of queries, \(K\) is the number of classes.

• boxes - Tensor of shape (N, Q, 4), where \(Q\) is the number of queries.

static from_config(config)

Construct RT-DETR Object Detection model from a config object.

Returns the model in evaluation mode when config.checkpoint loads weights and in training mode otherwise (see Pretrained models).

Parameters:

config (RTDETRConfig) – configuration object for RT-DETR.

Return type:

RTDETR

Note

For config.neck_hidden_dim, config.neck_dim_feedforward, config.neck_expansion, and config.head_num_decoder_layers, if they are None, their values will be replaced with the default values depending on the config.model_type. See the source code for the default values.

load_checkpoint(checkpoint, device=None)

Load checkpoint from a given url or file.

Either way the file is loaded with weights_only=True, so it may hold tensors, primitive types and plain containers only; see kornia.core.download.load_state_dict_from_url().

Parameters:

• checkpoint (str | list[str]) – The url or filepath for the respective checkpoint

• device (Optional[device], optional) – The desired device to load the weights and move the model Default: None

Return type:

None

class kornia.models.rt_detr.DETRPostProcessor(confidence_threshold=None, num_classes=80, num_top_queries=300, confidence_filtering=True, filter_as_zero=False)

Convert raw DETR model outputs into final bounding box detections.

This module applies the sigmoid function to the class logits, as RT-DETR was trained with a focal loss, keeps the num_top_queries highest (query, class) scores of each image, and transforms the normalized (cx, cy, w, h) boxes into (x, y, w, h) in the pixel coordinate system of each input image, scaled by that image’s own size.

Parameters:

• num_classes (int, optional) – The number of object classes. Default: 80

• confidence_threshold (Optional[float], optional) – The threshold to filter out low-confidence detections. Default: None

• num_top_queries (int, optional) – The number of top queries to consider for each image. Default: 300

• confidence_filtering (bool, optional) – Whether to apply confidence-based filtering. Default: True

• filter_as_zero (bool, optional) – If True, boxes below the confidence threshold are set to zero instead of being removed. Default: False

forward(logits, boxes, original_sizes)

Post-process outputs from DETR.

Parameters:

• logits (Tensor) – tensor with shape (N, Q, K), where \(N\) is the batch size, \(Q\) is the number of queries, \(K\) is the number of classes.

• boxes (Tensor) – tensor with shape (N, Q, 4), where \(N\) is the batch size, \(Q\) is the number of queries.

• original_sizes (Tensor) – tensor with shape (N, 2), where \(N\) is the batch size and each element represents the image size of (img_height, img_width).

Return type:

Union[Tensor, list[Tensor]]

Returns:

Processed detections, each a row (class_id, confidence_score, x, y, w, h). A tensor of shape (N, T, 6), \(T\) = num_top_queries, when confidence_filtering is False or confidence_threshold is 0; with filter_as_zero, the same tensor with the detections at or below the threshold set to zero. Otherwise a list of \(N\) tensors of shape (D, 6), where \(D\) is the number of detections above the threshold in that image.

Image Segmentation

kornia.contrib.connected_components(image, num_iterations=100)

Compute the Connected-component labelling (CCL) algorithm.

Try in browser

The implementation is an adaptation of the following repository:

https://gist.github.com/efirdc/5d8bd66859e574c683a504a4690ae8bc

Warning

This is an experimental API subject to changes and optimization improvements.

Note

See a working example here.

Convention:

• Components are 8-connected: pixels that touch at a corner are connected. Choosing a labeling method maps this onto SciPy, OpenCV and scikit-image.

• Foreground is a value exactly equal to 1: 0.5 and 2 are background. Other values and dtypes are accepted.

• The label of a component is 1 plus the largest raster index of its pixels in the batch flattened to (N, H, W); background is 0. Labels are therefore not contiguous, and an image gets other labels at another batch position or transposed. On MPS an image of more than \(2^{24}\) pixels raises ValueError.

• num_iterations is a fixed number of max-pooling steps, not a convergence test, and a label travels one pixel per step: the default 100 already splits a 1 x 102 line, and a filled 120 x 150 rectangle comes out with 1000 labels. Choosing a labeling method compares this with connected_components_union_find().

• The labels carry no gradient.

Parameters:

• image (Tensor) – the binary input image with shape (*, 1, H, W).

• num_iterations (int, optional) – the number of max-pooling steps, at least 1. Default: 100

Return type:

Tensor

Returns:

int64 labels with the shape of the input.

Example

>>> img = torch.rand(2, 1, 4, 5)
>>> img_labels = connected_components(img, num_iterations=100)

kornia.contrib.connected_components_union_find(image)

Label 8-connected foreground components exactly using block union-find.

This pure PyTorch implementation contracts each 2x2 block, following the block-based approach of Allegretti, Bolelli and Grana, Optimized Block-Based Algorithms to Label Connected Components on GPUs, TPDS 2020, https://doi.org/10.1109/TPDS.2019.2934683. Root hooking with scatter_reduce_ and pointer jumping replace the paper’s custom CUDA union kernels.

Unlike connected_components(), this function runs until convergence. There is no iteration budget to choose based on component diameter.

Convention:

Choosing a labeling method maps its 8-connectivity onto SciPy, OpenCV and scikit-image, and compares it with connected_components().

Parameters:

image (Tensor) – Binary image of shape (*, 1, H, W). Boolean, integer and floating-point tensors are accepted. Only values equal to 1 are foreground, as in connected_components().

Return type:

Tensor

Returns:

A torch.int64 tensor on the input device with the same shape. Background is zero. Foreground labels are the minimum 1-based block index in each component, in raster order over flattened batches and the \(\lceil H/2\rceil \times \lceil W/2\rceil\) block grid. Labels are deterministic and distinct across images, but not contiguous and not equal to the labels of the pooling implementation.

Note

This discrete operation is not differentiable. Working memory is linear in the number of pixels. Convergence checks synchronize the device with the host; CUDA graph capture and torch.compile(fullgraph=True) are not supported. All mask, edge and label computation stays on the input device. The existing pooling API remains useful for fixed-work graphs.

Example

>>> mask = torch.tensor([[[1, 0, 0], [0, 0, 1]]], dtype=torch.bool)
>>> connected_components_union_find(mask)
tensor([[[1, 0, 0],
         [0, 0, 2]]])

Choosing a labeling method

Both functions label 8-connected components: two foreground pixels (pixels equal to 1) that touch at an edge or a corner belong to one component. SciPy, OpenCV and scikit-image take any non-zero pixel as foreground, so the same partition is scipy.ndimage.label(mask == 1, structure=np.ones((3, 3))) (SciPy’s default structure is 4-connected), cv2.connectedComponents((mask == 1).astype(np.uint8), connectivity=8) (OpenCV’s default connectivity) and skimage.measure.label(mask == 1, connectivity=2) (scikit-image’s default for 2-D images). Label values differ between these libraries and between the two functions: compare foreground partitions, not label values.

connected_components_union_find computes the exact 8-connected partition without an iteration budget. It accepts boolean masks and returns int64 labels, so half-precision inputs do not cause distinct component IDs to collide. It is useful for segmentation masks containing long structures such as vessels, roads and text strokes. Its convergence checks synchronize the accelerator with the host, and the result is not differentiable.

connected_components retains its fixed-count pooling implementation; it returns int64 labels like connected_components_union_find but keeps its own label convention. When using it, choose enough iterations for the largest component’s graph diameter: the default 100 steps can leave a long or large component split into multiple labels. A small fixed budget may be faster than union-find for small or fragmented masks, but does not guarantee convergence.

Segment Anything (SAM)

class kornia.models.sam.SamModelType(value)

Map the SAM model types.

vit_h = 0

vit_l = 1

vit_b = 2

mobile_sam = 3

class kornia.models.sam.SamConfig(model_type=None, checkpoint=None, pretrained=False, encoder_embed_dim=None, encoder_depth=None, encoder_num_heads=None, encoder_global_attn_indexes=None)

Encapsulate the Config to build a SAM model.

Parameters:

• model_type (Union[str, int, SamModelType, None], optional) –

  the available models are: Default: None

  • 0, ‘vit_h’ or kornia.models.sam.SamModelType.vit_h

  • 1, ‘vit_l’ or kornia.models.sam.SamModelType.vit_l

  • 2, ‘vit_b’ or kornia.models.sam.SamModelType.vit_b

  • 3, ‘mobile_sam’, or kornia.models.sam.SamModelType.mobile_sam

  With None, the image encoder is a ViT built from the four encoder_* fields below, which are ignored when model_type is set.

• checkpoint (Optional[str], optional) – URL or a path for a file with the weights of the model Default: None

• pretrained (bool, optional) – load the released weights of model_type, which must then be set. A checkpoint takes precedence, with a warning. With False (the default) and no checkpoint, the weights are randomly initialised. Default: False

• encoder_embed_dim (Optional[int], optional) – Patch embedding dimension. Default: None

• encoder_depth (Optional[int], optional) – Depth of ViT. Default: None

• encoder_num_heads (Optional[int], optional) – Number of attention heads in each ViT block. Default: None

• encoder_global_attn_indexes (Union[int, tuple[int, ...], list[int], None], optional) – Encoder indexes for blocks using global attention: a tuple or a list of indexes, or one index as an int. Default: None

checkpoint: Optional[str] = None

encoder_depth: Optional[int] = None

encoder_embed_dim: Optional[int] = None

encoder_global_attn_indexes: Union[int, tuple[int, ...], list[int], None] = None

encoder_num_heads: Optional[int] = None

model_type: Union[str, int, SamModelType, None] = None

pretrained: bool = False

class kornia.models.sam.Sam(image_encoder, prompt_encoder, mask_decoder)

Implement the Segment Anything Model (SAM) wrapper.

This class coordinates the image encoder, prompt encoder, and mask decoder.

__init__(image_encoder, prompt_encoder, mask_decoder)

SAM predicts object masks from an image and input prompts.

Parameters:

• image_encoder (ImageEncoderViT | TinyViT) – The backbone used to encode the image into image embeddings that allow for efficient mask prediction.

• prompt_encoder (PromptEncoder) – Encodes various types of input prompts.

• mask_decoder (MaskDecoder) – Predicts masks from the image embeddings and encoded prompts.

forward(images, batched_prompts, multimask_output)

Predicts masks end-to-end from provided images and prompts.

This method expects that the images have already been pre-processed, at least been normalized, resized and padded to be compatible with the self.image_encoder.

Note

For each image (3, H, W), it is possible to input a batch (\(K\)) of \(N\) prompts, the results are batched by the number of prompts batch. So given a prompt with \(K=5\), and \(N=10\), the results will look like \(5xCxHxW\) where \(C\) is determined by multimask_output. And within each of these masks \((5xC)\), it should be possible to find \(N\) instances if the model succeed.

Parameters:

• images (Tensor) – The image as a torch tensor in (B, 3, H, W) format, already transformed for input to the model.

• batched_prompts (list[dict[str, Any]]) –

  A list over the batch of images (list length should be \(B\)), each a dictionary with

  the following keys. If it does not have the respective prompt, it should not be included in this dictionary. The options are:

  • ”points”: tuple of (torch.Tensor, torch.Tensor) within the coordinate keypoints and their respective labels. The tuple should look like (keypoints, labels), where the keypoints (a tensor) are a batched point prompts for this image, with shape (K, N, 2). Already transformed to the input frame of the model. The labels (a tensor) are a batched labels for point prompts, with shape (K, N). Where 1 indicates a foreground point and 0 indicates a background point.

  • ”boxes”: (torch.Tensor) Batched box inputs, with shape (K, 4). Already transformed to the input frame of the model.

  • ”mask_inputs”: (torch.Tensor) Batched mask inputs to the model, in the form (K, 1, H, W).

• multimask_output (bool) – Whether the model should predict multiple disambiguating masks, or return a single mask.

Returns:

• logits: Low resolution logits with shape (K, C, H, W). Can be passed as mask input to subsequent iterations of prediction. Where \(K\) is the number of input prompts, \(C\) is determined by multimask_output, and \(H=W=256\) are the model output size.

• scores: The predicted IoU of each mask, (K, C).

Return type:

A list over input images, where each element is as SegmentationResults the following

static from_config(config)

Build/load the SAM model based on it’s config.

Parameters:

config (SamConfig) – The SamConfig data structure. If the model_type is available, build from it, otherwise will use the parameters set.

Return type:

Sam

Returns:

The respective SAM model, in evaluation mode when it loads weights (config.pretrained or config.checkpoint) and in training mode otherwise (see Pretrained models).

Example

>>> from kornia.models.sam import SamConfig
>>> sam_model = Sam.from_config(SamConfig('vit_b'))

load_checkpoint(checkpoint, device=None)

Load checkpoint from a given url or file.

Either way the file is loaded with weights_only=True, so it may hold tensors, primitive types and plain containers only; see kornia.core.download.load_state_dict_from_url().

Parameters:

• checkpoint (str | list[str]) – The url or filepath for the respective checkpoint

• device (Optional[device], optional) – The desired device to load the weights and move the model Default: None

Return type:

None

Image Patches

kornia.contrib.compute_padding(original_size, window_size, stride=None)

Compute required padding to ensure chaining of extract_tensor_patches() and combine_tensor_patches() produces expected result.

Parameters:

• original_size (Union[int, Tuple[int, int]]) – the size of the original torch.Tensor.

• window_size (Union[int, Tuple[int, int]]) – the size of the sliding window used while extracting patches.

• stride (Union[int, Tuple[int, int], None], optional) – The stride of the sliding window. Optional: if not specified, window_size will be used. Default: None

Returns:

(top, bottom, left, right)

Return type:

The required padding as a tuple of four ints

Convention:

See the Convention block of ExtractTensorPatches, which states where the padding goes.

Example

>>> image = torch.arange(12).view(1, 1, 4, 3)
>>> padding = compute_padding((4,3), (3,3))
>>> out = extract_tensor_patches(image, window_size=(3, 3), stride=(3, 3), padding=padding)
>>> combine_tensor_patches(out, original_size=(4, 3), window_size=(3, 3), stride=(3, 3), unpadding=padding)
tensor([[[[ 0,  1,  2],
          [ 3,  4,  5],
          [ 6,  7,  8],
          [ 9, 10, 11]]]])

Note

This function will be implicitly used in extract_tensor_patches() and combine_tensor_patches() if allow_auto_(un)padding is set to True and no non-zero (un)padding is given.

kornia.contrib.extract_tensor_patches(input, window_size, stride=1, padding=0, allow_auto_padding=False)

Extract patches from tensors and stacks them.

Convention:

See the Convention block of ExtractTensorPatches.

Parameters:

• input (Tensor) – torch.Tensor image where to extract the patches with shape (B, C, H, W).

• window_size (Union[int, Tuple[int, int]]) – the size of the sliding window and the output patch size.

• stride (Union[int, Tuple[int, int]], optional) – stride of the sliding window. Default: 1

• padding (Union[int, Tuple[int, int], Tuple[int, int, int, int]], optional) – zero padding, or cropping where negative: an int, (ph, pw) or (top, bottom, left, right). Default: 0

• allow_auto_padding (bool, optional) – whether to allow automatic padding if the window and stride do not fit into the image. Default: False

Return type:

Tensor

Returns:

the torch.Tensor with the extracted patches with shape \((B, N, C, H_{out}, W_{out})\).

Examples

>>> input = torch.arange(9.).view(1, 1, 3, 3)
>>> patches = extract_tensor_patches(input, (2, 3))
>>> input
tensor([[[[0., 1., 2.],
          [3., 4., 5.],
          [6., 7., 8.]]]])
>>> patches[:, -1]
tensor([[[[3., 4., 5.],
          [6., 7., 8.]]]])

kornia.contrib.combine_tensor_patches(patches, original_size, window_size, stride, allow_auto_unpadding=False, unpadding=0, eps=None)

Restore input from patches.

Convention:

See the Convention block of CombineTensorPatches.

Parameters:

• patches (Tensor) – patched torch.Tensor with shape \((B, N, C, H_{out}, W_{out})\).

• original_size (Union[int, Tuple[int, int]]) – the size of the original torch.Tensor and the output size.

• window_size (Union[int, Tuple[int, int]]) – the size of the sliding window used while extracting patches.

• stride (Union[int, Tuple[int, int]]) – stride of the sliding window. Required.

• unpadding (Union[int, Tuple[int, int], Tuple[int, int, int, int]], optional) – the padding used while extracting the patches, removed from the output. Default: 0

• allow_auto_unpadding (bool, optional) – whether to allow automatic unpadding of the input if the window and stride do not fit into the original_size. Default: False

• eps (Optional[float], optional) – deprecated, has no effect. Each covered pixel is divided by its exact overlap count and a pixel no patch covers is 0. Default: None

Return type:

Tensor

Returns:

The combined patches in an image torch.Tensor with shape (B, C, H, W).

Example

>>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2))
>>> combine_tensor_patches(out, original_size=(4, 4), window_size=(2, 2), stride=(2, 2))
tensor([[[[ 0,  1,  2,  3],
          [ 4,  5,  6,  7],
          [ 8,  9, 10, 11],
          [12, 13, 14, 15]]]])

Note

This function is supposed to be used in conjunction with extract_tensor_patches().

class kornia.contrib.ExtractTensorPatches(window_size, stride=1, padding=0, allow_auto_padding=False)

nn.Module that extracts patches from tensors and stacks them.

In the simplest case, the output value of the operator with input size (B, C, H, W) is \((B, N, C, H_{out}, W_{out})\).

where

• \(B\) is the batch size.

• \(N\) denotes the total number of extracted patches.

• \(C\) denotes the number of input channels.

• \(H\), \(W\) the input height and width of the input in pixels.

• \(H_{out}\), \(W_{out}\) denote the patch size defined in the function signature.

An empty batch (\(B = 0\)) gives an empty output with the same \(N\), \(C\), \(H_{out}\) and \(W_{out}\) as a non-empty one. \(C\) must be non-zero and the window sizes and strides positive.

• window_size is the size of the sliding window and controls the shape of the output torch.Tensor and defines the shape of the output patch.

• stride controls the stride to apply to the sliding window and regulates the overlapping between the extracted patches.

• padding controls the amount of implicit zero padding on each side of each dimension.

• allow_auto_padding allows automatic calculation of the padding required to fit the window and stride into the image.

The parameters window_size, stride and padding can be either:

• a single int – in which case the same value is used for the height and width dimension.

• a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension.

padding can also be a tuple of four ints – in which case, the first two ints are for the height dimension while the last two ints are for the width dimension.

Convention:

• Patches are stacked row-major over the grid of window positions, left to right and then top to bottom, the order of torch.nn.functional.unfold().

• A positive padding adds zeros and a negative one crops. A 4-tuple is (top, bottom, left, right), not the (left, right, top, bottom) of torch.nn.functional.pad(); a 2-tuple (ph, pw) adds ph rows above and below and pw columns left and right.

• A window that does not fit warns and drops the last rows or columns; an image smaller than the window warns and then raises. allow_auto_padding=True pads instead, as compute_padding() computes: split evenly with the odd pixel at the bottom or right, and an image smaller than the window padded to one window.

• A non-zero padding switches off both the fit warning and allow_auto_padding, so a padding too small for the window drops the last rows or columns silently. An all-zero padding counts as none.

• A window_size or stride component below 1 raises ValueCheckError, a ValueError, up front (at construction for the module classes), here and in combine_tensor_patches(), compute_padding() and CombineTensorPatches, also under python -O and disable_checks().

Parameters:

• input – torch.Tensor image where to extract the patches with shape (B, C, H, W).

• window_size (Union[int, Tuple[int, int]]) – the size of the sliding window and the output patch size.

• stride (Union[int, Tuple[int, int]], optional) – stride of the sliding window. Default: 1

• padding (Union[int, Tuple[int, int], Tuple[int, int, int, int]], optional) – zero padding, or cropping where negative: an int, (ph, pw) or (top, bottom, left, right). Default: 0

• allow_auto_padding (bool, optional) – whether to allow automatic padding if the window and stride do not fit into the image. Default: False

Shape:

• Input: (B, C, H, W)

• Output: \((B, N, C, H_{out}, W_{out})\)

Returns:

the torch.Tensor with the extracted patches.

Examples

>>> input = torch.arange(9.).view(1, 1, 3, 3)
>>> patches = extract_tensor_patches(input, (2, 3))
>>> input
tensor([[[[0., 1., 2.],
          [3., 4., 5.],
          [6., 7., 8.]]]])
>>> patches[:, -1]
tensor([[[[3., 4., 5.],
          [6., 7., 8.]]]])

class kornia.contrib.CombineTensorPatches(original_size, window_size, stride=None, unpadding=0, allow_auto_unpadding=False)

nn.Module that combines patches back into full tensors.

In the simplest case, the output value of the operator with input size \((B, N, C, H_{out}, W_{out})\) is (B, C, H, W).

where

• \(B\) is the batch size.

• \(N\) denotes the total number of extracted patches.

• \(C\) denotes the number of input channels.

• \(H\), \(W\) the height and width of the output image in pixels.

• \(H_{out}\), \(W_{out}\) denote the patch size defined in the function signature.

An empty batch (\(B = 0\)) is reconstructed with the same spatial size and unpadding as a non-empty one.

• original_size is the size of the original image prior to extracting torch.Tensor patches, which is the size of the output image.

• window_size is the size of the sliding window used while extracting torch.Tensor patches.

• stride controls the stride to apply to the sliding window and regulates the overlapping between the extracted patches.

• unpadding is the amount of padding to be removed. If specified, this value must be the same as padding used while extracting torch.Tensor patches.

• allow_auto_unpadding allows automatic calculation of the padding required to fit the window and stride into the image. This must be used if the allow_auto_padding flag was used for extracting the patches.

The parameters original_size, window_size, stride, and unpadding can be either:

• a single int – in which case the same value is used for the height and width dimension.

• a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension.

unpadding can also be a tuple of four ints – in which case, the first two ints are for the height dimension while the last two ints are for the width dimension.

Convention:

• Overlapping patches are averaged: each pixel is divided by the number of patches that cover it, where torch.nn.functional.fold() sums them. A pixel that no patch covers (patches that do not fit) is 0.

• The output keeps the dtype of the patches. Integer patches are averaged exactly in int64, with the mean truncated toward zero (on MPS before torch 2.9 they go through float32, exact up to \(2^{24}\)); bool patches are True where any covering patch is True; float16 and bfloat16 patches are summed in float32.

• A stride larger than window_size raises, because the patches leave pixels uncovered.

• unpadding and allow_auto_unpadding follow the padding and allow_auto_padding rules of ExtractTensorPatches; pass the values used to extract the patches.

• stride defaults to window_size here and in compute_padding(), to 1 in ExtractTensorPatches and extract_tensor_patches(), and has no default in combine_tensor_patches(), which also takes allow_auto_unpadding before unpadding.

Parameters:

• patches – patched torch.Tensor with shape \((B, N, C, H_{out}, W_{out})\).

• original_size (Tuple[int, int]) – the size of the original torch.Tensor and the output size.

• window_size (Union[int, Tuple[int, int]]) – the size of the sliding window used while extracting patches.

• stride (Union[int, Tuple[int, int], None], optional) – stride of the sliding window. Default: window_size.

• unpadding (Union[int, Tuple[int, int], Tuple[int, int, int, int]], optional) – the padding used while extracting the patches, removed from the output. Default: 0

• allow_auto_unpadding (bool, optional) – whether to allow automatic unpadding of the input if the window and stride do not fit into the original_size. Default: False

Shape:

• Input: \((B, N, C, H_{out}, W_{out})\)

• Output: (B, C, H, W)

Example

>>> out = extract_tensor_patches(torch.arange(16).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2))
>>> combine_tensor_patches(out, original_size=(4, 4), window_size=(2, 2), stride=(2, 2))
tensor([[[[ 0,  1,  2,  3],
          [ 4,  5,  6,  7],
          [ 8,  9, 10, 11],
          [12, 13, 14, 15]]]])

Note

This function is supposed to be used in conjunction with ExtractTensorPatches.

Histogram Matching

kornia.contrib.histogram_matching(source, template)

Adjust the pixel values of an image to match its histogram towards a target image.

Histogram matching is the transformation of an image so that its histogram matches a specified histogram. Code referred to here.

Convention:

• One histogram of every element of source, all batch items and channels together, is matched to one histogram of every element of template, which can have any shape. This is skimage.exposure.match_histograms(source, template, channel_axis=None), which needs both inputs to have the same number of dimensions (flatten them otherwise); to match each sample or channel on its own, call the function on each slice.

• A source value maps to the template value at its quantile, interpolated linearly between template values with interp(). The result stays inside the template’s range once that range is rounded to the output dtype, so with a float64 template and a float32 source it can pass the unrounded ends by that rounding. It can also fall between template values: a uint8 template gives fractional values, as in skimage.

• A floating source keeps its dtype; an integer or bool source returns float32. The quantiles are computed in float32 (float64 for a float64 source).

• It is not differentiable: with a non-empty source, no gradient reaches it and a backward pass into template raises.

• An empty source with a non-empty template returns an empty tensor, still attached to its autograd graph; an empty template raises ValueError, even when source is empty too.

Parameters:

• source (Tensor) – Image to transform.

• template (Tensor) – Template image. It can have different dimensions to source.

Return type:

Tensor

Returns:

The transformed output image, with the shape of the source image.

kornia.contrib.interp(x, xp, fp)

One-dimensional linear interpolation for monotonically increasing sample points.

Returns the one-dimensional piecewise linear interpolant to a function with given discrete data points \((xp, fp)\), evaluated at \(x\).

Convention:

• Outside [xp[0], xp[-1]] the end segments are extrapolated linearly, where numpy.interp() clamps to fp[0] and fp[-1]; SciPy’s interp1d(xp, fp, fill_value="extrapolate") makes the same choice. At a knot the result is that knot’s fp.

• xp and fp are 1-D; x can have any shape, which the output keeps. xp must be increasing, which is not checked. The output dtype is the promoted dtype of x, xp and fp, or the default float dtype when all three are integer.

• A float16 or bfloat16 result is computed in float32 and cast to the output dtype once, so coordinate differences above 65504 do not overflow the weight and a segment whose fp difference overflows the output dtype stays finite. Inside a segment the result is within about half an ulp of the segment’s larger end value, and the knots are exact. float32 and float64 results are computed in their own dtype. An integer x or xp mixed with floating coordinates is converted to the dtype the result is computed in, so in float32 it is exact only up to \(2^{24}\): larger integer knots can merge and larger integer queries are rounded.

Parameters:

• x (Tensor) – the \(x\)-coordinates at which to evaluate the interpolated values.

• xp (Tensor) – the \(x\)-coordinates of the data points, must be increasing.

• fp (Tensor) – the \(y\)-coordinates of the data points, same length as xp.

Return type:

Tensor

Returns:

the interpolated values, same size as x.

Image Classification

class kornia.models.vit.VisionTransformer(image_size=224, patch_size=16, in_channels=3, embed_dim=768, depth=12, num_heads=12, dropout_rate=0.0, dropout_attn=0.0, backbone=None)

Vision transformer (ViT) module.

The module is expected to be used as operator for different vision tasks.

The method is inspired from existing implementations of the paper [DBK+21].

Without a backbone, the input must be image_size square with in_channels channels. With one, the backbone’s token grid must be the one it gives for an image_size input, which the position embeddings are built for; a backbone that ends in adaptive pooling, for example, gives that grid at any input size. The AugReg weights that from_config() loads were trained on images normalized with mean = std = 0.5 per channel, so a \([0, 1]\) image goes in as kornia.enhance.normalize(image, torch.tensor([0.5, 0.5, 0.5]), torch.tensor([0.5, 0.5, 0.5])).

Warning

This is an experimental API subject to changes in favor of flexibility.

Parameters:

• image_size (int, optional) – the size of the input image. Default: 224

• patch_size (int, optional) – the size of the patch to compute the embedding. Default: 16

• in_channels (int, optional) – the number of channels for the input. Default: 3

• embed_dim (int, optional) – the embedding dimension inside the transformer encoder. Default: 768

• depth (int, optional) – the depth of the transformer. Default: 12

• num_heads (int, optional) – the number of attention heads. Default: 12

• dropout_rate (float, optional) – dropout rate. Default: 0.0

• dropout_attn (float, optional) – attention dropout rate. Default: 0.0

• backbone (Module | None, optional) – an nn.Module to compute the image patches embeddings, as a \((B, D, h, w)\) map. The position embeddings are sized from its output for an image_size input. Default: None

Example

>>> img = torch.rand(1, 3, 224, 224)
>>> vit = VisionTransformer(image_size=224, patch_size=16)
>>> vit(img).shape
torch.Size([1, 197, 768])

property encoder_results: list[Tensor]

Return intermediate outputs captured by the transformer encoder.

Returns:

List of tensors produced by the encoder blocks. Each tensor stores token embeddings for a layer, typically shaped (B, N, D), where \(B\) is batch size, \(N\) is token count, and \(D\) is embedding dimension.

forward(x)

Encode an image batch into Vision Transformer token embeddings.

Parameters:

x (Tensor) – Image tensor with shape (B, C, H, W). Without a backbone, \(C\) must equal self.in_channels and \(H\) and \(W\) must both equal self.image_size; with one, the backbone’s token grid must be the one it gives for an image_size input. The position embeddings are a fixed table, one row per token of that grid.

Return type:

Tensor

Returns:

Normalized token embedding tensor produced by patch embedding and the transformer encoder. The output shape follows the encoder layout, usually (B, N, D).

Raises:

ValueError – without a backbone, if \(C\) is not in_channels or (H, W) is not (image_size, image_size); with one, if the backbone’s token grid is not the one it gives for an image_size input.

static from_config(variant, pretrained=False, **kwargs)

Build ViT model based on the given config string.

The format is vit_{size}/{patch_size}. E.g. vit_b/16 means ViT-Base, patch size 16x16. If pretrained=True, AugReg weights are loaded. The weights are hosted on HuggingFace’s model hub: https://huggingface.co/kornia.

Note

The available weights are: vit_l/16, vit_b/16, vit_s/16, vit_ti/16, vit_b/32, vit_s/32.

Parameters:

• variant (str) – ViT model variant e.g. vit_b/16.

• pretrained (bool, optional) – whether to load pre-trained AugReg weights. Default: False

• kwargs (Any) – other keyword arguments that will be passed to kornia.models.vit.VisionTransformer().

Return type:

VisionTransformer

Returns:

The respective ViT model, in evaluation mode when pretrained=True loads weights and in training mode otherwise (see Pretrained models).

Example

>>> from kornia.models.vit import VisionTransformer
>>> vit_model = VisionTransformer.from_config("vit_b/16", pretrained=True)

class kornia.models.vit_mobile.MobileViT(mode='xxs', in_channels=3, patch_size=(2, 2), dropout=0.0)

Module MobileViT. Default arguments is for MobileViT XXS.

Paper: https://arxiv.org/abs/2110.02178 Based on: chinhsuanwu/mobilevit-pytorch

Parameters:

• mode (str, optional) – ‘xxs’, ‘xs’ or ‘s’, defaults to ‘xxs’. Default: "xxs"

• in_channels (int, optional) – the number of channels for the input image. Default: 3

• patch_size (Tuple[int, int], optional) – the patch size \((p_h, p_w)\) the MobileViT blocks unfold their feature maps into, at strides 8, 16 and 32 of the input. Each of those maps must split into whole patches, which an input height and width that are multiples of \(32 p_h\) and \(32 p_w\) guarantee: with the default (2, 2), 192, 256 and 320 px run, and 224 and 288 px raise. Default: (2, 2)

• dropout (float, optional) – dropout ratio in Transformer. Default: 0.0

Raises:

ValueError – if mode is not 'xxs', 'xs' or 's'.

Example

>>> img = torch.rand(1, 3, 256, 256)
>>> mvit = MobileViT(mode='xxs')
>>> mvit(img).shape
torch.Size([1, 320, 8, 8])

class kornia.contrib.TinyViT(img_size=224, in_chans=3, num_classes=1000, embed_dims=(96, 192, 384, 768), depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24), window_sizes=(7, 7, 14, 7), mlp_ratio=4.0, drop_rate=0.0, drop_path_rate=0.0, use_checkpoint=False, mbconv_expand_ratio=4.0, local_conv_size=3, activation=nn.GELU, mobile_sam=False)

TinyViT model, as described in https://arxiv.org/abs/2207.10666.

The pretrained checkpoints of from_config() were trained on images normalized with the ImageNet mean and standard deviation, so a \([0, 1]\) image goes in as kornia.enhance.normalize(image, torch.tensor([0.485, 0.456, 0.406]), torch.tensor([0.229, 0.224, 0.225])).

Parameters:

• img_size (int, optional) – Size of the square input image, a multiple of 32 (of 16 with mobile_sam=True). The attention stages are built for the token grid of an img_size image (see forward()). Default: 224

• in_chans (int, optional) – Number of input image’s channels. Default: 3

• num_classes (int, optional) – Number of output classes. Default: 1000

• embed_dims (Sequence[int], optional) – List of embedding dimensions. Default: (96, 192, 384, 768)

• depths (Sequence[int], optional) – List of block count for each downsampling stage Default: (2, 2, 6, 2)

• num_heads (Sequence[int], optional) – List of attention heads used in self-attention for each downsampling stage. Default: (3, 6, 12, 24)

• window_sizes (Sequence[int], optional) – List of self-attention’s window size for each downsampling stage. Default: (7, 7, 14, 7)

• mlp_ratio (float, optional) – Ratio of MLP dimension to embedding dimension in self-attention. Default: 4.0

• drop_rate (float, optional) – Dropout rate. Default: 0.0

• drop_path_rate (float, optional) – Stochastic depth rate. Default: 0.0

• use_checkpoint (bool, optional) – Whether to use activation checkpointing to trade compute for memory. Default: False

• mbconv_expand_ratio (float, optional) – Expansion ratio used in MBConv block. Default: 4.0

• local_conv_size (int, optional) – Kernel size of convolution used in TinyViTBlock Default: 3

• activation (type[Module], optional) – activation function. Default: nn.GELU

• mobile_sam (bool, optional) – Whether to use modifications for MobileSAM. Default: False

forward(x)

Classify images if mobile_sam=False, produce feature maps if mobile_sam=True.

Parameters:

x (Tensor) – images (B, C, H, W) with \(H = W =\) img_size. The strided convolutions before the first attention stage round up, so \(H\) and \(W\) from img_size - 7 to img_size give the same token grid and run too.

Raises:

ValueError – if the input gives a token grid other than that of an img_size image.

Return type:

Tensor

static from_config(variant, pretrained=False, **kwargs)

Create a TinyViT model from pre-defined variants.

Returns the model in evaluation mode when pretrained loads the complete weights, classification head included, and in training mode otherwise (see Pretrained models).

Parameters:

• variant (str) – TinyViT variant. Possible values: '5m', '11m', '21m'.

• pretrained (bool | str, optional) – whether to use pre-trained weights. Possible values: False, True, 'in22k', 'in1k'. For TinyViT-21M (variant='21m'), 'in1k_384', 'in1k_512' are also available; named, they build at the img_size you pass (224 by default), so pass img_size=384 or 512 to run them at their resolution. Default: False

• **kwargs (Any) – other keyword arguments that will be passed to TinyViT. They override the variant’s embed_dims, depths, num_heads, window_sizes and drop_path_rate.

Return type:

TinyViT

Note

pretrained=True loads the ImageNet-1k checkpoint ('in1k'); for TinyViT-21M it loads 'in1k_384' when img_size is at least 384 and 'in1k_512' when it is at least 512. For feature extraction or fine-tuning, the ImageNet-22k checkpoint ('in22k') is preferred.

img_size does not change the window sizes: they are the 224-px [7, 7, 14, 7] at every img_size, except that the TinyViT-21M 'in1k_384' and 'in1k_512' checkpoints get the windows they were trained with, [12, 12, 24, 12] and [16, 16, 32, 16], and the stochastic depth rate they were fine-tuned with, drop_path_rate=0.1 instead of TinyViT-21M’s 0.2. A window_sizes or drop_path_rate argument overrides them. The checkpoint’s attention biases are interpolated (bicubic) only when the window sizes differ from the checkpoint’s. img_size must be a multiple of 32 (of 16 with mobile_sam=True), and a checkpoint’s accuracy is not preserved away from the resolution it was trained at: on the panda image of the TinyViT example, TinyViT-5M’s probability of the right class drops from 0.37 at 224 px to 0.02 at img_size=256.

With pretrained and a num_classes other than the checkpoint’s, the classification head is reset to zeros, and the model, a pretrained backbone to fine-tune, stays in training mode. The 'in22k' checkpoints have a 21841-class head: pass num_classes=21841 to get the complete model in evaluation mode. With the default num_classes=1000 the head is reset and the model stays in training mode, so call .eval() before extracting features.

Image Stitching

class kornia.contrib.ImageStitcher(matcher, estimator='ransac', blending_method='naive')

Stitch images with overlapping fields of view.

Convention:

• The images are RGB, (B, 3, H, W), all of one height, ordered left to right. matcher is a LoFTR or a LocalFeatureMatcher, which gets grayscale copies; any other matcher raises NotImplementedError.

• The panorama is in the first image’s pixel frame and has its height: content that maps left of its first column, or above or below its rows, is lost. Each next image is warped onto the panorama so far with a homography from their matches; every sample needs 4 matches, or a RuntimeError names the sample.

• "naive" blending keeps the warped image wherever it covers a pixel completely, so in an overlap the right image replaces the left one; a pixel it covers only in part keeps the panorama’s value.

• The result drops the empty trailing columns, those after the last column that any sample covers; an uncovered column before it, the first one included, is kept. A batch shares one crop, so it is cropped to its widest sample and a narrower sample keeps zero columns on its right.

• estimator="ransac" uses RANSAC("homography") (inl_th=2.0); "vanilla" fits every match with iterated DLT.

Parameters:

• matcher (Module) – image feature matching module, LoFTR or LocalFeatureMatcher.

• estimator (str, optional) – method to compute homography, either “vanilla” or “ransac”. “ransac” is slower with a better accuracy. Default: "ransac"

• blending_method (str, optional) – method to blend two images together. Only “naive” is currently supported. Default: "naive"

import torch
import kornia.feature as KF
from kornia.contrib import ImageStitcher

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
stitcher = ImageStitcher(KF.LoFTR(pretrained="outdoor"), estimator="ransac").to(device)
# img_left and img_right: (1, 3, H, W) RGB images of one height on `device`
with torch.inference_mode():
    panorama = stitcher(img_left, img_right)

See Image Stitching for a complete example.

Lambda

class kornia.contrib.Lambda(func)

Applies user-defined lambda as a transform.

Convention:

forward() calls func(img, *args, **kwargs) and returns its result unchecked; it need not be a tensor.

Parameters:

func (Callable[..., Any]) – Callable function.

Returns:

The output of the user-defined lambda.

Example

>>> import kornia
>>> x = torch.rand(1, 3, 5, 5)
>>> f = Lambda(lambda x: kornia.color.rgb_to_grayscale(x))
>>> f(x).shape
torch.Size([1, 1, 5, 5])

Distance Transform

kornia.contrib.distance_transform(image, kernel_size=3, h=0.35)

Approximates the Euclidean distance transform of images/volumes using cascaded convolution operations.

Try in browser

It uses the method described in [PDP20]. The transformation is applied independently across the channel dimension.

Convention:

• The distance is to the nearest positive pixel, where SciPy’s distance_transform_edt and OpenCV’s distanceTransform measure it to the nearest zero: distance_transform(x) approximates scipy.ndimage.distance_transform_edt(x <= 0), which is distance_transform_edt(x == 0) for a mask of 0 and 1.

• image is a binary mask whose sources are the pixels equal to 1, read after clamping it to [0, 1]: a value above 1 is a full source, as 1 is, and a negative value is background, as 0 is. Values above 1 or below 0 are clamped, so their gradient is 0, while values in [0, 1], the bounds included, pass the gradient through. A value \(0 < v < 1\) is a weaker source: it adds \(-h \log v\) to a distance within kernel_size // 2 of it whose window holds no other source, less where the window holds one, as the soft minimum sums them before its logarithm, and does not change the distances further away.

• A mask without any positive pixel gives 0 everywhere, as a full mask does, which keeps the output and its gradient finite.

• Known defect: the result is not Euclidean, and there is no general bound on its error (#5956). From an isolated source it is exact only within kernel_size // 2; beyond, it follows the number of cascade steps, so a diagonal comes out short by up to a fraction \(1 - 1/\sqrt{2}\) (29 %) of its length in 2-D and \(1 - 1/\sqrt{3}\) (42 %) in 3-D, at any kernel_size. Separately, the soft minimum pulls a value down wherever several foreground or already reached pixels share its window (replicate padding counts a pixel on the image border twice), the more the larger h: at the default h by 0.167 px beside a straight edge, 0.243 px midway between two sources 2 px apart and 0.579 px in a one-pixel hole, and in a hole towards 0 as h nears its upper bound.

Parameters:

• image (Tensor) – Image or volume with shape (B,C,H,W) or (B,C,D,H,W).

• kernel_size (int, optional) – size of the convolution kernel. Must be an odd number of at least 3. Default: 3

• h (float, optional) – value that influences the approximation of the min function. The cascade’s kernel taps are exp(-dist/h). Accepted: h large enough that the corner tap (distance sqrt(ndim) * (kernel_size // 2)) stays at or above the smallest normal number of the working precision, and small enough that the off-centre taps sum to less than 1 by a margin for the working precision’s rounding (2-D: below 0.5624 at kernel_size 3, 0.5626 in float64, about 0.50 at larger kernels; 3-D: below 0.4070 at 3, 0.4071 in float64, about 0.38 at larger kernels). Any other h raises: a smaller one would give wrong distances, a larger one would give 0 at background pixels whose window is full of reached pixels. float16 and bfloat16 inputs run the cascade in float32 and cast the result back, which widens the safe h range for those dtypes to match float32. Default: 0.35

Return type:

Tensor

Returns:

tensor with the shape and dtype of the input.

Example

>>> # 2D example:
>>> tensor = torch.zeros(1, 1, 5, 5)
>>> tensor[:,:, 1, 2] = 1
>>> dt = distance_transform(tensor)
>>> # 3D example:
>>> volume = torch.zeros(1, 1, 5, 5, 5)
>>> volume[:, :, 2, 2, 2] = 1
>>> dt = distance_transform(volume)

class kornia.contrib.DistanceTransform(kernel_size=3, h=0.35)

Module that approximates the Euclidean distance transform of images/volumes using convolutions.

Convention:

See the Convention block of distance_transform().

Parameters:

• kernel_size (int, optional) – size of the convolution kernel. Default: 3

• h (float, optional) – value that influences the approximation of the min function. See distance_transform() for the safe range of h relative to kernel_size and working precision. Default: 0.35

Plasma Fractal Noise

kornia.contrib.diamond_square(output_size, roughness=0.5, random_scale=1.0, random_fn=torch.rand, normalize_range=None, device=None, dtype=None)

Generate Plasma Fractal Images using a variant of the diamond square algorithm.

It ports tormentor’s variant of the diamond-square algorithm.

Convention:

• The map is grown on a \(2^n + 1\) grid and cropped to output_size at its top-left corner, so it can be a non-contiguous view. The seed is a \(3 \times K\) grid, transposed to \(K \times 3\) for a map taller than wide, whose cell corners, the even nodes of its first and last rows (columns when transposed), are random_fn draws, and whose other nodes are the means of their neighbours. Each level sets a new node to \((1 - s_k)\,\overline{n} + s_k\,u\), the mean \(\overline{n}\) of its neighbours mixed with a random_fn draw \(u\) by the weight \(s_k = \mathrm{random\_scale} \cdot \mathrm{roughness}^k\) at level \(k \geq 1\). random_scale and roughness act only through \(s_k\), a mixing weight, not an amplitude. A border node of the square step averages its three neighbours; the diamond step’s nodes are all interior, so its reflect padding never reaches them.

• random_fn is called as random_fn(shape, device=device, dtype=dtype) and is the only source of randomness, so functools.partial(torch.rand, generator=g) leaves the global generator alone. With torch.rand the map stays in [0, 1] without normalization while random_scale and roughness are in [0, 1], and it is flat in mean (0.5 at every pixel) for every random_scale, though its variance varies across the map.

• roughness and random_scale are a float, a 0-d tensor or a (B,) tensor with one value per sample, and follow device and dtype; with dtype=None a tensor keeps its dtype and a float takes the default dtype. When random_fn returns the dtype it is asked for, the map has the dtype that dtype (the default dtype for None), roughness and random_scale promote to, at every map size. A random_fn that returns another dtype also promotes the levels’ arithmetic to its dtype, but not a map with a side of 3 px or less, which is the seed alone and stores the draws in the requested dtype.

• normalize_range=(min, max) rescales each \((b, c)\) plane to that range with normalize_min_max().

Parameters:

• output_size (Tuple[int, int, int, int]) – a tuple of integers with the BxCxHxW of the image to be generated.

• roughness (Union[float, Tensor], optional) – the factor that scales the mixing weight at each recursion step. Default: 0.5

• random_scale (Union[float, Tensor], optional) – scales the per-level mixing weight \(s_k = \mathrm{random\_scale} \cdot \mathrm{roughness}^k\), \(k \geq 1\). Default: 1.0

• random_fn (Callable[..., Tensor], optional) – the callable function to use to sample a random torch.Tensor. Default: torch.rand

• normalize_range (Optional[Tuple[float, float]], optional) – if given, the (min, max) range each plane is min-max normalized to. Default: None

• device (Optional[device], optional) – the torch device to place the output map. Default: None

• dtype (Optional[dtype], optional) – the torch dtype to place the output map. Default: None

Return type:

Tensor

Returns:

A torch.Tensor with shape (B,C,H,W) containing the fractal image.

KMeans

class kornia.contrib.KMeans(num_clusters, cluster_centers, tolerance=10e-4, max_iterations=0, seed=None)

Implements the kmeans clustering algorithm with euclidean distance as similarity measure.

Convention:

• fit() stops after the first update whose squared total centre shift, \((\sum_k \lVert c_k^{\text{new}} - c_k \rVert)^2\), is at most tolerance, or after max_iterations updates. max_iterations=0 sets no limit, so tolerance must not be None then.

• seed calls torch.manual_seed(seed), which reseeds the global generator. The initial centres, sampled from the data of each fit() when cluster_centers is None, and the points that empty clusters move to are drawn from it; there is no generator argument. A second fit() on the same data therefore starts from other centres unless the global generator is reseeded first.

• A point at the same distance from two centres goes to the one with the lower index. A cluster that gets no point moves to a random data point.

• Distances and means are computed in the promoted dtype of the data and the centres, at least float32 (int64 for integer data and centres). The centres keep the dtype of the data, so an integer mean is truncated toward zero.

Parameters:

• num_clusters (int) – number of clusters the data has to be assigned to, a positive integer. Required, also with cluster_centers; without them, fit() needs at least num_clusters samples.

• cluster_centers (Tensor | None) – the (num_clusters, D) starting centres, or None to sample them from the data in each fit().

• tolerance (float, optional) – the bound on the squared total centre shift that stops fit(), or None to stop only after max_iterations updates. Default: 10e-4

• max_iterations (int, optional) – the maximum number of updates; 0 sets no limit. Default: 0

• seed (int | None, optional) – if given, passed to torch.manual_seed when the object is created. Default: None

Example

>>> kmeans = kornia.contrib.KMeans(3, None, 10e-4, 100, 0)
>>> kmeans.fit(torch.rand((1000, 5)))
>>> predictions = kmeans.predict(torch.rand((10, 5)))

property cluster_assignments: Tensor

Return cluster labels assigned during the most recent fit call.

Returns:

A 1D tensor with shape (N,), where N is the number of samples given to fit(). Each value is the cluster index nearest to the corresponding sample in the final centers, matching predict() on the training data even when fitting stops early.

Raises:

TypeError – If fit has not been run yet.

property cluster_centers: Tensor

Return the current cluster centers.

Returns:

• C is the number of clusters.

• D is the feature dimension of each sample.

If fit() has already been called, this returns the learned final centers. Otherwise, it returns the initialization provided during construction.

Return type:

A tensor with shape (C, D)

Raises:

TypeError – If no initial centers were provided and fit has not been run.

fit(X)

Fit iterative KMeans clustering till a threshold for shift in cluster centers or a maximum no of iterations have reached.

Parameters:

X (Tensor) – 2D input torch.Tensor to be clustered

Return type:

None

predict(x)

Find the cluster center closest to each point in x.

Parameters:

x (Tensor) – 2D torch.Tensor

Return type:

Tensor

Returns:

1D torch.Tensor containing cluster id assigned to each data point in x

Super Resolution

class kornia.contrib.SuperResolution(model, pre_processor, post_processor, name=None)

SuperResolution is a module that wraps an super resolution model.

forward(images)

Forward pass of the super resolution model.

Parameters:

images (Union[Tensor, List[Tensor]]) – If list of RGB images. Each image is a torch.Tensor with shape (3, H, W). If torch.Tensor, a torch.Tensor with shape (B, 3, H, W).

Return type:

Union[Tensor, List[Tensor]]

Returns:

output torch.Tensor.

static from_config(config)

Build a super-resolution model from a configuration object.

Dispatches to SmallSRBuilder for the lightweight variant and to RRDBNetBuilder for the RRDB variants.

Parameters:

config (SuperResolutionConfig) – Super-resolution configuration. See SuperResolutionConfig.

Return type:

SuperResolution

Returns:

A SuperResolution wrapper ready for inference.

Raises:

ValueError – If config.model_name is not a supported variant.

save(images, edge_maps=None, directory=None, output_type='torch')

Save the super resolution results.

Parameters:

• images (Union[Tensor, List[Tensor]]) – input torch.Tensor.

• edge_maps (Union[Tensor, List[Tensor], None], optional) – super-resolved images; computed with forward() when None. Default: None

• directory (Optional[str], optional) – where to save outputs. Default: None

• output_type (str, optional) – “torch” or “pil”, passed to visualize(); the files hold the same pixels either way. Default: "torch"

Return type:

None

to_onnx(onnx_name=None, include_pre_and_post_processor=True, save=True, additional_metadata=None, **kwargs)

Export the current super resolution model to an ONNX model file.

Parameters:

• onnx_name (Optional[str], optional) – The name of the output ONNX file. If not provided, "kornia_<name>.onnx", with the model’s name, for example "kornia_small_sr.onnx". Default: None

• include_pre_and_post_processor (bool, optional) – Whether to include the pre-processor and post-processor in the exported model. Default: True

• save (bool, optional) – If True, write the model to onnx_name; it is returned either way. Default: True

• additional_metadata (Optional[List[Tuple[str, str]]], optional) – Additional metadata to add to the ONNX model. Default: None

• kwargs (Any) – Additional keyword arguments of torch.onnx.export().

Return type:

ModelProto

Returns:

The exported ONNX model.

Raises:

ValueError – With the pre- and post-processor, if the output is a list (or tuple) with one image per input image and the exported batch is dynamic, from a dynamic_axes or dynamic_shapes argument or, with the TorchScript exporter, because the list comes only for more than one image.

Note

With the pre- and post-processor, an output that is a list with one image per input image, or a tuple as long as the batch at two batch sizes, is exported for a batch of one image: the exporters turn such a sequence into one graph output per traced image, so a dynamic batch would return the traced batch’s images at any batch size. The list comes from a post-processor that returns one, from the network itself, or from a pre-processor that returns a list of images, which forward() runs through the network one at a time. Known tensor-output networks (the builders’ RRDB and SmallSR networks, nn.Identity and nn.Conv2d) with nn.Identity or ResizePreProcessor before them and nn.Identity or OutputRangePostProcessor after them keep a dynamic batch without a probe, provided the wrapper and each component retain their known forward implementations and no forward hook applies to them, a global one included. Other combinations are probed on one image at the export size in eval mode, without gradients and outside the export. A list result requires a fixed batch; a tensor result is checked again on the actual tracing batch before accepting a dynamic graph. A dynamic graph with per-image outputs raises a ValueError after the export, before anything is saved.

visualize(images, edge_maps=None, output_type='torch')

Draw the super resolution results.

Parameters:

• images (Union[Tensor, List[Tensor]]) – input torch.Tensor.

• edge_maps (Union[Tensor, List[Tensor], None], optional) – super-resolved images; computed with forward() when None. Default: None

• output_type (str, optional) – type of the output. Default: "torch"

Return type:

Union[Tensor, List[Tensor], List[Image]]

Returns:

output torch.Tensor.

class kornia.contrib.SuperResolutionConfig(model_name='small_sr', pretrained=True, upscale_factor=3, image_size=None)

Configuration to construct a SuperResolution model.

The fields mirror the arguments accepted by RRDBNetBuilder and SmallSRBuilder, so a config can drive either builder.

Parameters:

• model_name (str, optional) – Name of the model variant to build. RRDB variants are "RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B" and "RealESRGAN_x2plus"; the lightweight variant is "small_sr". Default: 'small_sr'

• pretrained (bool, optional) – If True, load pretrained weights for the selected variant. Default: True

• upscale_factor (int, optional) – Integer scale used by the lightweight variant. Ignored by the RRDB variants, whose scale is fixed by model_name. The lightweight variant resizes every input to 224 x 224 first, discarding its aspect ratio, so its output is 224 * upscale_factor square (672 x 672 for 3, from a 100 x 50 input as from a 480 x 640 one). Its pretrained weights exist only for 3: another factor with pretrained=True raises a size-mismatch RuntimeError. Default: 3

• image_size (Optional[int], optional) – Optional fixed input size for the lightweight variant. It declares the size fed to the wrapper; it does not change the pre-processor, which always resizes to 224, so the output side stays 224 * upscale_factor. Default: None

class kornia.contrib.SmallSRBuilder

Provide a static factory to build lightweight Super-Resolution models.

static build(model_name='small_sr', pretrained=True, upscale_factor=3, image_size=None)

Build a lightweight super-resolution (SR) model wrapper.

SR stands for super-resolution, which means generating a higher resolution output from a lower resolution input.

Parameters:

• model_name (str, optional) – Name of the model variant to build. Currently only "small_sr" is supported. Default: "small_sr"

• pretrained (bool, optional) – If True, load pretrained weights for the selected model. Default: True

• upscale_factor (int, optional) – Integer scale used by SmallSRNetWrapper. The pre-processor first resizes every input to 224 x 224, discarding its aspect ratio, so the output is 224 * upscale_factor square: with 3, a 100 x 50 and a 480 x 640 input both give 672 x 672. Pretrained weights exist only for 3; another factor with pretrained=True raises a size-mismatch RuntimeError. Default: 3

• image_size (Optional[int], optional) – Optional fixed input size. It declares the size fed to the wrapper and does not change the pre-processor, which always resizes to 224; the output side is therefore 224 * upscale_factor. Default: None

Return type:

SuperResolution

Returns:

A SuperResolution instance configured with resizing pre-processing and identity post-processing.

Raises:

ValueError – If model_name is not supported.

class kornia.contrib.RRDBNetBuilder

Provide a static factory to build Residual-in-Residual Dense Block (RRDB) networks.

static build(model_name='RealESRNet_x4plus', pretrained=True)

Build a preconfigured RRDB (Residual-in-Residual Dense Block) model.

The returned object is a SuperResolution wrapper configured for inference. Internally, this factory selects architecture hyperparameters based on model_name, and optionally loads pretrained weights.

Parameters:

• model_name (str, optional) – Name of the RRDB variant to construct. Supported options are: "RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B", and "RealESRGAN_x2plus". Default: "RealESRNet_x4plus"

• pretrained (bool, optional) – If True, download and load pretrained weights for the selected model variant. Default: True

Return type:

SuperResolution

Returns:

A SuperResolution wrapper ready for inference.

Raises:

ValueError – If model_name is not one of the supported variants.