PyTorch library to accelerate super-resolution research
Project description
StudioSR
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)
Installation
From PyPI
pip install studiosr
From source (Editable)
git clone https://github.com/veritross/studiosr.git
cd studiosr
python3 -m pip install -e .
Documentation
Documentation along with a quick start guide can be found in the docs/ directory.
Quick Example
$ python -m studiosr --image image.png --scale 4 --model swinir
from 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)
Train
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()
Evaluate
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)
Benchmark
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 |
License
StudioSR is an open-source library under the MIT license.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
No source distribution files available for this release.See tutorial on generating distribution archives.
Built Distribution
studiosr-0.1.14-py3-none-any.whl
(46.0 kB
view hashes)
Close
Hashes for studiosr-0.1.14-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 8cb316049bcb37782943b568ec9cacb10275619a8c451817ab785ca17b9ec220 |
|
MD5 | 1b39eb4c872837945867e2c4ebafd11e |
|
BLAKE2b-256 | 950516d4db7a140f4ad04da8ccf961ad94762a56a2009c794a242eaffcc1e2db |