轻量级图像超分辨率自适应加权学习网络（AWSRN）

AWSRN模型在DIV2K（800张训练图像，扩充到4000张，100张验证图像）上进行了预训练，用于2x、3x和4x的图像超分辨率。它是由王等人在2019年的论文 Lightweight Image Super-Resolution with Adaptive Weighted Learning Network 中介绍的，并首次在 this repository 中发布。

图像超分辨率的目标是从单个低分辨率（LR）图像中恢复高分辨率（HR）图像。下面的图像显示了真实值（HR）、双三次上采样和模型上采样。

模型描述

深度学习在单图像超分辨率（SISR）任务中取得了巨大成功，近年来在性能上表现出色。然而，大多数基于卷积神经网络的SR模型需要大量计算，限制了它们在现实世界中的应用。为了解决这个问题，本文提出了一种轻量级的SR网络，命名为自适应加权超分辨率网络（AWSRN）。AWSRN中设计了一种新颖的局部融合块（LFB）用于高效的残差学习，它包括堆叠的自适应加权残差单元（AWRU）和局部残差融合单元（LRFU）。此外，还提出了自适应加权多尺度（AWMS）模块，充分利用重建层中的特征。AWMS由几个不同尺度的卷积组成，可以根据AWMS中自适应权重的贡献来删除冗余的尺度分支，以实现轻量级网络。常用数据集上的实验结果表明，与具有相似参数和计算开销的最先进方法相比，所提出的轻量级AWSRN在×2、×3、×4和×8尺度因子上都达到了优越的性能。

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

使用目的和限制

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

使用方法

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

pip install super-image

以下是如何使用预训练模型放大您的图像：

from super_image import AwsrnModel, 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 = AwsrnModel.from_pretrained('eugenesiow/awsrn-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）图像的大小减小2倍、3倍和4倍，从而创建低分辨率（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, AwsrnModel, AwsrnConfig

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

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

评估数据集包括：

• Set5- Bevilacqua et al. (2012)
• Set14- Zeyde et al. (2010)
• BSD100- Martin et al. (2001)
• Urban100- Huang et al. (2015)

下面的结果列显示为PSNR/SSIM，它们与双三次基线进行比较。

|数据集|尺度|双三次|awsrn-bam||---|---|---|---||Set5|2x|33.64/0.9292|37.99/0.9606||Set5|3x|30.39/0.8678|35.05/0.9403||Set5|4x|28.42/0.8101|32.13/0.8947||Set14|2x|30.22/0.8683|33.66/0.918||Set14|3x|27.53/0.7737|31.01/0.8581||Set14|4x|25.99/0.7023|28.75/0.7851||BSD100|2x|29.55/0.8425|33.76/0.9253||BSD100|3x|27.20/0.7382|29.63/0.8188||BSD100|4x|25.96/0.6672|28.51/0.7647||Urban100|2x|26.66/0.8408|31.95/0.9265||Urban100|3x| |29.14/0.871||Urban100|4x|23.14/0.6573|26.03/0.7838|

您可以在下面找到一个笔记本，用于方便地对预训练模型进行评估：

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}
}

@article{wang2019lightweight,
  title={Lightweight Image Super-Resolution with Adaptive Weighted Learning Network},
  author={Wang, Chaofeng and Li, Zhen and Shi, Jun},
  journal={arXiv preprint arXiv:1904.02358},
  year={2019
}