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

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.14-py3-none-any.whl (46.0 kB view details)

Uploaded Python 3

File details

Details for the file studiosr-0.1.14-py3-none-any.whl.

File metadata

  • Download URL: studiosr-0.1.14-py3-none-any.whl
  • Upload date:
  • Size: 46.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/5.1.0 CPython/3.11.9

File hashes

Hashes for studiosr-0.1.14-py3-none-any.whl
Algorithm Hash digest
SHA256 8cb316049bcb37782943b568ec9cacb10275619a8c451817ab785ca17b9ec220
MD5 1b39eb4c872837945867e2c4ebafd11e
BLAKE2b-256 950516d4db7a140f4ad04da8ccf961ad94762a56a2009c794a242eaffcc1e2db

See more details on using hashes here.

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