本发明涉及图像修复技术领域,特别是涉及一种基于注意力机制的图像修复方法。
背景技术:
图像修复任务(imageinpainting),是从给定的受损图像中合成图像缺失区域中的替代内容,、且使得修复的图像在视觉上逼真和在语义上合理。图像修复可在其他应用中使用,如图像编辑,当图像中存在分散人注意力的场景元素时,如人或者物体(通常是不可避免的),允许用户移除图像中不需要的元素,同时在空白区域填充视觉和语义上合理的内容。
生成对抗网络启发自博弈论中二人零和博弈的思想,具有生成式网络和判别式网络两个网络,利用它们间相互竞争从而不断提升网络性能,最终达到平衡。基于对抗生成网络思想,衍生出许多变种网络,并且这些网络在图像合成、图像超分、图像风格转换和图像修复等方面都取得了显著的进步。图像修复的研究,包括图像补全、图像去水印、图像去雨滴和图像去雾都得到了研究者们的关注。
人类的注意力机制(attentionmechanism)是从直觉中得到,它是人类利用有限的注意力资源从大量信息中快速筛选出高价值信息的手段。深度学习中的注意力机制借鉴了人类的注意力思维方式,被广泛的应用在自然语言处理(naturelanguageprocessing,nlp)、图像分类及语音识别等各种不同类型的深度学习任务中,并取得了显著的成果。
随着科技不断发展,人们在不同领域的需求也在相应提高,包括电影广告动画制作和网络游戏等,逼真的图像修复技术对用户的良好体验具有重要意义。
因此,在此背景下,开发基于注意力机制的图像修复方法,使得修复后的图像在视觉上逼真和在语义上合理,具有重要的意义。
技术实现要素:
本发明的目的是针对现有技术中存在的技术缺陷,而提供一种基于注意力机制的图像修复方法,通过引入了重建损失函数、风格损失函数、全变分损失函数和对抗损失函数作为约束,提高网络的鲁棒性,使得能生成视觉逼真和语义合理的修复图像。
为实现本发明的目的所采用的技术方案是:
一种基于注意力机制的图像修复方法,包括步骤:
s1.预处理数据集中图像,得到训练数据集和测试数据集,预处理后的数据集中包括受损图像和真实图像(groundtruth),受损图像是真实图像通过乘以一个二值掩码得到,预处理后的图像尺寸大小一致;
s2.使用训练数据集训练模型,得到能对受损图像进行修复的神经卷积网络模型;
上述模型中包含一个生成器和一个判别器;生成器结构为u-net型网络结构,生成器包含16个卷积层,前8个卷积层后面均接有第一激活层,前8个卷积层逐渐减小特征图,前8个卷积层中的第2、5个卷积层后分别跟一个注意力层,用于学习特征图内的长远依赖关系以此获得更好的编码信息;后8个卷积层前后分别接有激活层和正则化层,后8个卷积逐渐增加特征图的大小到原图,后8个卷积层中的第3、7个卷积层后分别跟一注意力层,目的在于学习特征图内的长远依赖关系以此获得生成质量更好的生成图像。判别器由5个卷积层堆叠而成,用于判断图像是否真实;
训练集中受损图像作为模型输入,对应的真实图像作为目标生成图像,训练模型中生成器和判别器完成图像修复任务;通过生成器对输入的受损图像处理,得到修复图像;
判别器接受输入的目标生成图像(groundtruth)和模型生成的修复图像,通过修复图像与真实图像在判别器中进行对抗损失的计算,预测模型生成的修复图像是否逼近于目标生成图像;
模型迭代多次达到稳定后完成模型的训练,得到最终的图像修复模型;
s3.使用训练好的图像修复模型在测试数据集上测试其图像修复性能,给定受损图像模型即可自动输出修复好的图像。
其中,注意力模块的处理步骤如下:
首先将从之前隐藏层中得到的图像特征图x映射到两个隐空间f、g中,然后计算注意力得分,其中f(x)=wfx,g(x)=wgx,wf和wg均为可学习的参数,
注意力得分计算如下:
其中sij=f(xi)tg(xj),βj,i表示模型在生成第j个区域时对第i个位置的注意力程度,n表示特征图上的总区域数量,
注意力层的输出o=(o1,o2,...,oj,...,on),其中,oj可以表示为:
其中,h(xi)=whxi,v(xi)=wvxi,wh和wv均为可学习的参数,wf,wg,wh和wv均以卷积核为1×1的卷积层实现,
将注意力层的输出乘以一个缩放参数,并与输入的特征图相加,可得到:
yi=γoi+xi
其中,yi表示生成的第i个位置,oi表示注意力层的输出,xi表示输入的特征图,γ为平衡因子。
进一步的,步骤s2包括:
s21.使用凯明初始化方法初始化模型中的生成器和判别器的权重参数,其中生成器的损失函数为lg,判别器的损失函数为ld;
s22.将受损图像输入到生成器中,生成器输出与真实图像尺寸大小一致放入修复图像,将修复图像和真实图像作为判别器的输入,依次迭代训练使生成器的损失函数lg和判别器的损失函数为ld降至趋于稳定;
s23.交替训练模型中的生成器和判别器,当损失函数趋于稳定后得到最终的图像修复模型。
所述生成器和判别器的损失函数如下:
ld=-[ey~p(y)[logd(y)]+ex~p(x)[log(1-d(g(x)))]]]
其中,λ1,λ2,λ3,λ4为平衡因子,用于调整生成器中各个损失函数所占权重;x,y分别为受损图像和真实图像,g,d分别为生成器和判别器,l1为图像之间的l1范数度量,作为生成器的重建损失函数,lstyle为风格迁移损失函数,ltv为全变差损失函数,
所述训练数据集中输入图像和真实图像对记为[x,y],x为受损图像,y为目标生成图像;
其中,所述生成器的重建损失函数为:
l1=ex,y~p(x,y)||g(x)-y||1
其中x,y分别为从联合分布p(x,y)中采样出来的受损图像和对应的真实图像,e(*)表示取平均操作,||*||1表示l1范数,g(*)为生成器对应的映射函数。
其中,所述生成器的风格迁移损失函数为:
lstyle=αlc+βls
其中,α,β为平衡因子,lc,ls分别为内容损失函数和风格损失函数。
其中,所述的风格迁移损失函数中内容损失函数为:
其中,φj(*)为预训练的vgg网络的第j层,
其中,所述的风格迁移损失函数中风格损失函数为:
其中,ωi为计算在预训练的vgg网络第i层的风格损失函数,
其中,所述生成器的全变差损失函数为:
其中,hwc分别为图像的高、宽和通道数。
其中,所述生成器的对抗损失函数为:
其中,x,y分别为从联合分布p(x,y)中采样出来的受损图像和对应的真实图像,e(*)表示取平均操作,g,d分别为生成器和判别器对应的映射函数。
本发明基于注意力机制的图像修复方法,使用u-net作为生成器的网络结构,结合多种损失函数(重建损失函数,风格迁移损失函数,全变差损失函数和对抗损失函数),使得模型收敛更快,效果更好,泛化能力更强;可生成视觉上逼真和语义上合理的修复图像。
本发明使用生成网络,提高了模型容量和加快训练速度,提高模型的泛化能力和加快训练速度;引入了判别网络,使生成的修复图像更加接近真实图像,显著的提高了生成的修复图像的视觉质量。
本发明使用注意力机制,模仿人类的注意力思维方式,使得模型在补全受损图像时,获得需要重点关注的目标区域,使生成的修复图像更加接近真实图像,显著的提高了生成的修复图像的视觉质量。
附图说明
图1是本发明在测试数据集中一张受损图像上的测试结果,从左往右依次为受损图像,生成的修复图像和真实图像(groundtruth)。
图2是本发明基于注意力机制的图像修复方法的流程图;
其中:x表示输入的受损图像,y表示真实图像,
图3是本发明中所使用到的注意力层的示意图。
具体实施方式
为使本发明的目的、技术方案和优点更加清楚明白,以下结合具体实施例,并参照附图,对本发明作进一步的详细说明。应当理解,此处所描述的具体实施例仅仅用以解释本发明,并不用于限定本发明。
本发明通过一种基于注意力机制的图像修复方法来学习一组高度复杂的非线性的变换,用于把受损图像映射到修复图像上,同时使得修复后的图像在视觉上逼真和语义上合理。
参见附图所示,本发明基于注意力机制的图像修复方法,包括以下步骤:
步骤s1,先对imagenet数据集中的图像进行预处理。
第一,imagenet数据集中图像的尺寸大小并非统一,故本发明重新调整imagenet中训练数据集和测试数据集中的图像尺寸大小为256*256;
第二,生成各种类型的二值掩码(mask);
第三,对重新调整大小后的训练数据集增广,以增加训练数据集中的图像数量,包括随机水平翻转,随机颜色变换;
第四,将二值掩码乘上训练数据集中重新调整大小后的图像上生成受损图像,受损图像和尺寸调整后的图像(真实图像)作为模型的输入;将二值掩码乘上测试数据集中重新调整大小后的图像上生成测试图像。使用该测试图像,测试其模型的泛化性能。
步骤s2,利用步骤s1所述的训练数据集,训练模型中的生成器和判别器,以用来完成图像修复任务。
在模型的生成器中,利用前8个卷积层结构进行特征提取,每经过一个卷积层特征图的大小减半,直至特征图大小降为1*1,然后通过后8个卷积层进行图像的重建,每经过一个卷积层特征图的大小增大一倍,直复原为输入图像的大小(256*256)。前8个卷积层中每个卷积层后均皆有激活层(leakyrelu);后8个卷积层中,先有激活层(relu)激活后输入给卷积层,卷积层后紧跟正则化层。在前8个卷积层中,第2和5个卷积层后跟有一个注意力层,用于学习特征图内的长远依赖关系,获得重要关注区域。在后8个卷积层中,第3和7个卷积层后跟有一个注意力层,以生成质量更好的修复图像
其中,生成器中的第一个卷积层的输入通道数,输出通道数,滤波器大小,步长和填充分别是3,64,4,2,1。特征图尺寸大小减小这一阶段中卷积层的输出通道逐渐增大一倍,直到输出通道为512后保持不变,其滤波器大小,步长和填充均为4,2,1保持不变。在特征图增大这一阶段,卷积层为特殊的反卷积层,其作用和普通卷积层相反,用于增大特征图的大小。前四个反卷积层的的输出通道均为512,后3个反卷积层的输出通道数逐渐减小一倍,最后一层反卷积层的输出通道数为3,用于生成rgb的三通道图像。反卷积层的滤波器大小,步长和填充分别为4,2,1,且整个特征图尺寸大小正大阶段保持不变。在u-net网络结构中,除了最外围的卷积层和反卷积层外,其余的卷积层和对应的反卷积层之间均有跳跃连接。跳跃连接可以解决梯度消失和梯度爆炸的问题,也可加快模型的收敛的速度。
其中,注意力模块的处理步骤如下:
首先将从之前隐藏层中得到的图像特征图x映射到两个隐空间f、g中,然后计算注意力得分,其中f(x)=wfx,g(x)=wgx,wf和wg均为可学习的参数,
注意力得分计算如下:
其中sij=f(xi)tg(xj),βj,i表示模型在生成第j个区域时对第i个位置的注意力程度,n表示特征图上的总区域数量,
注意力层的输出o=(o1,o2,...,oj,...,on),其中,oj可以表示为:
其中,h(xi)=whxi,v(xi)=wvxi,wh和wv均为可学习的参数,wf,wg,wh和wv均以卷积核为1×1的卷积层实现,
将注意力层的输出乘以一个缩放参数,并与输入的特征图相加,可得到:
yi=γoi+xi
其中,yi表示生成的第i个位置,oi表示注意力层的输出,xi表示输入的特征图,γ为平衡因子。
判别器网络结构为由卷积层,正则化层(batchnorm)和激活层(leakyrelu)堆叠而成,其中卷积层滤波器大小,步长,填充分别为4,2,1,在本发明中卷积层个数为5,这一部分作为图像的特征提取,最后一个卷积层的输出通道数为1,用于判断输入图像是否真实。该判别器的输入为生成器生成的修复图像
该步骤中,利用受损图像x模型的输入,真实的无受损图像作为生成目标,交替训练模型中的生成器和判别器完成图像修复任务。在图像修复任务中,通过生成器对输入进行处理得到生成的修复图像,与真实图像在判别器中进行对抗损失计算。迭代多次达到稳定后完成模型的训练。
本发明中利用所述对抗生成网络的高度非线性拟合能力和注意力机制的长远依赖学习能力,针对图像修复这个任务,构造以受损图像和真实图像作为输入,修复图像作为输出的深度网络。特别的,网络通过外加损失函数的限制可以使得修复后的图像在视觉上逼真和在语义上合理。这样,通过如图2所示的网络,可以训练一个基于注意力机制的图像修复模型,从而完成图像的修复任务。在测试阶段,直接使用受损图像作为模型的输入,得到可视化的效果图,如图1所示
具体地,基于注意力机制的图像修复模型中生成器和判别器的损失函数如下所示:
ld=-[ey~p(y)[logd(y)]+ex~p(x)[log(1-d(g(x)))]]]
其中,λ1,λ2,λ3,λ4为平衡因子,用于调整生成器中各个损失函数所占的权重。x,y分别为受损图像和真实图像。g,d分别为生成器和判别器。l1为图像之间的l1范数度量,作为生成器的重建损失函数,lstyle为风格迁移损失函数,ltv为全变差损失函数,
上述的基于注意力机制的深度神经网络模型,主要是完成图像的修复任务,所述模型中生成器和判别器的最终目标为lg和ld这两个损失函数降至最低并且保持稳定。
其中,所述训练数据集中输入图像和真实图像对记为[x,y],x为受损图像,y为目标真实图像;
其中,所述生成器的重建损失函数为:
l1=ex,y~p(x,y)||g(x)-y||1
其中x,y分别为从联合分布p(x,y)中采样出来的受损图像和对应的真实图像,e(*)表示取平均操作,||*||1表示l1范数,g(*)为生成器对应的映射函数。
其中,所述生成器的风格迁移损失函数为:
lstyle=αlc+βls
其中,α,β为平衡因子,lc,ls分别为内容损失函数和风格损失函数。
其中,所述的风格迁移损失函数中内容损失函数为:
其中,φj(*)为预训练的vgg网络的第j层,
其中,所述的风格迁移损失函数中风格损失函数为:
其中,ωi为计算在预训练的vgg网络第i层的风格损失函数,
其中,所述生成器的全变差损失函数为:
其中,hwc分别为图像的高、宽和通道数。
其中,所述生成器的对抗损失函数为:
其中,x,y分别为从联合分布p(x,y)中采样出来的受损图像和对应的真实图像,e(*)表示取平均操作,g,d分别为生成器和判别器对应的映射函数。
本发明中,所述基于注意力机制的图像修复模型具体进行如下训练的步骤:
步骤s21:初始化图像修复任务的网络权重参数,λ1,λ2,λ3,λ4分别设置为10,240,0.1,1,批处理大小设为64,学习率设为10-4,并在整个训练过程中保持不变;
步骤s22.将受损图像输入到生成器中,生成器输出与真实图像大小一致的修复图像,将修复图像和真实图像作为判别器的输入,依次迭代训练使生成器的损失函数lg和判别器的损失函数为ld降至趋于稳定。
s23.交替训练深度模型中的生成器和判别器,当损失函数趋于稳定后得到最终的图像修复模型。
为了详细说明本发明的具体实施方式及验证本发明的有效性,将本发明提出的方法应用于一个公开的数据集(image),该数据集中是根据wordnet层次结构组织的图像数据集,大概有一千四百多万张图像,两千多万个类别,大类别包括鸟、蔬菜和食物等等。imagenet中的训练数据集大概有一百二多万张图像。
本发明选择imagenet的训练数据集训练所述模型中的生成器和判别器,选择imagenet中的测试集构造本发明所需的测试集,用于测试模型的泛化性能。
所述基于注意力机制的图像修复方法中数据处理方式如下:
第一,imagenet数据集中图像的尺寸大小并非统一,故本发明重新调整imagenet中训练数据集和测试数据集中的图像尺寸大小为256*256;
第二,生成各种类型的二值掩码(mask);
第三,对重新调整大小后的训练数据集增广,以增加训练数据集中的图像数量,包括随机水平翻转,随机颜色变换;
第四,将二值掩码乘上训练数据集中重新调整大小后的图像上生成受损图像,受损图像和尺寸调整后的图像(真实图像)作为模型的输入;将二值掩码乘上测试数据集中重新调整大小后的图像上生成测试图像。使用该测试图像,测试其模型的泛化性能。可视化结果如图1所示。该实验结果有效的证明了本发明所提出的方法对图像修复的有效性。
以上所述的具体实施例,对本发明的目的、技术方案和有益效果进行了进一步详细说明,应理解的是,以上所述仅为本发明的具体实施例而已,并不用于限制本发明,凡在本发明的精神和原则之内,所做的任何修改、等同替换、改进等,均应包含在本发明的保护范围之内。