基于金字塔修复模块以及双边注意力机制的单阶段图像修复方法

4 years ago

source link: https://bbs.cvmart.net/articles/1387
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

基于金字塔修复模块以及双边注意力机制的单阶段图像修复方法

11个月前 ⋅ 1015 ⋅ 0 ⋅ 0

介绍一下我最新的一份工作，希望能够给大家提供一些参考。做的任务依旧是图像修复，但是其中的atttention模块其实可以嵌套到所有生成任务，希望大家有机会可以尝试一下，由于deadline我没有将其拓展到其他领域，个人认为这是个有意思的idea，也希望可以和大家讨论。
论文链接：https://arxiv.org/abs/1912.08642
源码地址：https://github.com/KumapowerLIU/One-Stage-Inpainting-with-Bilateral-Attention-and-Pyramid-Filling-Block

模型架构：

One-stage：
整个模型训练分为两步，但是测试的时候只有一个网络，其中第一次训练是利用简单的U-net修复出目标的结构，在第一次训练中我们的label没有纹理。第一次训练完成后，我们利用第一次训练的得到的参数进行第二次训练，第二次的label就是原图（有纹理），这样就能够在测试的时候节省很多时间，我们最后单张图片的修复的时间只有8.8ms（图片大小256×256）是已知最快的。同时在第二阶段训练时，我们加入了Bilateral Attention 和 Pyramid Filling Block提升效果。
双边注意力：

传统的attention只从值的相似度来考虑特征之间的相似性，但是这往往不准确，我们辅以距离信息来更加精确特征之间的相似性，所以我们的attention模块从值和距离两个角度来刻画特征性相似程度，类似于双边滤波，所以我们的attention模块叫做Bilateral Attention。具体来说，以当前的特征点为中心，计算其与周围8个点包括它自己的值的相似度（其实就是3×3的kernel），这里我们可以用点积来计算，并将这9个值分别乘以这些权重然后相加得到新的特征点。同时，我们以当前特征点为中心，利用高斯分布，来刻画其与周围24个点（5×5的kernel）的相似度，然后同样进行加权相加，得到新的特征点，最后这两个特征点拼接并降维得到最后的特征图。
金字塔修复模块：

我们利用已经被填满的高级语义特征逐步填充浅层的特征，具体来说我们从高到低采用金字塔架构逐步填充特征，并把这些特征利用skip-connection的方式连接到decoder，同时我们采用short-cut的方式将原特征图与填充完的特征图进行连接，这是因为在经过第一次训练阶段之后，我们的马赛克区域并不是完全无效的。

最后附上我们的效果图：