英文

像素关注网络(PAN)

PAN模型在DIV2K上进行了预训练(800张训练图像,扩充到4000张,100张验证图像),用于2x、3x和4x图像超分辨率。该模型在赵等人的论文 Efficient Image Super-Resolution Using Pixel Attention 中提出,并于 this repository 首次发布。

图像超分辨率的目标是从一张低分辨率(LR)图像中恢复出一张高分辨率(HR)图像。下面的图像显示了真实图像(HR)、双三次上采样和模型上采样。

模型描述

PAN模型提出了一种用于图像超分辨率的轻量级卷积神经网络。像素关注(PA)在形式上类似于通道关注和空间关注。然而,PA生成的是3D关注图而不是1D关注向量或2D图。这种注意力机制引入了更少的额外参数,但能产生更好的SR结果。

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

该模型非常轻量级,模型参数只有260k至270k(约1MB)。

使用目的和限制

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

使用方法

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

pip install super-image

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

from super_image import PanModel, 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 = PanModel.from_pretrained('eugenesiow/pan-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)图像。在训练过程中,使用大小为64×64的RGB图像块从LR输入中提取,同时使用相应的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, PanModel, PanConfig

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

config = PanConfig(
    scale=4,                                # train a model to upscale 4x
    bam=True,                               # apply balanced attention to the network
)
model = PanModel(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表示,并与双三次插值进行比较。

|数据集|放大倍数|双三次插值|PAN-BAM||---|---|---|---||Set5|2x|33.64/0.9292|37.7/0.9596||Set5|3x|30.39/0.8678|34.62/0.9371||Set5|4x|28.42/0.8101|31.9/0.8911||Set14|2x|30.22/0.8683|33.4/0.9161||Set14|3x|27.53/0.7737|30.83/0.8545||Set14|4x|25.99/0.7023|28.54/0.7795||BSD100|2x|29.55/0.8425|33.6/0.9234||BSD100|3x|27.20/0.7382|29.47/0.8153||BSD100|4x|25.96/0.6672|28.32/0.7591||Urban100|2x|26.66/0.8408|31.35/0.92||Urban100|3x| |28.64/0.861||Urban100|4x|23.14/0.6573|25.6/0.7691|

你可以在下面找到一个笔记本,轻松运行对预训练模型的评估:

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}
}
@misc{zhao2020efficient,
      title={Efficient Image Super-Resolution Using Pixel Attention}, 
      author={Hengyuan Zhao and Xiangtao Kong and Jingwen He and Yu Qiao and Chao Dong},
      year={2020},
      eprint={2010.01073},
      archivePrefix={arXiv},
      primaryClass={eess.IV}
}