kornia.models#

Builders for Kornia’s ready-to-use models (object detection, edge detection, semantic segmentation and Kimi-VL). Each builder returns a configured model with pretrained weights. For the papers behind the models, see the Models section. Pretrained weights are downloaded on first use, and the model builders return a regular nn.Module that accepts a batched (B, 3, H, W) float image in [0, 1]. 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. If None, no resizing will be inferred from config file. 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. 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.

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)

The architecture is vendored from BasicSR (Apache-2.0, Copyright 2018-2022 BasicSR Authors); no extra package is required to use it.

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.

Note

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

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.