Implementation:Kornia Kornia SmallSR
| Knowledge Sources | |
|---|---|
| Domains | Vision, Super_Resolution |
| Last Updated | 2026-02-09 15:00 GMT |
Overview
SmallSR implements a compact super-resolution neural network using efficient sub-pixel convolution for real-time image upscaling, based on the paper Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network by Shi et al.
Description
This module provides two classes within the Kornia library: SmallSRNet, the core super-resolution network that operates on single-channel (luminance) images, and SmallSRNetWrapper, which wraps the network with RGB-to-YCbCr and YCbCr-to-RGB color space conversions for end-to-end processing of color images. The network architecture consists of three convolutional layers with ReLU activations followed by a sub-pixel shuffle layer (nn.PixelShuffle). Pre-trained weights are available and downloaded automatically via CachedDownloader. The module also provides a weight_init function using orthogonal initialization.
Usage
Import this module when you need a lightweight, real-time super-resolution model. Use SmallSRNet directly for grayscale/luminance upscaling, or SmallSRNetWrapper for seamless RGB image super-resolution.
Code Reference
Source Location
- Repository: Kornia
- File: kornia/models/small_sr.py
- Lines: 1-99
Signature
class SmallSRNet(nn.Module):
def __init__(self, upscale_factor: int, inplace: bool = False, pretrained: bool = True) -> None: ...
def load_from_file(self, path_file: str) -> None: ...
def forward(self, x: torch.Tensor) -> torch.Tensor: ...
class SmallSRNetWrapper(nn.Module):
def __init__(self, upscale_factor: int = 3, pretrained: bool = True) -> None: ...
def forward(self, input: torch.Tensor) -> torch.Tensor: ...
def weight_init(model: nn.Module) -> None: ...
Import
from kornia.models.small_sr import SmallSRNet, SmallSRNetWrapper
I/O Contract
Inputs (SmallSRNet)
| Name | Type | Required | Description |
|---|---|---|---|
| upscale_factor | int | Yes | Factor by which to increase image resolution. |
| inplace | bool | No | Whether ReLU should operate in-place (default False). |
| pretrained | bool | No | Whether to load pre-trained weights (default True). |
Inputs (SmallSRNetWrapper)
| Name | Type | Required | Description |
|---|---|---|---|
| upscale_factor | int | No | Factor by which to increase image resolution (default 3). |
| pretrained | bool | No | Whether to load pre-trained weights (default True). |
Outputs
| Name | Type | Description |
|---|---|---|
| output (SmallSRNet) | torch.Tensor | Upscaled single-channel image of shape (B, 1, H*factor, W*factor). |
| output (SmallSRNetWrapper) | torch.Tensor | Upscaled RGB image of shape (B, 3, H*factor, W*factor). |
Architecture
The SmallSRNet consists of:
- conv1: Conv2d(1, 64, 5x5) with ReLU
- conv2: Conv2d(64, 64, 3x3) with ReLU
- conv3: Conv2d(64, 32, 3x3) with ReLU
- conv4: Conv2d(32, upscale_factor^2, 3x3)
- pixel_shuffle: PixelShuffle(upscale_factor)
The SmallSRNetWrapper processes RGB images by:
- Converting RGB to YCbCr
- Applying SmallSRNet to the Y (luminance) channel
- Bicubic-interpolating the Cb and Cr (chrominance) channels
- Concatenating and converting back to RGB
Usage Examples
import torch
from kornia.models.small_sr import SmallSRNet, SmallSRNetWrapper
# Single-channel super-resolution
sr_net = SmallSRNet(upscale_factor=3, pretrained=True)
y_channel = torch.rand(1, 1, 64, 64)
y_upscaled = sr_net(y_channel) # shape: (1, 1, 192, 192)
# Full RGB super-resolution
sr_wrapper = SmallSRNetWrapper(upscale_factor=3)
rgb_image = torch.rand(1, 3, 64, 64)
rgb_upscaled = sr_wrapper(rgb_image) # shape: (1, 3, 192, 192)