CN109712203A - 一种基于自注意力生成对抗网络的图像着色方法 - Google Patents

Abstract

本发明公开一种基于自注意力生成对抗网络的图像着色方法，其包括以下步骤：步骤1，训练灰度图片着色模型；步骤2，将训练数据集中的灰度图像输入对抗网络执行特征提取阶段、特征融合阶段、反卷积计算阶段以及自注意力学习阶段重建出其对应的彩色图像；步骤3，将自注意力学习后重建的彩色图像与对应的原始色彩图像进行比较，计算惩罚函数；步骤4，基于损失函数作为GAN的优化损失；步骤5，将训练过程分为预设的多个子训练周期，采用步进增长策略依次进行子训练周期的训练以获取生成器网络。本发明采用对抗生成网络来从一张黑白或者灰度图像中重构出符合人类主观视觉偏好的彩色图像，使得更加逼真。

Description

一种基于自注意力生成对抗网络的图像着色方法

技术领域

本发明涉及图像着色及增强领域，尤其涉及一种基于自注意力生成对抗网络的图像着色方法。

背景技术

图像着色是一种图像增强的基本手段，旨在没有任何色彩提示的情况下为灰度图补充色彩信息，得到更佳完善的观感和视觉体验。随着时代的发展，彩色图像及视频对于普通消费者已经是稀松平常的体验，相对于早期只有黑白或者灰度信息的图片和影像资料显得丰富多彩。然而，老旧的黑白或者灰度资料由于缺少颜色信息，极难恢复成为体验更佳的彩色资料供现代观众观看。再者，随着硬件技术的飞跃，为了获得更好的视觉体验，人们对于4K乃至8K的图片和影像资料的需求越来越大。因此，为了使灰度图片及视频资料重焕生机，以及提高用户体验的情况下，我们亟需一种灰度图像和视频的人工智能着色方法。

在许久以前，图形学专家和计算机科学家就不断地在考虑图像着色问题。图像着色问题基本分成两类，一类是带有用户提示的辅助着色方法，第二类是无需任何着色提示的智能着色方法。基于提示的辅助着色方法需要人类的监督来完成上色，其也有两个分支：颜色涂抹法和色彩迁移法。Levin^【1】提出的基于颜色涂抹的方法十分高效，也得到了广泛应用。该模型以凸优化的方法，利用用户在局部区域涂抹的颜色色调对灰度图片进行上色，而实际上并没有学习到对于特定物件的上色方法。Yatziv^【2】在Levin的基础上提出了色度混合，将应用扩展到视频的上色。Qu^【3】在漫画上色领域另有建树，该方法提升了在相同纹理下的颜色连贯性和一致性。在色彩迁移法上，Welsh^【4】提出了基于另一张彩色图片作为参考，为灰度图像上色的方法。该模型能够为灰度图和参考图之间的信息产生一种色彩映射。

另一方面，智能着色方法更加吸引科学家的投入。由于极大地减少了图像上色的先验知识，自动上色的工作尤为困难，但其优势在于减去了大量的人力成本和上色时间，从而达到更为高效和便捷地实现图像着色问题。传统的智能着色方法采用了HoG^【5】，DAISY^【6】或色彩直方图^【7】来作为特征产生彩色图片。然而，随着深度学习的引入和大数据的加成，卷积神经网络已经在计算机视觉方面显示出其极大的层级特征提取能力，也逐渐取代了传统的特征工程。目前，基于卷积神经网络可以只用单纯的灰度图像作为端到端的模型输入来重构出其对应的彩色图像。然而，较为普遍的方法是使用逐像素的L1或者L2范数来计算重构的彩色图像与其对应的真实彩色图像间的差异，难免会产生色彩模糊及色彩溢出的问题。Cheng^【8】提出了利用小神经网络模型来提取不同特征及不同像素块进行上色的方法，同时也已用了语义直方图对图片进行区块处理。

随着Goodfellow提出的生成对抗网络^【5】(Generative adversarial networks:GAN)的出现，使用生成器和判别器之间的对抗学习，minimax损失函数也与之前提到的L2损失极为不同。近年来，许多基于GAN的方法如春笋般提出，其中包括最基础的DCGAN^【9】，条件GAN ^【10】，以及Pix2Pix^【11】较为复杂的CycleGAN^【12】等。虽然基于GAN的相关研究已经取得了较好的效果，但其方法基本都应用于学习待转换的图像与参考图像之间映射关系，达到生成的图像更接近真实图像，使得更加逼真。

发明内容

本发明的目的在于提供一种基于自注意力生成对抗网络的图像着色方法,利用自注意力生成对抗网络的技术将把黑白或者灰度图片重构出其对应的彩色图像的方法，目的在于提高老旧图片和影视资料的观感体验。本发明显著地提高了图像着色效果。。

本发明采用的技术方案是：

一种基于自注意力生成对抗网络的图像着色方法，其包括如下步骤：

步骤1：为了训练灰度图片着色模型，选取大量的彩色图像，在每张彩色原始图像进行旋转,镜像，和裁剪操作，得到大量的增广图片数据，并对每张彩色原始图像转其对应的灰度图像。为了提高彩色图像重构效果将把灰色图像堆叠成三通道的图像。之后将把灰度图像I_G和彩色原始图像I_C分别匹配作为深度学习的训练数据集，图片初始大小均为3×512×512，分别对应颜色通道，图片宽度和高度。

步骤2：将训练数据集中的灰度图像输入生成网络执行特征提取阶段、特征融合阶段、反卷积计算阶段以及自注意力学习阶段，最终重建出其对应的彩色图像。所述步骤2具体包括以下几个步骤：

步骤2.1：特征提取阶段

特征提取阶段由六大块卷积层组成，每个卷积块内又包括三个相同维度输入输出的卷积层。输入的灰度图像进入特征提取层后，进行卷积运算、谱归一化、批归一化和激活函数运算得到特征提取层输出结果，其中一个大卷积块的计算公式是：

F＝g{BN{SN{W₃*g{BN{SN{W₂*g{BN{SN{W₁*I_G+B₁}}}+B₂}}}+B₃}}} (1)

其中g表示非线性激活函数，W₁,W₂,W₃,B₁，B₂,B₃分别表示特征提取阶段中第一个、第二个和第三个卷积层的权重和偏置，SN表示谱归一化函数，BN表示批归一化函数，I_g表示输入图片，F表示特征提取阶段得到的输出结果。

步骤2.2特征融合阶段：

特征融合阶段由一个卷积层和一个跳跃式连接构成。通过跳跃式连接将该层的输出与特征提取层的输出相连接，避免梯度弥散并且通过保留原始特征以增强信息。表示公式如下：

F_g＝g(BN(W_LF_L+B_L))+F₂ (2)

其中F_L表示特征提取卷积模块的输出，W_L和B_L表示特征融合阶段卷积层的权重和偏置参数,F₂是前一阶段的特征提取层的输出，g表示非线性激活函数。

步骤2.3反卷积计算阶段：

本发明的反卷积计算阶段由一个或者多个反卷积层组成，反卷积层将特征融合阶段所得特征图进行上采样操作，根据放大倍数选择反卷积的数量，同时减少特征层的通道数减小计算量。其计算公式是：

I′_O＝((W_d1*F_g+B_d1)*W_d2+B_d2)…*W_dn+B_dn (3)

其中W_d1，W_d2，W_dn,B_d1,B_d2,和B_dn分别表示第一个、第二个和第n个反卷积层的权重和偏置参数，F_g是步骤2.2的特征融合层的输出，I′_O表示反卷积阶段的输出。

步骤2.4自注意力学习阶段：

本发明的自注意力学习阶段由多个卷积层组成。自注意力引入到卷积GAN中是对卷积的有机补充，有助于模拟跨越图像区域的长距离、多级别依赖关系。先前提出的基于卷积网络的GAN由于卷积层容易生成大量相同纹理而产生的失真效果明显，而自注意力网络能对每个位置的细节都与远处的细节进行准确协调。其计算公式是：

其中i为输出位置，j为参与运算的位置，Cx)为归一化参量，f(x_i,x_j)表示向量x_i与x_j处的相似关系函数，g(x_j)为x_j处的特征。

步骤3：利用步骤2.4自注意力学习后重建的上色图像与步骤1中所对应的原始色彩图像进行比较，计算Perceptual Loss惩罚函数。该损失函数可以表示为：

其中，j表示网络的第j层，C_jH_jW_j表示第j层特征图的大小。损失网络使用在ImageNet上训练好的VGG16网络，用φ表示网络。

步骤4，基于HingeLoss损失函数作为GAN的优化损失，损失函数为：

其中，L_D表示判别器的损失函数，L_G表示生成器的损失函数。损失函数的计算都可以看作是在数据P分布情况下的期望；

步骤5，将训练过程分为预设的多个子训练周期，采用步进增长策略依次进行子训练周期的训练；训练开始时将彩色原始图像缩放成小图片并佐以大学习率开始训练，且每完成一个子训练周期后逐步增大彩色原始图像并逐步减小学习率；

当完成一个子训练周期后重建的彩色图像与其对应的原始色彩图像比较时没有获得预先设定的重建效果，则继续进行反向传播，利用梯度下降优化算法更新卷积权值参数和偏置参数再执行步骤2；当在完成一个子训练周期后重建的彩色图像达到预期或完成所有预设的子训练周期的次数后，则获得最终结果。

具体地，与以往单个重复单个训练周期不断训练GAN的生成器和辨别器的训练方式不同，采用步进增长策略。在原始图片缩放成小图片的基础上开始训练，并佐以大学习率。在训练周期结束后增大输入图片，再减小学习率在此进行训练。以此类推，这样可以让分辨率更高的图片在低分辨率图片的基础上增强精度，减少因卷积网络生成导致的失真和不合理的颜色效果。

本发明采用以上技术方案，采用对抗生成网络来从一张黑白或者灰度图像中重构出其对应的较为符合人类主观视觉偏好的彩色图像，使得更加逼真。本发明结合了自注意力技术与渐进增大彩色图像分辨率的思想，从极小彩色图像到极大彩色图像的分辨率渐进增大，重构出更好的着色效果。

附图说明

以下结合附图和具体实施方式对本发明做进一步详细说明；

图1是本发明基于自注意力生成对抗网络的图像着色方法的网络结构；

图2是本发明所提出网络结构中自注意力学习模块；

图3是本发明与PaintsChainer,LTBC^[1]和RTIC^[2]在公开数据集Konachan上测试的效果对比图；

图4是本发明与PaintsChainer,LTBC^[1]和RTIC^[2]在公开数据集Danbooru上测试的效果对比图；

图5是本发明与PaintsChainer,LTBC^[1]和RTIC^[2]在公开数据集ImageNet上测试的效果对比图。

具体实施方式

如图1-5之一所示，本发明公开了一种基于自注意力生成对抗网络的图像着色方法，其包括如下步骤：

步骤1：为了训练灰度图片生成模型，选取Konachan高清动漫图片数据集，对原2K或4K分辨率的图片数据进行随机截取成为彩色原始图像，再对每张彩色原始图像进行旋转,镜像操作后，并经过RGB转灰度图操作，获得其对应的灰度图像。然后将灰度图像I_C和彩色原图像I_C分别切割成1×512×512和3×512×512的子图像，并进行归一化处理将图像像素值映射到[-1,1区间，得到训练数据集。

步骤2：将训练数据集中的灰度图像扩充到三维，和所期望的彩色输出图像维度一致，输入网络分别执行特征提取和数据生成。训练采取不同的输入图像大小、批量大小(为了适应GPU显存)、学习率等，将训练数据集中的灰度图像输入生成网络执行特征提取阶段、特征融合阶段、反卷积计算阶段以及自注意力学习阶段，最终重建出其对应的彩色图像。所述步骤2具体包括以下几个步骤：

步骤2.1特征提取阶段：

特征提取阶段由六大块卷积层组成，除了第一个卷积层的卷积核为7×7大小，其他卷积核均为3×3。每个卷积块内又包括三个相同维度输入输出的输入的灰度图像进入特征提取层后，进行卷积运算、谱归一化、批归一化和激活函数运算得到特征提取层输出结果。第一个卷积层通道数设置为64，之后逐块递增至512。其中一个大卷积块的计算公式是：

F＝g{BN{SN{W₃*g{BN{SN{W₂*g{BN{SN{W₁*I_G+B₁}}}+B₂}}}+B₃}}} (1)

其中g表示非线性激活函数，在本网络中指代ReLU激活函数，W₁,W₂,W₃,B₁，B₂,B₃分别表示特征提取阶段中第一个、第二个和第三个卷积层的权重和偏置，SN表示谱归一化函数，BN表示批归一化函数，I_g表示输入图片，F表示特征提取阶段得到的输出结果。使用谱归一化方法来限制GAN中分辨器D的优化过程，从而达到整个模型能学习到更好的生成器G的结果。

步骤2.2特征融合阶段：

特征融合阶段由一个卷积层和一个跳跃式连接构成。通过跳跃式连接将该层的输出与特征提取层的输出相连接，避免梯度弥散并且通过保留原始特征以增强信息。其中卷积层的卷积核大小设置为3×3，滤波器个数为20，卷积层输出接批归一化以及非线性激活函数ReLU。表示公式如下：

F_g＝g(BN(W_LF_L+B_L))+F (2)

其中F_L表示特征提取卷积模块的输出，W_L和B_L表示特征融合阶段卷积层的权重和偏置参数,BN表示批归一化层，F是特征提取层的输出，g表示非线性激活函数。

步骤2.3反卷积计算阶段：

本发明的反卷积计算阶段由一个或者多个反卷积层组成，反卷积层将特征融合阶段所得特征图进行上采样操作，根据放大倍数选择反卷积的数量，其计算公式是：

I′_O＝((W_d1*F_g+B_d1)*W_d2+B_d2)…*W_dn+B_dn (3)

其中W_d1，W_d2，W_dn,B_d1,B_d2,和B_dn分别表示第一个、第二个和第n个反卷积层的权重和偏置参数，F_g是步骤2.2的特征融合层的输出，I′_O表示反卷积阶段的输出。反卷积层的滤波器大小均为4×4，步长为2。在实际操作中，特征融合和反卷积阶段同时进行，先有特征融合得到新的特征图再进行反卷积。

步骤2.4自注意力学习阶段：

本发明的自注意力学习阶段由多个卷积层组成。自注意力引入到卷积GAN中是对卷积的有机补充，有助于模拟跨越图像区域的长距离，多级别依赖关系。先前的基于卷积网络的GAN由于卷积层容易生成大量相同纹理而产生的失真效果明显，而自注意力网络能对每个位置的细节都与远处的细节进行准确协调。其计算公式是：

其中i为输出位置，j为参与运算的位置，C(x)为归一化参量，f(x_i,x_j)表示向量x_i与x_j处的相似关系函数，g(x_j)为x_j处的特征。在自注意力学习曾中运用到的卷积输入和输出通道分别为64和3，卷积核大小均为1×1。为了实现这个过程，可以将矩阵乘法运算加上卷积运算。

步骤3：利用步骤2.4自注意力学习后重建的上色图像与步骤1中所对应的原始色彩图像进行比较，计算Perceptual Loss惩罚函数。该损失函数可以表示为：

其中，j表示网络的第j层，C_jH_jW_j表示第j层特征图的大小。损失网络使用在ImageNet上训练好的VGG16网络，用φ表示网络。

步骤4，基于HingeLoss损失函数作为GAN的优化损失，损失函数为：

其中，L_D表示判别器的损失函数，L_G表示生成器的损失函数。损失函数的计算都可以看作是在数据P分布情况下的期望；

步骤5，训练采用步进增长策略，将训练过程分为多个子训练过程。与以往单个重复单个训练周期不断训练GAN的生成器和辨别器的训练方式不同，采用步进增长策略。先将输入图片缩放至64×64，在初始学习率为1e^-3的情况下进行第一次迭代，当达到应有的迭代次数之后，再增大图片大小至96×96，并将学习率减半进行进一步拟合。以此类推，直至输入图片为原始图片大小。采用步进增长策略可以让分辨率更高的图片在低分辨率图片的基础上增强精度，减少因卷积网络生成导致的失真和不合理的颜色效果。

当重建后的图像与其对应的高分辨率图像进行比较时没有获得预先设定的重建效果，则继续进行反向传播，利用梯度下降优化算法更新卷积权值参数和偏置参数，再执行步骤2；在达到预期或完成步骤5中所有的子训练之后，求得步骤2所获得的生成对抗网络的生成器网络则为最终结果。

本发明采用峰值信噪比(PSNR:Peak Signal to Noise Ratio)和结构相似性(SSIM:Structural Similarity Index)来衡量图像上色性能。PSNR越高代表图像客观质量越好，SSIM值反映了图片的主观质量，该值越高表示图片还原度越高，视觉效果越好。比较结果如表一所示。

表一本发明与现有技术对于不同数据集的PSNR和MS-SSIM平均值

从表1中可看出，本发明比现有技术中的PSNR值和SSIM值在Konachan，Danbooru和ImageNet中分别提高了1.71dB、0.97dB、0.88dB和0.2883、0.1573、0.01964。由此可见，本发明采用以上技术方案，较这些方法的重建效果更好。

参考文献：

1.A.Levin,D.Lischinski,and Y.Weiss.Colorization using optimization.InACM Transactions on Graphics(ToG),volume 23,pages 689–694.ACM,2004

2.L.Yatziv and G.Sapiro.Fast image and video col-orizationusingchrominance blending.IEEE Transac-tions on Image Processing,15(5):1120–1129,2006.

3.Y.Qu,T.-T.Wong,and P.-A.Heng.Manga coloriza-tion.In ACMTransactionson Graphics(TOG),vol-ume 25,pages 1214–1220.ACM,2006.

4.T.Welsh,M.Ashikhmin,andK.Mueller.Transferring color to greyscaleimages.In ACM Transactions on Graphics(TOG),volume 21,pages 277–280.ACM,2002.

5.N.Dalal and B.Triggs.Histograms of oriented gra-dients for humandetection.In Computer Vision and Pattern Recognition,2005.CVPR 2005.IEEE Com-puter Society Conference on,volume 1,pages 886–893.IEEE,2005.

6.E.Tola,V.Lepetit,and P.Fua.Daisy:An efficient dense descriptorapplied to wide-baseline stereo.IEEE transactions on pattern analysis andmachine intelli-gence,32(5):815–830,2010.

7.J.Hafner,H.S.Sawhney,W.Equitz,M.Flickner,and W.Niblack.Efficientcolor histogram indexing for quadratic form distance functions.IEEE trans-actions on pattern analysis and machine intelligence,17(7):729–736,1995.

8.CHENG,Z.,YANG,Q.,AND SHENG,B.2015.Deep colorization.In Proceedingsof ICCV 2015,29–43.

9.I.Goodfellow,J.Pouget-Abadie,M.Mirza,B.Xu,D.Warde-Farley,S.Ozair,A.Courville,and Y.Ben-gio.Generative adversarial nets.In Advances in neu-ralinformation processing systems,pages 2672–2680,2014.

10.A.Radford,L.Metz,and S.Chintala.Unsu-pervised representationlearning with deep convolu-tional generative adversarial networks.arXivpreprint arXiv:1511.06434,2015.

11.M.Mirza and S.Osindero.Conditional generative ad-versarialnets.2014.

12.P.Isola,J.-Y.Zhu,T.Zhou,and A.A.Efros.Image-to-image translationwith conditional adversarial net-works.arXiv preprint arXiv:1611.07004,2016.

13.Zhu JY,Park T,Isola P,Efros AA.Unpaired image-to-image translationusing cycle-consistent adversarial networks.arXiv preprint.2017.

14.Iizuka,Satoshi,Edgar Simo-Serra,and Hiroshi Ishikawa."Let there becolor！:joint end-to-end learning of global and local image priors forautomatic image colorization with simultaneous classification."ACMTransactions on Graphics (TOG)35.4(2016):110.

15.Zhang,Richard,et al."Real-time user-guided image colorization withlearned deep priors."arXiv preprint arXiv:1705.02999(2017)。

Claims (6)

1.一种基于自注意力生成对抗网络的图像着色方法，其特征在于：其包括如下步骤：

步骤1：训练灰度图片着色模型：选取大量的彩色原始图像处理得到大量的增广图片数据，并对每张彩色原始图像转其对应的灰度图像，将灰色图像堆叠成三通道的图像，将把灰度图像I_G和彩色原始图像I_C分别匹配作为深度学习的训练数据集，

步骤2：将训练数据集中的灰度图像作为输入并执行特征提取阶段、特征融合阶段、反卷积计算阶段以及自注意力学习阶段，最终重建出其对应的彩色图像，具体包括以下几个步骤：

步骤2.1，特征提取阶段：输入的灰度图像进入特征提取阶段后，进行卷积运算、谱归一化、批归一化和激活函数运算得到特征提取层输出结果；特征提取阶段由六大块卷积层组成，每个卷积块内包括三个相同维度输入输出的卷积层；其中一个大卷积块的计算公式是：

F＝g{BN{SN{W₃*g{BN{SN{W₂*g{BN{SN{W₁*I_G+B₁}}}+B₂}}}+B₃}}} (1)

其中g表示非线性激活函数，W₁，W₂，W₃，B₁，B₂，B₃分别表示特征提取阶段中第一个、第二个和第三个卷积层的权重和偏置，SN表示谱归一化函数，BN表示批归一化函数，I_g表示输入图片，F表示特征提取阶段得到的输出结果。

步骤2.2，特征融合阶段：

特征融合阶段由一个卷积层和一个跳跃式连接构成，通过跳跃式连接将该卷积层的输出与特征提取层的输出相连接，特征融合阶段的公式如下：

F_g＝g(BN(W_LF_L+B_L))+F₂ (2)

其中F_L表示特征提取卷积模块的输出，W_L和B_L表示特征融合阶段卷积层的权重和偏置参数，F₂是前一阶段的特征提取层的输出，g表示非线性激活函数。

步骤2.3，反卷积计算阶段：反卷积计算阶段由一个以上反卷积层组成，反卷积层将特征融合阶段所得特征图进行上采样操作，其计算公式是：

I′_O＝((W_d1*F_g+B_d1)*W_d2+B_d2)…*W_dn+B_dn (3)

其中W_d1，W_d2，W_dn，B_d1，B_d2，和B_dn分别表示第一个、第二个和第n个反卷积层的权重和偏置参数，F_g是特征融合阶段的输出，I′_O表示反卷积阶段的输出。

步骤2.4，自注意力学习阶段：自注意力学习阶段由多个卷积层组成，其计算公式是：

其中i为输出位置，j为参与运算的位置，C(x)为归一化参量，f(x_i，x_j)表示向量x_i与x_j处的相似关系函数，g(x_j)为x_j处的特征；

步骤3：将自注意力学习后重建的彩色图像与对应的原始色彩图像进行比较，计算Perceptual Loss惩罚函数，该损失函数可以表示为：

其中，j表示网络的第j层，C_jH_jW_j表示第j层特征图的大小，φ表示损失网络；

步骤4，基于HingeLoss损失函数作为GAN的优化损失，损失函数为：

其中，L_D表示判别器的损失函数，L_G表示生成器的损失函数。损失函数的计算都可以看作是在数据P分布情况下的期望；

步骤5，将训练过程分为预设的多个子训练周期，采用步进增长策略依次进行子训练周期的训练；

当完成一个子训练周期后重建的彩色图像与其对应的原始色彩图像比较时没有获得预先设定的重建效果，则继续进行反向传播，利用梯度下降优化算法更新卷积权值参数和偏置参数再执行步骤2；当在完成一个子训练周期后重建的彩色图像达到预期或完成所有预设的子训练周期的次数后，则获得最终结果。

2.根据权利要求1所述的一种基于自注意力生成对抗网络的图像着色方法，其特征在于：步骤1中每张所述彩色原始图像进行旋转、镜像和裁剪操作得到大量的增广图片数据。

3.根据权利要求1所述的一种基于自注意力生成对抗网络的图像着色方法，其特征在于：步骤1中增广后的训练图片初始大小均为(3×512×512)，分别对应颜色通道，图片宽度和高度。

4.根据权利要求1所述的一种基于自注意力生成对抗网络的图像着色方法，其特征在于：步骤2.3中反卷积计算阶段中根据放大倍数选择反卷积的数量。

5.根据权利要求1所述的一种基于自注意力生成对抗网络的图像着色方法，其特征在于：步骤3中损失网络使用在ImageNet上训练好的VGG16网络。

6.根据权利要求1所述的一种基于自注意力生成对抗网络的图像着色方法，其特征在于：步骤5中步进策略具体为：训练开始时将彩色原始图像缩放成小图片并佐以大学习率开始训练，且每完成一个子训练周期后逐步增大彩色原始图像并逐步减小学习率。