英文

多尺度残差网络用于图像超分辨率 (MSRN)

MSRN模型是在DIV2K上预训练的(800张训练图像,增强到4000张图像,100张验证图像),用于2倍、3倍和4倍图像超分辨率。该模型由Li等人在2018年的论文 Multi-scale Residual Network for Image Super-Resolution 中提出,并在 this repository 中首次发布。

图像超分辨率的目标是从单个低分辨率 (LR) 图像恢复出高分辨率 (HR) 图像。下图显示了真实图像 (HR)、双三次插值放大 x2 和模型放大 x2。

模型描述

MSRN模型提出了一种特征提取结构,称为多尺度残差块。该模块可以"自适应地检测不同尺度上的图像特征"并"发现图像的潜在特征"。

该模型还应用了由 Wang et al. (2021) 发明的平衡注意力(BAM)方法,以进一步提高结果。

使用和限制

您可以使用预训练模型将图像放大 2倍、3倍和4倍。您还可以使用训练器在自己的数据集上进行模型训练。

使用方法

该模型可以与 super_image 库一起使用:

pip install super-image

以下是使用预训练模型放大图像的方法:

from super_image import MsrnModel, 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 = MsrnModel.from_pretrained('eugenesiow/msrn-bam', scale=2)      # scale 2, 3 and 4 models available
inputs = ImageLoader.load_image(image)
preds = model(inputs)

ImageLoader.save_image(preds, './scaled_2x.png')                        # save the output 2x scaled image to `./scaled_2x.png`
ImageLoader.save_compare(inputs, preds, './scaled_2x_compare.png')      # save an output comparing the super-image with a bicubic scaling

训练数据

2倍、3倍和4倍图像超分辨率的模型是在 DIV2K 上进行预训练的。该数据集包含800张高质量(2K分辨率)的图像用于训练,增强到4000张图像,并使用100张验证图像(图像编号801至900)作为开发集。

训练过程

预处理

我们遵循 Wang et al. 的预处理和训练方法。通过使用双三次插值作为调整图像大小的方法,从高分辨率 (HR) 图像中减小 LR (低分辨率) 图像的尺寸,分别减小x2、x3和x4倍。在训练过程中,使用大小为64×64的LR输入的RGB块及其对应的HR块。数据增强应用于训练集的预处理阶段,从原始图像的四个角和中心创建五个图像。

我们需要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

预训练

该模型在GPU上进行训练。下面提供了训练代码:

from super_image import Trainer, TrainingArguments, MsrnModel, MsrnConfig

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

config = MsrnConfig(
    scale=4,                                # train a model to upscale 4x
    bam=True,                               # apply balanced attention to the network
)
model = MsrnModel(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()

评估结果

评估指标包括 PSNR SSIM

评估数据集包括:

下面的结果列表示为 PSNR/SSIM,它们与双三次插值基准进行比较。

|数据集 |放大倍数 |双三次插值 |msrn-bam ||--- |--- |--- |--- ||Set5 |2倍 |33.64/0.9292 | 38.02/0.9608 ||Set5 |3倍 |30.39/0.8678 | 35.13/0.9408 ||Set5 |4倍 |28.42/0.8101 | 32.26/0.8955 ||Set14 |2倍 |30.22/0.8683 | 33.73/0.9186 ||Set14 |3倍 |27.53/0.7737 | 31.06/0.8588 ||Set14 |4倍 |25.99/0.7023 | 28.78/0.7859 ||BSD100 |2倍 |29.55/0.8425 | 33.78/0.9253 ||BSD100 |3倍 |27.20/0.7382 | 29.65/0.8196 ||BSD100 |4倍 |25.96/0.6672 | 28.51/0.7651 ||Urban100 |2倍 |26.66/0.8408 | 32.08/0.9276 ||Urban100 |3倍 | | 29.26/0.8736 ||Urban100 |4倍 |23.14/0.6573 | 26.10/0.7857 |

您可以在预训练模型上轻松运行评估的笔记本如下:

BibTeX条目和引用信息

@misc{wang2021bam,
    title={BAM: A Lightweight and Efficient Balanced Attention Mechanism for Single Image Super Resolution}, 
    author={Fanyi Wang and Haotian Hu and Cheng Shen},
    year={2021},
    eprint={2104.07566},
    archivePrefix={arXiv},
    primaryClass={eess.IV}
}
@InProceedings{Li_2018_ECCV,
    author = {Li, Juncheng and Fang, Faming and Mei, Kangfu and Zhang, Guixu},
    title = {Multi-scale Residual Network for Image Super-Resolution},
    booktitle = {The European Conference on Computer Vision (ECCV)},
    month = {September},
    year = {2018}
}