Overview
Provides a super-resolution pipeline wrapping RRDB-based (RealESRGAN) and lightweight (SmallSR) models with pre-processing, inference, post-processing, ONNX export, and visualization.
Description
The super_resolution module in the Kornia contrib package implements a complete image super-resolution pipeline. The SuperResolution class extends ModelBase and ONNXExportMixin to wrap a super-resolution model with pre-processing and post-processing stages. Two builder classes are provided: RRDBNetBuilder constructs RealESRGAN/RealESRNet models (RRDB architecture) with support for 2x and 4x upscaling using pretrained weights, while SmallSRBuilder constructs a lightweight super-resolution model. The module supports ONNX export, visualization (saving both source and super-resolved images), and configurable input/output image sizes.
Usage
Import this module when you need to upscale images using deep learning-based super-resolution, either with full-quality RRDB models (RealESRGAN) or lightweight models (SmallSR). Also use it when you need to export a super-resolution model to ONNX format.
Code Reference
Source Location
Signature
class SuperResolution(ModelBase, ONNXExportMixin):
name: str = "super_resolution"
input_image_size: Optional[int]
output_image_size: Optional[int]
pseudo_image_size: Optional[int]
def forward(self, images: Union[torch.Tensor, List[torch.Tensor]]) -> Union[torch.Tensor, List[torch.Tensor]]: ...
def visualize(self, images, edge_maps=None, output_type="torch"): ...
def save(self, images, edge_maps=None, directory=None, output_type="torch") -> None: ...
def to_onnx(self, onnx_name=None, include_pre_and_post_processor=True,
save=True, additional_metadata=None, **kwargs): ...
class RRDBNetBuilder:
@staticmethod
def build(model_name: str = "RealESRNet_x4plus", pretrained: bool = True) -> SuperResolution: ...
class SmallSRBuilder:
@staticmethod
def build(model_name: str = "small_sr", pretrained: bool = True,
upscale_factor: int = 3, image_size: Optional[int] = None) -> SuperResolution: ...
Import
from kornia.contrib import SuperResolution
from kornia.contrib.super_resolution import RRDBNetBuilder, SmallSRBuilder
I/O Contract
Inputs (SuperResolution.forward)
| Name |
Type |
Required |
Description
|
| images |
torch.Tensor or List[torch.Tensor] |
Yes |
Input images: list of (3, H, W) tensors or a batched tensor of shape (B, 3, H, W)
|
Inputs (RRDBNetBuilder.build)
| Name |
Type |
Required |
Description
|
| model_name |
str |
No |
Model variant: "RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B", or "RealESRGAN_x2plus" (default: "RealESRNet_x4plus")
|
| pretrained |
bool |
No |
Whether to load pretrained weights (default: True)
|
Inputs (SmallSRBuilder.build)
| Name |
Type |
Required |
Description
|
| model_name |
str |
No |
Model name, currently only "small_sr" supported (default: "small_sr")
|
| pretrained |
bool |
No |
Whether to load pretrained weights (default: True)
|
| upscale_factor |
int |
No |
Upscaling factor (default: 3)
|
| image_size |
int |
No |
Input image size; if None, uses default of 224
|
Outputs
| Name |
Type |
Description
|
| super_resolved |
torch.Tensor or List[torch.Tensor] |
Upscaled image tensor(s) with higher spatial resolution than input
|
Available Pretrained Models
| Model Name |
Scale |
Architecture |
Num RRDB Blocks
|
| RealESRGAN_x4plus |
4x |
RRDB |
23
|
| RealESRNet_x4plus |
4x |
RRDB |
23
|
| RealESRGAN_x4plus_anime_6B |
4x |
RRDB |
6
|
| RealESRGAN_x2plus |
2x |
RRDB |
23
|
| small_sr |
3x |
SmallSRNet |
N/A
|
Usage Examples
import torch
from kornia.contrib.super_resolution import RRDBNetBuilder, SmallSRBuilder
# Build a RealESRNet 4x super-resolution model
sr_model = RRDBNetBuilder.build(model_name="RealESRNet_x4plus", pretrained=True)
# Upscale an image
image = torch.rand(1, 3, 128, 128)
sr_output = sr_model(image)
print(sr_output.shape) # Approximately (1, 3, 512, 512) for 4x upscaling
# Save source and super-resolved images
sr_model.save(image, directory="sr_output")
# Build a lightweight SmallSR model
small_sr = SmallSRBuilder.build(model_name="small_sr", upscale_factor=3)
small_output = small_sr(image)
# Export to ONNX
onnx_model = sr_model.to_onnx(onnx_name="sr_model.onnx")
Related Pages
