模型:
eugenesiow/msrn-bam
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,它们与双三次插值基准进行比较。
|数据集 |放大倍数 |双三次插值 |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 |
您可以在预训练模型上轻松运行评估的笔记本如下:
@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} }