kornia.models#

Builders for Kornia’s ready-to-use models (object detection, edge detection, super-resolution, semantic segmentation and Kimi-VL), each returning a regular nn.Module. The object detection, edge detection and super-resolution builders load pretrained weights by default and KimiVLBuilder.from_pretrained_hf loads them always, downloading them on first use. SegmentationModelsBuilder downloads nothing: it wraps a network you built yourself. The page also covers the vision tower of Qwen2.5-VL, whose weights kornia does not host: Qwen25VisionTransformer.from_pretrained reads a checkpoint you downloaded yourself (see Qwen25VisionTransformer). For the papers behind the models, see the Models section. Every model here accepts a batched (B, 3, H, W) float image in [0, 1], except Kimi-VL, Qwen2.5-VL and a SegmentationModelsBuilder model built without preproc_params, which passes its input to the network unchanged. Kimi-VL takes that image normalized with mean = std = 0.5 per channel (2 * x - 1), as its checkpoint’s preprocessor_config.json specifies, with H and W multiples of 28 (see kornia.models.kimi_vl.KimiVLModel.forward). The Qwen2.5-VL vision tower takes it normalized with the CLIP mean and std. kornia.io.get_sample_images() provides a couple of sample images for quick experiments.

RTDETRDetectorBuilder#

The RTDETRDetectorBuilder class is a builder for constructing a detection model based on the RT-DETR architecture, which is designed for real-time object detection. It is capable of detecting multiple objects within an image and provides efficient inference suitable for real-world applications.

Key Methods:

build: Constructs and returns an instance of the RTDETR detection model.
visualize: Draws the detected boxes on the input images.

class kornia.contrib.object_detection.RTDETRDetectorBuilder[source]#

Bases: object

A builder class for constructing RT-DETR object detection models.

Try in browser

This class provides static methods to:

Build an object detection model from a model name or configuration.
Export the model to ONNX format for inference.

Note

To use this model, load image tensors and call model.save(images).

Example

The following code demonstrates how to use RTDETRDetectorBuilder to detect objects in an image:

import kornia
from kornia.contrib.object_detection import RTDETRDetectorBuilder

image = kornia.io.get_sample_images()[0][None]
model = RTDETRDetectorBuilder.build()
detections = model(image)  # list of (D, 6) tensors: class id, score, x, y, w, h
drawn = model.visualize(image, detections)  # the boxes drawn on the image

static build(model_name=None, config=None, pretrained=True, image_size=None, confidence_threshold=None, confidence_filtering=None)[source]#

Build and returns an RT-DETR object detector model.

Either model_name or config must be provided. If neither is provided, a default pretrained model (rtdetr_r18vd) will be built.

Parameters:

model_name (Optional[str], optional) – Name of the RT-DETR model to load. Can be one of the available pretrained models. Including ‘rtdetr_r18vd’, ‘rtdetr_r34vd’, ‘rtdetr_r50vd_m’, ‘rtdetr_r50vd’, ‘rtdetr_r101vd’. Default: None
config (Optional[Any], optional) – A custom configuration object for building the RT-DETR model. Default: None
pretrained (bool, optional) – Whether to load a pretrained version of the model (applies when model_name is provided). Default: True
image_size (Optional[int], optional) – The size to which input images will be resized during preprocessing, used only with config. If None, config.input_size is used. With model_name, or without either, it is ignored: the input is resized to the 640 px the published weights were trained for. Recommended scales include [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]. Default: None
confidence_threshold (Optional[float], optional) – Threshold to filter results based on confidence scores. Default: None
confidence_filtering (Optional[bool], optional) – Whether to filter results based on confidence scores. None filters. Default: None

Return type:

ObjectDetector

Returns:

ObjectDetector: An object detector instance initialized with the specified model, preprocessor, and post-processor.

EdgeDetectorBuilder#

The EdgeDetectorBuilder class implements a state-of-the-art edge detection model based on DexiNed, which excels at detecting fine-grained edges in images. This model is well-suited for tasks like medical imaging, object contour detection, and more.

Key Methods:

build: Builds and returns an instance of the DexiNed edge detection model.
visualize: Returns the edge maps as images for further processing or display.

class kornia.contrib.edge_detection.EdgeDetectorBuilder[source]#

Bases: object

EdgeDetectorBuilder is a class that builds an edge detection model.

This is a high-level API that builds edge detection models like kornia.models.DexiNed and wraps them with EdgeDetector.

Note

To use this model, load image tensors and call model.save(images).

Example

The following code shows how to use the EdgeDetectorBuilder to detect edges in an image:

import kornia
from kornia.contrib.edge_detection import EdgeDetectorBuilder

image = kornia.io.get_sample_images()[0][None]
model = EdgeDetectorBuilder.build()
edges = model(image)  # list with one (1, 1, H, W) edge map per image

static build(model_name='dexined', pretrained=True, image_size=352)[source]#

Build an edge detection model.

The detector takes RGB images in \([0, 1]\). For "dexined" it flips them to BGR, scales them to \([0, 255]\) and subtracts the BGR mean [103.939, 116.779, 123.68], the input the pretrained DexiNed was trained on (xavysp/DexiNed@d944b70), then applies a sigmoid to the network’s logits.

Parameters:

model_name (str, optional) – Name of the model to build. Currently only “dexined” is supported. Default: "dexined"
pretrained (bool, optional) – If True, loads pretrained weights. Default: True
image_size (int, optional) – Size to which input images will be resized during preprocessing. Default: 352

Return type:

EdgeDetector

Returns:

EdgeDetector instance configured with the specified model.

Example

>>> detector = EdgeDetectorBuilder.build(pretrained=True, image_size=352)
>>> img = torch.rand(1, 3, 320, 320)
>>> out = detector(img)

RRDBNet#

The RRDBNet class is the Residual-in-Residual Dense Block generator behind ESRGAN and Real-ESRGAN. It is a plain nn.Module that upsamples a batched (B, 3, H, W) image by a factor of 1, 2 or 4, and its module and parameter names match the reference implementation, so the published Real-ESRGAN checkpoints load with strict=True. kornia.contrib.super_resolution.RRDBNetBuilder configures it for the released Real-ESRGAN variants and downloads their weights, and returns a SuperResolution wrapper ready for inference:

from kornia.contrib.super_resolution import RRDBNetBuilder

model = RRDBNetBuilder.build("RealESRNet_x4plus")
upscaled = model(images)  # (B, 3, H, W) -> (B, 3, 4H, 4W)

class kornia.models.RRDBNet(num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32)[source]#

Bases: Module

Network consisting of Residual in Residual Dense Blocks, as used in ESRGAN and Real-ESRGAN.

ESRGAN is extended here for scale x2 and scale x1. For those scales the input is first pixel-unshuffled – the inverse of a pixel shuffle – to reduce the spatial size and enlarge the channel size before it is fed into the main ESRGAN architecture, so the network always upsamples by a factor of 4 internally.

Parameters:

num_in_ch (int) – Channel number of inputs.
num_out_ch (int) – Channel number of outputs.
scale (int, optional) – Upsampling factor. One of 1, 2 or 4. Default: 4
num_feat (int, optional) – Channel number of intermediate features. Default: 64
num_block (int, optional) – Block number in the trunk network. Default: 23
num_grow_ch (int, optional) – Channels for each growth. Default: 32

Shape:

Input: \((B, C_{in}, H, W)\). For scale=2 both spatial sizes must be divisible by 2, and for scale=1 by 4.
Output: \((B, C_{out}, H \cdot scale, W \cdot scale)\).

Raises:: ValueError – If scale is not 1, 2 or 4.

Example

>>> import torch
>>> model = RRDBNet(num_in_ch=3, num_out_ch=3, scale=4, num_feat=8, num_block=1, num_grow_ch=4)
>>> model(torch.rand(1, 3, 8, 8)).shape
torch.Size([1, 3, 32, 32])

Example

The following code upsamples an image by a factor of 4 with a randomly initialized generator:

import torch
from kornia.models import RRDBNet

model = RRDBNet(num_in_ch=3, num_out_ch=3, scale=4, num_feat=64, num_block=23).eval()
upsampled = model(torch.rand(1, 3, 32, 32))  # (1, 3, 128, 128)

forward(x)[source]#

Super-resolve x of shape \((B, C_{in}, H, W)\) by self.scale.

Return type:: Tensor

SegmentationModelsBuilder#

The SegmentationModelsBuilder class wraps a segmentation network you have already built – typically one from segmentation_models_pytorch (smp), but any nn.Module mapping (B, 3, H, W) to (B, C, H, W) works – in a SemanticSegmentation container, prepending the ONNX-friendly preprocessing (BGR-to-RGB, range rescaling, mean/std normalization) that the encoder’s pretrained weights expect. Kornia does not import smp; you build the network and fetch its preprocessing parameters yourself.

Key Methods:

build: Wraps a constructed segmentation network and its encoder’s preprocessing parameters.
get_preprocessing_pipeline: Turns a preprocessing-parameter dictionary into an ImageSequential.

Main parameters of build:

model: (nn.Module) The segmentation network.
preproc_params: (dict | None) The encoder’s preprocessing parameters, in the shape returned by smp.encoders.get_preprocessing_params(encoder_name): input_space, input_range, mean and std. None means the input is fed to the network unchanged.
name: (str) The name of the wrapped model, used by save.

class kornia.models.segmentation.segmentation_models.SegmentationModelsBuilder[source]#

Bases: object

Wrap a segmentation network and its encoder’s preprocessing in a SemanticSegmentation.

The builder is written for networks from segmentation_models_pytorch (smp), whose encoders ship the preprocessing parameters their pretrained weights expect, but any nn.Module mapping a (B, 3, H, W) image batch to a (B, C, H, W) prediction works. Kornia does not import smp: you build the network and fetch its preprocessing parameters, and the builder supplies the ONNX-friendly preprocessing pipeline and the container. Give the network a softmax head (activation="softmax2d" in smp) if you want SemanticSegmentation.visualize(): it expects per-pixel class probabilities and raises on raw logits.

Example

>>> import segmentation_models_pytorch as smp
>>> from kornia.models.segmentation import SegmentationModelsBuilder
>>> net = smp.Unet(
...     encoder_name="resnet34", encoder_weights="imagenet", classes=2, activation="softmax2d"
... )
>>> params = smp.encoders.get_preprocessing_params("resnet34")
>>> model = SegmentationModelsBuilder.build(net, params, name="Unet_resnet34")
>>> model(torch.rand(1, 3, 64, 64)).shape
torch.Size([1, 2, 64, 64])

Example

Here’s an example of how to use SegmentationModelsBuilder with an smp UNet for two-class segmentation:

import kornia
import segmentation_models_pytorch as smp
from kornia.models.segmentation import SegmentationModelsBuilder

net = smp.Unet(encoder_name="resnet34", encoder_weights="imagenet", classes=2, activation="softmax2d")
params = smp.encoders.get_preprocessing_params("resnet34")
model = SegmentationModelsBuilder.build(net, params, name="Unet_resnet34")

input_tensor = kornia.io.get_sample_images()[0][None]
segmented_output = model(input_tensor)
print(segmented_output.shape)  # (1, 2, H, W)

The softmax head is what visualize() needs: it colours each pixel by its most probable class and raises on raw logits.

static build(model, preproc_params=None, name='segmentation_model')[source]#

Wrap a constructed segmentation network in a SemanticSegmentation.

Parameters:

model (Module) – The segmentation network, e.g. smp.Unet(...). It is put in eval mode.
preproc_params (Optional[dict[str, Any]], optional) – The preprocessing parameters of the network’s encoder, in the shape returned by smp.encoders.get_preprocessing_params(encoder_name): the keys input_space ("RGB" or "BGR"), input_range ([0, 1] or [0, 255]), mean and std (per-channel lists, or None for no normalization). See get_preprocessing_pipeline(). None feeds the input to the network unchanged. Default: None
name (str, optional) – Name of the wrapped model; SemanticSegmentation.save() uses it for file names. Default: "segmentation_model"

Return type:

SemanticSegmentation

Returns:

The container running preprocessing, the network and an identity post-processor.

static get_preprocessing_pipeline(preproc_params)[source]#

Build the preprocessing pipeline expected by a segmentation model.

Parameters:: preproc_params (dict[str, Any]) – Dictionary from the segmentation-model metadata, e.g. smp.encoders.get_preprocessing_params(encoder_name). It must carry the keys input_space ("RGB" or "BGR": the color order the network was trained on, so a "BGR" network gets its RGB input flipped), input_range ([0, 1] or [0, 255]: the range the mean/std are expressed in, so [0, 255] multiplies the [0, 1] input by 255 first), and mean and std (per-channel lists, or None for no normalization).
Return type:: ImageSequential
Returns:: ImageSequential containing ONNX-friendly color conversion, rescaling, and normalization steps.

Note

Set pipeline.disable_features = True before exporting the returned pipeline to ONNX. This disables convenience input/output conversion and output caching, whose tensor attribute mutation is rejected by some versions of torch.export.

Raises:

BaseError – If one of the four keys is missing (a KORNIA_CHECK()).
ValueError – If input_space or input_range is not one of the supported values.

class kornia.models.segmentation.SemanticSegmentation(model, pre_processor, post_processor, name=None)[source]#

Semantic Segmentation is a module that wraps a semantic segmentation model.

It runs pre_processor, model and post_processor in turn on a batch or a list of images. SegmentationModelsBuilder builds one around a segmentation_models_pytorch network and its encoder’s preprocessing.

Parameters:

model (Module) – The segmentation network, mapping a (B, 3, H, W) batch to (B, C, H, W) predictions. visualize() expects per-pixel class probabilities (a softmax head); raw logits raise.
pre_processor (Module) – Pre-processing module applied to the input images.
post_processor (Module) – Post-processing module applied to the network output.
name (Optional[str], optional) – Optional name, used by save() for file names. Default: None

forward(images)[source]#

Forward pass of the semantic segmentation 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 tensor.

visualize(images, semantic_masks=None, output_type='torch', colormap='random', manual_seed=2147)[source]#

Visualize the segmentation masks.

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)\).
semantic_masks (Union[Tensor, list[Tensor], None], optional) – If list of segmentation masks. Each mask is a torch.Tensor with shape \((C, H, W)\). If torch.Tensor, a torch.Tensor with shape \((B, C, H, W)\). Default: None
output_type (str, optional) – The type of output, can be “torch” or “pil”. Default: "torch"
colormap (str, optional) – The colormap to use, can be “random” or a custom color map. Default: "random"
manual_seed (int, optional) – The manual seed to use for the colormap. Default: 2147

Return type:

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

KimiVLBuilder#

The KimiVLBuilder class constructs Kimi-VL models from a configuration or downloads pretrained weights. Pretrained loading currently supports only the converted Kimi-VL-A3B-Instruct vision encoder and projector checkpoint.

Key Methods:

from_config: Constructs a randomly initialized Kimi-VL model from a KimiVLConfig.
from_pretrained_hf: Downloads and strictly loads the supported pretrained checkpoint.

class kornia.models.kimi_vl.KimiVLBuilder[source]#

Bases: object

Builder for Kimi-VL models.

Provides convenient methods to create Kimi-VL models from configs or load pretrained weights.

Example

The following code loads the supported pretrained Kimi-VL vision model:

from kornia.models.kimi_vl import KimiVLBuilder

model = KimiVLBuilder.from_pretrained_hf().eval()

static from_config(config)[source]#

Build model from configuration.

Parameters:: config (KimiVLConfig) – Model configuration.
Return type:: KimiVLModel
Returns:: KimiVLModel instance.

static from_pretrained_hf(cache_dir=None)[source]#

Load pretrained Kimi-VL-A3B-Instruct vision weights from Hugging Face Hub.

Downloads the vision encoder and projector weights of moonshotai/Kimi-VL-A3B-Instruct from the Kornia-owned safetensors checkpoint at https://huggingface.co/kornia/kimi-vl-a3b-instruct-vision. The checkpoint values are bitwise-identical to the original release (bf16), including the full 64x64 positional-embedding grid, which the model interpolates at runtime for other input resolutions.

Parameters:: cache_dir (Optional[str], optional) – Optional cache directory for downloaded files. Defaults to torch’s hub cache, which is where every other kornia checkpoint is cached. Default: None
Return type:: KimiVLModel
Returns:: KimiVLModel instance with pretrained weights, in evaluation mode, since it always loads weights (see Pretrained models).

Note

Only Kimi-VL-A3B-Instruct is currently supported.

Qwen25VisionTransformer#

The Qwen25VisionTransformer class is the vision tower of Qwen2.5-VL. It takes a batched (B, 3, H, W) image, with H and W multiples of 28, already normalized with the CLIP mean and std the checkpoint’s preprocessor_config.json names, and returns the merged vision tokens. kornia does not host the weights: from_pretrained reads the visual.* tensors of a Qwen2.5-VL checkpoint you downloaded yourself.

class kornia.models.qwen25.Qwen25VisionTransformer(config=None)[source]#

Vision tower of Qwen2.5-VL.

A ViT with a temporal patch embedding, a 2D rotary position embedding, RMS-normalized blocks with gated MLPs, attention restricted to windows except in config.fullatt_block_indexes, and a merger that turns every 2 x 2 square of patches into one output token. It follows Qwen2_5_VisionTransformerPretrainedModel in Hugging Face transformers v4.51.0, with two differences in interface and none in the result:

It takes images rather than flattened patches. Each image is repeated over temporal_patch_size frames, which is what the reference image processor does for a still image.
Windowed attention uses a mask over tokens kept in place, instead of reordering the tokens into windows and back. Attention with a rotary embedding is unchanged by permuting the tokens, so the output matches.

The input must already be resized and normalized as the checkpoint’s preprocessor_config.json describes: for the released checkpoints, values in \([0, 1]\) normalized with mean (0.48145466, 0.4578275, 0.40821073) and std (0.26862954, 0.26130258, 0.27577711).

Parameters:: config (Optional[Qwen25VisionConfig], optional) – Vision tower configuration. Defaults to the Qwen2.5-VL-3B-Instruct values. Default: None

Example

>>> config = Qwen25VisionConfig(depth=2, hidden_size=32, intermediate_size=48, num_heads=2,
...                             out_hidden_size=16, fullatt_block_indexes=(1,))
>>> model = Qwen25VisionTransformer(config)
>>> model(torch.rand(1, 3, 56, 84)).shape
torch.Size([1, 6, 16])

Example

from kornia.models.qwen25 import Qwen25VisionTransformer

model = Qwen25VisionTransformer.from_pretrained("path/to/Qwen2.5-VL-3B-Instruct")

forward(images)[source]#

Encode a batch of images into vision tokens.

Parameters:: images (Tensor) – Normalized images with shape \((B, C, H, W)\). H and W must be multiples of patch_size * spatial_merge_size (28 for the released checkpoints).
Return type:: Tensor
Returns:: Tokens with shape \((B, H W / (p s)^2, D_{out})\) in raster order over the merged grid, where \(p\) is the patch size and \(s\) the spatial merge size.

classmethod from_pretrained(path, config=None)[source]#

Build the tower and load the visual.* weights of a Qwen2.5-VL checkpoint the user downloaded.

kornia does not host these weights. Only the vision tensors are read, so the language model in the same shards is never copied into memory. The keys are mapped strictly: a missing or unexpected key raises.

Parameters:

path (str | PathLike[str]) – A .safetensors file, or a checkpoint directory containing model.safetensors.index.json or model.safetensors.
config (Optional[Qwen25VisionConfig], optional) – Configuration matching the checkpoint. Defaults to Qwen2.5-VL-3B-Instruct. The 7B checkpoint needs out_hidden_size=3584; the 32B and 72B checkpoints need intermediate_size=3456 and out_hidden_size 5120 and 8192. Default: None

Return type:

Qwen25VisionTransformer

Returns:

The tower with the checkpoint’s weights, in float32 and in evaluation mode.

class kornia.models.qwen25.Qwen25VisionConfig(depth=32, hidden_size=1280, intermediate_size=3420, num_heads=16, in_channels=3, patch_size=14, temporal_patch_size=2, spatial_merge_size=2, window_size=112, fullatt_block_indexes=(7, 15, 23, 31), out_hidden_size=2048, rms_norm_eps=1e-06, rope_theta=10000.0)[source]#

Configuration for the Qwen2.5-VL vision tower.

The defaults are the vision_config of the released Qwen/Qwen2.5-VL-3B-Instruct checkpoint. The 7B release differs only in out_hidden_size (3584); the 32B and 72B releases also use intermediate_size 3456, with out_hidden_size 5120 and 8192.

Parameters:

depth (int, optional) – Number of transformer blocks. Default: 32
hidden_size (int, optional) – Token dimension inside the blocks. Default: 1280
intermediate_size (int, optional) – Hidden dimension of the gated MLP in each block. Default: 3420
num_heads (int, optional) – Number of attention heads. Default: 16
in_channels (int, optional) – Number of input image channels. Default: 3
patch_size (int, optional) – Spatial size of a patch, in pixels. Default: 14
temporal_patch_size (int, optional) – Number of frames one patch spans. A still image is repeated this many times, as the reference image processor does. Default: 2
spatial_merge_size (int, optional) – Side of the square of neighbouring patches the merger concatenates into one output token. Default: 2
window_size (int, optional) – Side of the attention window, in pixels. Default: 112
fullatt_block_indexes (Tuple[int, ...], optional) – Blocks that attend over the whole image rather than within a window. Default: (7, 15, 23, 31)
out_hidden_size (int, optional) – Dimension of the output tokens. Default: 2048
rms_norm_eps (float, optional) – Epsilon of the RMS normalizations. Default: 1e-06
rope_theta (float, optional) – Base of the rotary position embedding frequencies. Default: 10000.0

Note

This documentation provides detailed information about each model class, its methods, and usage examples. For further details on individual methods and arguments, refer to the respective code documentation.