Skip to main content

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

import torch

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

  • The evaluation metric is PSNR.
  • You can check the full benchmark here.
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


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.10-py3-none-any.whl (42.5 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page