super-image

适用于PyTorch的先进图像超分辨率模型。

安装

使用pip

pip install super-image

演示

立即尝试各种 模型 对您的图片进行测试。

快速开始

快速使用预训练模型将图片放大2倍、3倍和4倍。查看完整的模型列表 如下。

from super_image import EdsrModel, ImageLoader
from PIL import Image
import requests

url = 'https://paperswithcode.com/media/datasets/Set5-0000002728-07a9793f_zA3bDjj.jpg'
image = Image.open(requests.get(url, stream=True).raw)

model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)
inputs = ImageLoader.load_image(image)
preds = model(inputs)

ImageLoader.save_image(preds, './scaled_2x.png')
ImageLoader.save_compare(inputs, preds, './scaled_2x_compare.png')

预训练模型

预训练模型在各种规模下可用，并托管在令人惊叹的 huggingface_hub。默认情况下，这些模型在 DIV2K 上进行了预训练，这是一个用于训练的800张高质量（2K分辨率）图像的数据集，经过增强至4000张，并使用100张验证图像（编号为801至900）的dev集。

下方的排行榜显示了各种测试集（Set5、Set14、BSD100、Urban100）上各种规模下每个模型的PSNR / SSIM指标。数值越高越好。所有训练都进行到1000个周期（一些出版物，如a2n，在他们的实验中将训练周期增加到>1000个周期）。

缩放x2

缩放 x3

缩放 x4

您可以在下方找到笔记本，以便轻松运行预训练模型的评估

训练模型

我们需要huggingface datasets库来下载数据

pip install datasets

以下代码获取数据并对数据进行预处理/增强。

from datasets import load_dataset
from super_image.data import EvalDataset, TrainDataset, augment_five_crop

augmented_dataset = load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='train')\
    .map(augment_five_crop, batched=True, desc="Augmenting Dataset")                                # download and augment the data with the five_crop method
train_dataset = TrainDataset(augmented_dataset)                                                     # prepare the train dataset for loading PyTorch DataLoader
eval_dataset = EvalDataset(load_dataset('eugenesiow/Div2k', 'bicubic_x4', split='validation'))      # prepare the eval dataset for the PyTorch DataLoader

训练代码如下

from super_image import Trainer, TrainingArguments, EdsrModel, EdsrConfig

training_args = TrainingArguments(
    output_dir='./results',                 # output directory
    num_train_epochs=1000,                  # total number of training epochs
)

config = EdsrConfig(
    scale=4,                                # train a model to upscale 4x
)
model = EdsrModel(config)

trainer = Trainer(
    model=model,                         # the instantiated model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=eval_dataset            # evaluation dataset
)

trainer.train()