MDSR模型在DIV2K上进行了预训练(800张训练图像,扩充为4000张,100张验证图像),用于2x,3x和4x图像超分辨。它在Lim等人(2017)的论文中首次推出,并在 this repository 中首次发布。
图像超分辨的目标是从单个低分辨率(LR)图像恢复高分辨率(HR)图像。下图显示了真实图像(HR),双三次插值和模型放大。
MDSR是一个使用更深和更宽的架构(32个ResBlocks和256个通道)来提高性能的模型。它使用全局和局部跳跃连接,并且放大是在网络末尾完成的。它不使用批归一化层(输入和输出具有类似的分布,归一化中间特征可能不可取),而是使用恒定缩放层来确保稳定的训练。使用L1损失函数(绝对误差)而不是L2(均方误差),作者通过经验显示出更好的性能,并且需要更少的计算。
此模型还应用了 Wang et al. (2021) 发明的平衡注意力(BAM)方法,以进一步改善结果。
您可以使用预训练的模型将图像放大2x、3x和4x。您还可以使用训练器在自己的数据集上训练模型。
该模型可以与 super_image 库一起使用:
pip install super-image
以下是使用预训练模型放大图像的方法:
from super_image import MdsrModel, 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 = MdsrModel.from_pretrained('eugenesiow/mdsr-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
2x、3x和4x图像超分辨率的模型是在 DIV2K 上预训练的,该数据集包含800个高质量(2K分辨率)图像进行训练,并扩充为4000个图像,使用100个验证图像(图像编号为801到900)。
我们遵循 Wang et al. 的预处理和训练方法。通过使用双三次插值作为调整大小方法,从高分辨率(HR)图像缩小x2、x3和x4倍生成低分辨率(LR)图像。在训练过程中,使用来自LR输入的大小为64×64的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, MdsrModel, MdsrConfig training_args = TrainingArguments( output_dir='./results', # output directory num_train_epochs=1000, # total number of training epochs ) config = MdsrConfig( scale=4, # train a model to upscale 4x bam=True, # apply balanced attention to the network ) model = MdsrModel(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表示,与双三次插值基线进行比较。
|Dataset |Scale |Bicubic |mdsr-bam ||--- |--- |--- |--- ||Set5 |2x |33.64/0.9292 |38/0.9607 ||Set5 |3x |30.39/0.8678 |35.07/0.9402 ||Set5 |4x |28.42/0.8101 |32.19/0.8949 ||Set14 |2x |30.22/0.8683 |33.68/0.9182 ||Set14 |3x |27.53/0.7737 |31.04/0.8582 ||Set14 |4x |25.99/0.7023 |28.73/0.7847 ||BSD100 |2x |29.55/0.8425 |33.77/0.9253 ||BSD100 |3x |27.20/0.7382 |29.62/0.8188 ||BSD100 |4x |25.96/0.6672 |28.5/0.7645 ||Urban100 |2x |26.66/0.8408 |32.04/0.9272 ||Urban100 |3x | |29.16/0.8717 ||Urban100 |4x |23.14/0.6573 |26.02/0.7834 |
您可以在下面找到一个轻松运行预训练模型评估的笔记本:
@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} }
@article{ahn2018fast, title={Fast, Accurate, and Lightweight Super-Resolution with Cascading Residual Network}, author={Ahn, Namhyuk and Kang, Byungkon and Sohn, Kyung-Ah}, journal={arXiv preprint arXiv:1803.08664}, year={2018} }