论文速递 | 基于金字塔及双边注意力机制的图像修复新方法

加入极市专业CV交流群,与6000+来自腾讯,华为,百度,北大,清华,中科院等名企名校视觉开发者互动交流!更有机会与李开复老师等大牛群内互动!

同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总,行业技术交流。关注 极市平台 公众号 ,回复 加群,立刻申请入群~

本文为极市开发者投稿,如需转载,请联系小助手微信( cv-mart )

介绍一下我最新的一份工作,希望能够给大家提供一些参考。做的任务依旧是图像修复,但是其中的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)是已知最快的。同时在第二阶段训练时,我们加入了Bilateral Attention 和 Pyramid Filling Block提升效果。

双边注意力(Bilateral Attention):

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

金字塔修复模块(Pyramid Filling Block)

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

总结:我觉得最重要的是那个attention的想法,希望对大家的工作有用。

(0)

相关推荐