StudioSR is a PyTorch library providing implementations of training and evaluation of super-resolution models. StudioSR aims to offer an identical playground for modern super-resolution models so that researchers can readily compare and analyze a new idea. (inspired by PyTorch-StudioGan)
From PyPI
pip install studiosrgit clone https://github.com/veritross/studiosr.git
cd studiosr
python3 -m pip install -e .Documentation along with a quick start guide can be found in the docs/ directory.
$ python -m studiosr --image image.png --scale 4 --model swinirfrom studiosr.models import SwinIR
from studiosr.utils import imread, imwrite
model = SwinIR.from_pretrained(scale=4).eval()
image = imread("image.png")
upscaled = model.inference(image)
imwrite("upscaled.png", upscaled)from studiosr import Evaluator, Trainer
from studiosr.data import DIV2K
from studiosr.models import SwinIR
dataset_dir="path/to/dataset_dir",
scale = 4
size = 64
dataset = DIV2K(
dataset_dir=dataset_dir,
scale=scale,
size=size,
transform=True, # data augmentations
to_tensor=True,
download=True, # if you don't have the dataset
)
evaluator = Evaluator(scale=scale)
model = SwinIR(scale=scale)
trainer = Trainer(model, dataset, evaluator)
trainer.run()from studiosr import Evaluator
from studiosr.models import SwinIR
from studiosr.utils import get_device
scale = 2 # 2, 3, 4
dataset = "Set5" # Set5, Set14, BSD100, Urban100, Manga109
device = get_device()
model = SwinIR.from_pretrained(scale=scale).eval().to(device)
evaluator = Evaluator(dataset, scale=scale)
psnr, ssim = evaluator(model.inference)| Method | Scale | Training Dataset | Set5 | Set14 | BSD100 | Urban100 |
|---|---|---|---|---|---|---|
| EDSR | x 4 | DIV2K | 32.485 | 28.814 | 27.721 | 26.646 |
| RCAN | x 4 | DIV2K | 32.639 | 28.851 | 27.744 | 26.745 |
| SwinIR | x 4 | DF2K | 32.916 | 29.087 | 27.919 | 27.453 |
| HAT | x 4 | DF2K | 33.055 | 29.235 | 27.988 | 27.945 |
| Method | Scale | Training Dataset | Set5 | Set14 | BSD100 | Urban100 |
|---|---|---|---|---|---|---|
| EDSR | x 3 | DIV2K | 34.680 | 30.533 | 29.263 | 28.812 |
| RCAN | x 3 | DIV2K | 34.758 | 30.627 | 29.302 | 29.009 |
| SwinIR | x 3 | DF2K | 34.974 | 30.929 | 29.456 | 29.752 |
| HAT | x 3 | DF2K | 35.097 | 31.074 | 29.525 | 30.206 |
| Method | Scale | Training Dataset | Set5 | Set14 | BSD100 | Urban100 |
|---|---|---|---|---|---|---|
| EDSR | x 2 | DIV2K | 38.193 | 33.948 | 32.352 | 32.967 |
| RCAN | x 2 | DIV2K | 38.271 | 34.126 | 32.390 | 33.176 |
| SwinIR | x 2 | DF2K | 38.415 | 34.458 | 32.526 | 33.812 |
| HAT | x 2 | DF2K | 38.605 | 34.845 | 32.590 | 34.418 |
StudioSR is an open-source library under the MIT license.
