百度网盘AI大赛-水印智能消除赛第3名方案

Posted 2022-06-13 ACALJJ32

百度网盘AI大赛-水印智能消除赛第3名方案

项目简介

本项目是百度网盘AI大赛——图像处理挑战赛：百度水印智能消除赛第3名的解决方案。在该项目中，我们对本次赛题的任务与数据进行了分析，并在此基础上详细介绍了我们模型的搭建思路，同时开源了我们的代码。我们本次比赛的思路更多的还是反映一些比较基础的探索上，没有非常强的创新性，所以比较适合刚刚入门深度学习以及图像处理比赛的选手，掌握一些最最基本思路，即使是新手也可以拿到一个非常理想的名次。

比赛地址：https://aistudio.baidu.com/aistudio/competition/detail/209/0/introduction

我们的代码：
Github（https://github.com/ACALJJ32/watermark_remove_paddle_slbr.git）
AI Studio：（https://aistudio.baidu.com/aistudio/projectdetail/4086295?contributionType=1&shared=1）

赛题介绍

日常生活中带有水印的图片很常见，即使是PS专家，也很难快速且不留痕迹的去除水印。而使用智能去除水印的算法，可以快速自动去除图片中的水印。本次比赛我们需要通过深度学习技术训练模型，对给定的真实场景下采集得到的带有水印的图片进行处理，并最终输出处理后的扫描结果图片。

数据分析

在本次比赛最新发布的数据集中，所有的图像数据均由真实场景采集得到，再通过技术手段进行相应处理，生成可用的脱敏数据集。本次比赛也是一个非常典型的image-to-image的任务。我们下载了官方提供的数据集后，发现训练集是非常大的，训练集共1012550个样本，A榜测试集共5000个样本，我们随机抽取了一部分样本如下：

通过观察可以发现，官方提供的训练数据集里面的水印图案还是有非常明显的特点的，我们总结了几点如下：

• 训练集中的水印图案基本都是字体，基本没有其他图案的水印；一张图像中的水印颜色是不会变化的。

• 不同图像中的水印在一张图中的分布范围是不一样的。分布范围广的几乎整张图像都会有水印；分布范围小的水印会位于图像的中心位置。

• 所有图像中的水印图案分布是类似，都是以一种类似的方向倾斜的。

当我们总结了这些特点后，大概会思考到用什么样子的网络结构去处理不同的问题。比如说为了解决分布范围不同，我们通常会优先考虑U-Net structure或者dilation conv来提取不同尺度的特征；水印的颜色变换很小，可以尝试用channel-attention来让模型捕捉到不同通过通道之间的关系。不过，最快的解决方法还是好好看看官方提供的学习资料，本次比赛官方给了三个可以参考的开源代码，我们选择了时间比较近的SLBR。

SLBR：https://github.com/bcmi/SLBR-Visible-Watermark-Removal

数据处理

我们在数据处理上并没有花费太多的心思，处理流程如下：

水平翻转是一种常见的图像数据增强的方法。为什么使用随机裁剪而不是按照步长裁剪？因为这次的训练数据已经非常充足了，按照步长裁剪后数据集会更大，随机裁剪输入已经足够了。

模型选择

为什么这里讲模型选择而不是讲改进思路？因为一个模型的它一定是存在某个方面的局限性的，可以是性能也可以是速度，如果我们一开始就只选定了一个模型然后去不断改进，理论上当然可以得到一个不错的分数，但是这个过程是比较辛苦的。比赛的初期官方提供了一个baseline，其实就是一个非常简单的U-Net，我是非常不建议从这个baseline开始改的，因为提分的过程一定会非常吃力，所以我们还是把精力花在了阅读paddle文档，将SLBR的模型改写成了paddle的形式。我们也没有尝试过直接将.pth转换.pdparam，因为经验和直觉告诉我们，很可能遇到什么问题导致模型转换不过来。我们选择的模型结构如下：

改进思路

这里我们先说明一下，我们最终的比赛方案还是这个SLBR模型，里面的模块基本没有再做改进，这里可能让大家比较失望，什么都没动的模型还能拿第三名会不会太离谱，但是事实确实如此，而且我们也是做了一系列尝试发现没有效果后，最终又回到了起点。这部分将详细介绍我们是怎么进行魔改最后又回到开始的。

SLBR的思路就是将水印去除的思路分成两个阶段，第一阶段尝试生成一张粗处理的图像以及一张水印的mask图，第二阶段将第一阶段的图像和mask图一起输入再训练。我们在实践的时候先是尝试去实现Coarse Stage，Refinement Stage的网络其实我们也是同步去实现的，但是第二阶段的网络会占用非常大的显存，直接GPU out of memory，所以SLBR的Refinenment Stage网络结构我们在很长一段时间都是放弃的。第一阶段的Coarse Stage改好后，我们在训练集上直接训练了一个epoch，提交验证的分数如下：

可以看到这是一个非常理想的开始，MSSSIM有0.99735，PSNR有40.7081dB，要比baseline高好几个level，这也给了我们很大的信心。

紧接着最容易想到的就是用其他网络去替换原论文的Refinement Stage结构，我们最先想到的还是用NAFNet（Simple Baselines for Image Restoration）去替换，该网络结构如下：

我们的对比实验结果如下：

可以看到，加上NAFNet后，我们的PSNR直接高了将近2dB.

在接下来的改进中，我们尝试在NAFNet中用skip connection，因为对于图像处理来讲，skip connection在信息的补充上也有重要的作用，但是对于本次任务，直接用skip connnection会把水印的图案引入，所以我们考虑加两层conv，起到一个过滤的作用，示意图如下：

得到的对比实验如下：

可以看到添加skip connection后，又会高将近0.5dB. 以上基本都是有效的尝试，我们也进行过一些其他的尝试，有一次大的改动是这样的，我们想用一种并行的连接方式来进一步提高模型性能，网络结构是我们NTIRE2022 HDR的思路：

这样的一个设计是为了让模型可以捕捉到不同分辨率特征，从而学习到更强的表征，实验结果如下：

可以看到这次改进的分数只有35.30dB，这说明一个过大的模型不一定就能得到很高的分数。我们还进行过很多的尝试，这里列了一张excel表格，是我们做过的相关实验：

这里想说明的意思就是，对于模型的改进，我们还是要脚踏实地一点点来，一下子拍脑门想到的方案大概率是不靠谱，好的模型都是做实验做出来的。

小结

• 待改进的点

  本次比赛中，我们的思路还是存在很多改进的地方。水印消除的任务中，我们了解到mask feature对实验结果的影响是非常高的，如何生成准确稳定的mask是任务的关键，我们也尝试过对训练数据中的数据对直接进行相减取绝对值来生成mask，在用这样的mask引导模型生成对应的mask，但是这样做的效果也并不理想，所以我们暂时放弃考虑了这个点，如果这个问题处理得比较好，那么分数一定会得到更明显的提升。我们也尝试直接用vgg loss来改进生成的mask图，但是验证的分数也不理想，提交结果如下：

• 关于打比赛

  打比赛注意节奏，基本上是：数据分析+调研 => 对各种模型进行实验对比，选择最好的模型，站在巨人的肩膀上 => 不断进行改进

  打比赛也是一个快速提高自己能力的过程，参加过一个完整的比赛后，我们的视野、能力、经验都会有非常显著的提升，所以在学生期间，还是非常建议大家去打打比赛，对自己帮助很大。