【百度网盘AI大赛 文档图像超分比赛】 SwinIR方案

【百度网盘AI大赛 文档图像超分比赛】 SwinIR方案

一、项目背景

        该项目基于百度网盘AI大赛——图像处理挑战赛：文档图像超分 比赛进行开发，下面将从比赛介绍和提交排名两个方面对项目背景进行具体阐述。

1.1 比赛介绍

• 背景：大家平时在使用手机进行拍照、扫描的时候，往往会需要将图片放大的场景，可是分辨率放大，图片中的元素都变得模糊起来，进而难以使用。本次比赛希望选手们通过算法等知识技术，帮助人们将因为放大而模糊的图片复原，提高图片分辨率，实现文档图像的“无损放大”，让其重新发挥作用。

• 分析：本次比赛提供了4篇参考文献：SwinIR: Image Restoration Using Swin Transformer（SwinIR）、Activating More Pixels in Image Super-Resolution Transformer（HAT）、Image Super-Resolution Using Very Deep Residual Channel Attention Networks（RCAN）和 Learned Image Downscaling for Upscaling using Content Adaptive Resampler（CAR）。CAR模型需要先对高分辨率图像下采样，再对下采样图像进行超分辨率，研究的更多是如何让下采样图像与真实的低分辨率图像尽可能相似。但在比赛测试中不可能提供高分辨率图像，因此首先排除选取该网络作为baseline。从目前论文中公布的自然图像超分辨评估结果来看，HAT要优于Swin IR优于RCAN，但是在实际应用中需要对算力、推理速度进行多方面考虑，Swin IR和HAT都是在多GPU上训练极长的时间才获得了论文中的SOTA性能，因此需要慎重选择。由于目前公开的baseline提供了RCAN模型，本项目基于Swin IR模型进行训练，并提供了HAT模型的paddle脚本文件，来解决当前公开项目缺乏另外2个模型训练案例的问题。

• 训练策略： 原始的Swin IR 采用 $64\times 64$ 的图像块进行训练，但是对于1K到2K测评图像来说，预测时间过长，这里我们采用 $128\times 128$ 的图像块进行多阶段训练。根据之前“模糊文档恢复比赛” 经验 ， Transformer类网络的优化往往存在着一些技巧型操作，比如“模糊文档恢复比赛”第一名对图片进行下采样之后直接获得了近三个点的增长。因此参赛者可以从域变换、数据增强策略等方面进行挖掘。截止11月8日，榜一大神的分数已经高达67.69，因此当前模型仍有极大的提升空间，同时 Transformer类网络也许并不是最优网络，需要参赛者自行探索。

• 比赛链接：https://aistudio.baidu.com/aistudio/competition/detail/493/0/introduction

1.2 提交排名

        当前公开版本在 A 榜分数为 0.64625, PSNR为 29.92764， SSIM为 0.9932, 单张图片测试耗时为 7.665s。在11月8日时，排名大约在11名左右。可作为 Swin IR 的预训练模型继续优化。

        当前项目的后续优化版本在 A 榜分数为 0.65014， PSNR为 30.6457， SSIM为 0.99382, 单张图片测试耗时为 6.6519s。在11月8日, 排名为第9名。

二、项目方案

        该项目采用SwinIR进行图像超分辨率，该模型的网络架构图如下:

        该网络由三个部分组成。浅层特征提取、深层特征提取和高质量图像重建。其中，深层特征提取模块由一系列残差 Swin Transformer模块组成。Swin IR在 图像超分辨率、图像降噪 和JPEG 压缩图像增强三个任务取得了SOTA的成绩，因此选作本次比赛的baseline模型。

三、数据说明

        该项目的数据来源于百度网盘AI大赛——图像处理挑战赛：文档图像超分的官方数据集。为了让网络更好地收敛，我们对数据集进行了数据增广、数据筛选等处理。

3.1 数据集介绍

        在本次比赛最新发布的数据集中，所有的图像数据均由真实场景采集得到，再通过技术手段进行相应处理，生成可用的脱敏数据集。该任务为image-to-image的形式，因此源数据和对应的超分后数据均以图片的形式来提供。

数据集构成

训练数据集构成

|- train_data  
    |- x
    |- x2
    |- x4

测试数据集构成

|- test_data  
    |- x

• 本次比赛最新发布的数据集共包含训练集、A榜测试集、B榜测试集三个部分，其中训练集共3000个样本，A榜测试集共200个样本，B榜测试集共200个样本；
• x 为原始图像数据，x2为原始数据分辨率放大两倍后的数据，x4为原始数据分辨率放大四倍后的数据（仅有训练集数据提供x2，x4图像 ，A榜测试集、B榜测试集数据均不提供）；
• x，x2 与 x4中的图片根据图片名称一一对应。

数据集下载

数据集类型	是否开放	百度网盘链接	AI Studio链接
训练集	已开放	百度网盘下载	AI studio引用
A榜测试集	已开放	百度网盘下载	AI studio引用
B榜测试集	未开放

数据集可视化展示

3.2 数据增广

• 由于Swin IR模型采用Transformer架构，因此对图像的尺寸较为敏感，过大的图像尺寸会导致显存溢出，过小的尺寸又会增加模型耗时，因此将模型裁剪为$128\times 128$ 的patch作为网络输入。

• 为了提高网络鲁棒性，在数据集读取过程中，对图像块进行了随机翻转和旋转操作。

3.3 数据筛选

• 采用“百度网盘AI大赛-模糊文档图像恢复比赛第1名方案”项目的方法，根据平均梯度，筛选掉图像纹理较少或者空白的图像块。

• 由于数据预处理过程较长，我们已经将裁剪后的图像块上传到 AI Studio数据集 中。 想要对数据进行预处理的同学可以参考/home/aistudio/work/data_preparation_all.py文件。

3.4 数据划分

• 裁剪后的数据集总共包含126837 张图像块, 不同阶段的数据集和划分比例如下表：

当前阶段	数据集内容	对应原始数据集压缩包	数据集大小	划分比例
Step 1	patch_dataset	01-03.zip	train 38507 val 786	0.98：0.02
Step 2	patch_dataset1	04-08.zip	train 66224 val 3486	0.95：0.05
Step 3	patch_dataset patch_dataset2	01-03.zip 09-10.zip	train 60613 val 6735	0.90：0.10

四、代码实现

• 训练细节: Swin IR 模型采用多阶段训练策略进行训练，训练环境为 V100 32G， 如需在 V100 16G 环境下训练，请降低 batch 数量。每个阶段训练细节如下表：

当前阶段	训练模式	损失函数	迭代次数	初始学习率	调整策略	batch
Step 1	eval	CharbonnierLoss + MS-SSIM	80k	1e-4	固定学习率	4
Step 2	train	L1Loss	210k	1e-4	MultiStepDecay	4
Step 3	train	CharbonnierLoss + MS-SSIM	540k	0.25*1e-4	固定学习率	4

• 与baseline差异： 与已有baseline不同，这里我们对文档图像的2倍超分辨和4倍超分辨进行联合训练, 2倍超分辨和4倍超分辨损失按 0.4 和 0.6 的比例进行求和。

• 模型保存： 网络每 500 次迭代保存一次临时模型,以便进行网络恢复; 每 5000 次迭代在测试集上进行一次测试, 选取测试分数最高的模型保存。目前最优模型保存在 /home/aistudio/work/best_model/model_Step3_515k.param 路径下。

# 解压数据集，虽然控制台可以解压，但是强烈建议在终端运行如下命令， patch_dataset_Step1.zip有近4万*3张图像块，输出过多
# 可能会导致控制台卡顿，整个解压过程耗时约半个小时左右
!unzip -d /home/aistudio/ /home/aistudio/data/data173301/patch_dataset_Step1.zip
!unzip -d /home/aistudio/ /home/aistudio/data/data173301/patch_dataset_Step2.zip
!unzip -d /home/aistudio/ /home/aistudio/data/data173301/patch_dataset_Step3.zip

# 安装paddle_msssim包
!pip install paddle_msssim

# 启动第一阶段训练代码
%cd /home/aistudio/work/
# 创建输出文件夹，用于模型保存
!mkdir output sample
!python train_Step1.py

# 启动第二阶段训练代码,此处需要加载Step1的最佳模型进行权重初始化
%cd /home/aistudio/work/
!python train_Step2.py

# 启动第三阶段训练代码,此处需要加载Step2的最佳模型进行权重初始化
%cd /home/aistudio/work/
/
!python train_Step3.py

五、效果展示

• 可视化结果保存于/home/aistudio/result/x2和/home/aistudio/result/x4文件夹。由于生成的超分图像过大,请下载到本地查看。

• 提交： 可直接下载/home/aistudio/submit_project.zip文件或自行打包/home/aistudio/predict_code内的脚本和模型进行提交。

%cd  /home/aistudio
!mkdir /home/aistudio/result/ /home/aistudio/result/x/  /home/aistudio/result/x2 /home/aistudio/result/x4
!cp /home/aistudio/work/best_model/model_Step3_515k.param /home/aistudio/predict_code
%cd /home/aistudio/predict_code/
!python predict.py /home/aistudio/result/x/  /home/aistudio/result/x2   /home/aistudio/result/x4

六、总结提高

• 目前Swin IR采用128×128 的图像块，batch size = 4 作为输入，耗时为 7.6s。 如果采用256×256的图像块或更大尺寸作为输入耗时可能降到5s以内; 因此可以先在128×128块图像块上进行预训练,再在256×256图像块上进行微调。

• Swin IR对图像块进行预测试时存在部分空白或纹理较少图像块，可以使用平均梯度筛选出这些图像块进行区分处理。对于纹理较少图像块采用cubic插值法,其他部分采用Swin IR进行超分辨率，进一步缩短模型处理时间。

• 此外,我们还实现了HAT模型, 该模型位于/home/aistudio/work/hat_arc.py路径下, 算力充足的小伙伴可以尝试使用HAT模型进行图像超分辨率。

七、参考项目

[1] 百度网盘AI大赛-模糊文档图像恢复比赛第1名方案

https://aistudio.baidu.com/aistudio/projectdetail/4500726?channelType=0&channel=0

此文章为搬运
原项目链接