SwinIR x4 Image Super-Resolution
This repository contains a fine-tuned SwinIR (Swin Transformer for Image Restoration) model for 4Γ image super-resolution.
The model takes a low-resolution RGB image as input and generates a 4Γ higher-resolution image.
Model Details
| Property | Value |
|---|---|
| Architecture | SwinIR |
| Task | Image Super-Resolution |
| Upscale Factor | 4Γ |
| Input Channels | 3 (RGB) |
| Input Image/Patch Size | 64 Γ 64 |
| Window Size | 8 |
| Embedding Dimension | 180 |
| Depths | [6, 6, 6, 6, 6, 6] |
| Attention Heads | [6, 6, 6, 6, 6, 6] |
| MLP Ratio | 2 |
| Upsampler | nearest + convolution |
| Residual Connection | 1conv |
| Framework | PyTorch |
Model Architecture
The model was created using the following SwinIR configuration:
model = net(
upscale=4,
in_chans=3,
img_size=64,
window_size=8,
img_range=1.0,
depths=[6, 6, 6, 6, 6, 6],
embed_dim=180,
num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2,
upsampler="nearest+conv",
resi_connection="1conv"
)
Training
The model was initialized from a pretrained SwinIR Γ4 checkpoint and fine-tuned on paired low-resolution and ground-truth images.
Training setup
- Loss: L1 Loss
- Optimizer: AdamW
- Learning Rate: 2e-5
- Weight Decay: 1e-4
- Gradient Clipping: 1.0
- Learning Rate Scheduler: Cosine Annealing
- Training Patch: 64 Γ 64 LR
- Corresponding Ground Truth Patch: 256 Γ 256
- Data Augmentation: Random horizontal and vertical flips
Model Checkpoint
The main model file is:
best_swinir.pth
This file contains the fine-tuned SwinIR model checkpoint selected based on validation PSNR.
Usage
To use the model, first create the SwinIR architecture with the configuration described above and load the checkpoint.
Example:
import torch
import sys
sys.path.append("/path/to/SwinIR")
from models.network_swinir import SwinIR as net
model = net(
upscale=4,
in_chans=3,
img_size=64,
window_size=8,
img_range=1.0,
depths=[6, 6, 6, 6, 6, 6],
embed_dim=180,
num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2,
upsampler="nearest+conv",
resi_connection="1conv"
)
checkpoint = torch.load(
"best_swinir.pth",
map_location="cpu"
)
if "model" in checkpoint:
model.load_state_dict(checkpoint["model"])
else:
model.load_state_dict(checkpoint)
model.eval()
Input and Output
Low-Resolution RGB Image
β
SwinIR
β
4Γ Super-Resolved Image
For example:
64 Γ 64 β 256 Γ 256
Repository Structure
.
βββ best_swinir.pth
βββ config.json
βββ README.md
Files
best_swinir.pthβ Fine-tuned SwinIR model checkpointconfig.jsonβ Model architecture configurationREADME.mdβ Model documentation
Intended Use
This model is intended for image super-resolution experiments where a low-resolution RGB image needs to be enhanced to a 4Γ higher spatial resolution.
Limitations
The model was trained and evaluated on the dataset used for the associated super-resolution competition. Performance may vary on images with characteristics significantly different from the training data.
License
This repository contains a fine-tuned model based on the SwinIR architecture. Please refer to the original SwinIR repository and its license for the underlying architecture and pretrained model.