基于CycleGAN和Pix2Pix的轨面缺陷图像智能生成技术

陈嘉欣，孙传猛，葛耀栋，李欣宇，靳书云，李 勇

(1.中北大学 省部共建动态测试技术国家重点实验室，山西 太原 030051；2.中北大学 电气与控制工程学院，山西 太原 030051；3.华北计算机系统工程研究所，北京 100083；4.重庆大学 煤矿灾害动力学与控制国家重点实验室，重庆 400044)

在国家“科技强国”“交通强国”战略支撑下,高速铁路等新型载运工具得到了普遍关注和认可,并积极进行布局。其中,高动态、高精度轨道巡检技术对高速铁路的智能运维至关重要[1]。然而,基于人工智能技术开展高速条件下钢轨轨面缺陷自主检测面临以下小样本学习难题:

1)数据分布不平衡[2]。已有轨道在通常情况下都是无缺陷的,对刚修建完或正在建设中的高速铁路轨道而言,获得轨面缺陷数据更是难上加难。显然,目前构建的轨面缺陷数据集存在正负样本分布不平衡问题,这严重降低了相关缺陷检测模型的性能,极端情况下模型会将缺陷识别为噪声从而导致缺陷检测结果无效。

2)数据缺乏多样性。轨面缺陷数据获取困难,相关数据集规模偏小,这意味着数据正负样本的形态偏少,背景与目标样式多样性不足,导致模型只能学到部分特征而影响模型的识别性能,甚至引发“过拟合”问题。

3)数据标注难度大[3]。轨面缺陷数据的标注存在鉴定难等问题,需要专业人员完成。此外,通过实地拍摄得到覆盖面广的数据集并准确标注,将会耗费大量的人力、物力。

针对上述小样本学习难题,一种可靠的解决方法是迁移学习,它通过人工生成缺陷形态符合真实分布的、目标与背景近似真实的模拟图像来使数据集的增强,训练增强数据集实现缺陷检测模型的卓越性能;随着新的、真实的缺陷数据不断加入训练集,不断微调已有模型,实现性能的最优化。

显然,模拟图像生成的质量直接关系到迁移学习的成败。生成对抗网络(generative adversarial network, GAN)[4-8]是一种基于博弈论的深度学习生成模型,可以学习样本数据高阶语义特征并生成具有该特征的新数据。GAN模型是在2014年由Goodfellow等提出,相对其他数据增强方法,其生成的图像具有清晰度、分辨率更高且速度更快等优点[9-10]。在图像生成领域,相关学者对GAN模型做了大量研究工作。2019年,Bailo等[11]采用条件生成对抗网络(conditional generative adversarial network, CGAN),对显微镜下的红细胞图像进行生成,得出在生成数据对检测任务的增益大于分割任务的结论。2020年,Suh等[12]提出分类增强型生成对抗网络(classification enhancement generative adversarial networks, CEGAN),可以在数据不平衡的情况下生成合成的少数群体数据来增强分类,提高了非平衡数据条件下分类的稳定性和性能。同年,Niu等[13]提出一种新的缺陷图像生成方法,引入两个鉴别器和周期一致性损失来生成换向器缺损图像。2021年,Belhiti等[14]改进了StackGAN,实现对人脸数据的生成。同年,Cui等[15]提出生成式协同训练网络GenCo,通过引入多个互补的鉴别器来缓解鉴别器过拟合的问题,这些鉴别器在训练中可以在多个不同的尺度上进行检测。

虽然轨面缺陷形态无定式,但类别是固定的。因此,生成的轨面缺陷模拟图像,既不能样式单一、与小样本数据集中缺陷形态完全一致,又不能形态过分夸张、跳脱已有类别限制。因此,本文采用CycleGAN和Pix2Pix相结合的方法,实现小样本数据集语义特征学习和轨面缺陷数据的智能生成。其中,Pix2Pix输入的是成对的缺陷图像和背景图像,生成特定类别的缺陷图像;CycleGAN模型则生成不受定式约束的缺陷图像。二者结合,实现模拟图像缺陷类别不变而形态各异,从而达到数据集增强的目的。

1 相关理论

1.1 GAN

GAN[4]结构见图1。

图1 GAN模型基本结构

通过生成器网络G和判别器网络D不断进行二元极大、极小博弈,达到同时优化两个模型的目的。其中,生成器G的目标是提高生成样本与真实样本的相似度,即实现生成样本pg与真实样本pd分布相同;判别器D的目标是分辨真实样本与生成样本。GAN损失函数(目标函数)LGAN的表达式为

E(ln{1-fD[fG(z)]})}

(1)

式中:x为图像数据;z为随机噪声;E为期望值;fD为判别器D的输出;fG为生成器G的输出。

1.2 Pix2Pix

经典GAN模型生成内容由参数fG和随机噪声z决定,无法控制生成内容。为控制生成器生成内容,引入条件信息y。在Pix2Pix模型中,y为与m对应的另一个图像域数据。在生成器中,随机噪声z和数据y同时作为输入,生成一个跨模态特征;在判别器中,数据m及对应的y作为输入,生成跨模态向量,同时判断m的真伪。这样,随着条件信息y的引入,生成器逐渐被控,生成特定结果。加入条件信息y,Pix2Pix的损失函数[16-17]Lpix2pix的表达式为

E(ln{1-fD[fG(z|y)]})}

(2)

通常,在损失函数中引入L1正则项可以使模型参数稀疏化,从而提升泛化能力,有助于生成更清晰的数据[18],则L1正则项损失函数LL1为

LL1(fG)=E[‖y-fG(m,z)‖1]

(3)

在正则约束情况下,添加惩罚因子λ为0.01的惩罚项,则模型参数fG*为

(4)

不同于经典GAN使用MLP作为模型结构,Pix2Pix使用了CNN中常用的卷积+BN+ReLU的模型结构。其中,生成器采用了经典的编码器-解码器结构[19],即UNet,其典型结构见图2。UNet可以保留不同分辨率下像素级的细节信息,实现将高分辨率输入网格映射到高分辨率输出网格;判别器采用PatchGAN[20],这不同于经典GAN,将整幅图作为判别器判别的目标,而是将输入图像分成若干个图像块(Patch),依次作为判别目标。当图像块的大小为1×1时,PatchGAN又被叫做PixelGAN。Pix2Pix对一张图片切割成不同的N×N大小的Patch,判别器对每一个Patch做真假判别,将一张图片所有Patch的结果取平均作为最终的判别器输出。PatchGAN判别器结构见图3。

图2 UNet典型结构

图3 PatchGAN判别器结构

最终,Pix2Pix形成如图4所示的结构。图4中,黑白Mask标注即为引入的条件信息y,该标注为二值图像,反映了缺陷所对应的位置和形状大小信息。

图4 Pix2Pix模型基本结构

1.3 CycleGAN

Pix2Pix模型训练需要输入成对的缺陷图像和背景图像。显然,实际成对的轨面图像数据是很难获取的。CycleGAN[21],即循环生成对抗网络,采用双向循环生成的结构,突破Pix2Pix对数据集图片一一对应的限制。

CycleGAN结构见图5。图5中包含两个生成器网络和两个判别器网络。生成器网络结构见图6。其中,两个生成器网络结构相同,包含下采样、转码和上采样3部分;两个判别器结构采用PatchGAN,同Pix2Pix模型。

图5 CycleGAN模型基本结构

图6 生成器网络结构

设无缺陷轨面图像为X域,有缺陷轨面图像为Y域,则生成器网络G1的映射函数fG1为

fG1(x)=y′x∈X

(5)

生成器网络G2的映射函数fG2为

fG2(y)=x′y∈Y

(6)

两个生成器可以实现跨域图像生成,两个判别器则分别判断两个生成器生成数据与真实数据之间的差异,这与经典GAN网络是一致的。

由于输入CycleGAN的两组(X、Y域)图像没有一一对应关系,即输入的两组图片完全是随机搭配在一起,常规GAN以及Pix2Pix学不到任何联系,而CycleGAN通过添加“循环生成”并优化循环一致性损失(consistencyloss)来代替Pix2Pix中的约束条件,使生成器网络保留原域图片特征。循环一致性体现在:由生成器网络G1生成的Y域数据,再由生成器网络G2恢复为X域数据x′,这与初始输入的X域x应该是一致的。因此,可由L1损失来判断x与x′的差异,生成器的循环一致性损失Lcyc的表达式为

Lcyc(fG1,fG2)=E{‖fG2[fG1(x)]-x‖1}+

E{‖fG1[fG2(y)]-y‖1}

(7)

同经典GAN一样,CycleGAN还包含对抗性损失。其中,由X域到Y域的对抗损失为

LGAN(fG1,fD2,X,Y)=E[lnfD2(y)]+

E(ln{1-fD2[fG1(x)]})

(8)

式中: fD2为判别器网络D2对应的映射函数。

由Y域到X域的对抗损失为

LGAN(fG2,fD1,Y,X)=E[lnfD1(x)]+

E(ln{1-fD1[fG2(y)]})

(9)

式中: fD1为判别器网络D1对应的映射函数。

此外,CycleGAN还引入了IdentityLoss来确保生成器网络所生成的数据尽量保留源图像特性(如色调)。IdentityLoss的表达式为

LIdent(fG1,fG2)=E[‖fG1(x)-x‖1]+E[‖fG2(y)-y‖1]

(10)

这样,CycleGAN总的损失函数表达式为

λ2LIdent(fG1,fG2)

(11)

式中:λ1、λ2为正则化系数。

1.4 智能生成总体流程

综上所述,采样Pix2Pix模型可以约束生成缺陷图像的类型,采用CycleGAN模型可以生成不受定式约束的轨面缺陷图像。这样,将二者结合即可实现模拟图像缺陷类别均衡而形态各异,达到数据集增强的目的,从而解决基于人工智能的高动态高精度轨道巡检技术的小样本学习难题,其具体过程见图7。

图7 总体流程

基于CycleGAN和Pix2Pix的轨面缺陷图像智能生成方法的伪代码如下:

初始化:初始化Mp、Mc。

1.for(t=1,…,T)do

2. 训练模型Mp中的生成器Gp、判别器Dp;

3. 损失函数反向传播,更新模型参数;

4.endfor

5.for(t=1,…,T)do

6. 更新模型Mc中的生成器G1、G2,判别器D1、D2;

7. 损失函数反向传播,更新模型参数;

8.endfor

9.设定生成黑白Mask标注类型,使用Gp得到缺陷图像。

10.输入轨面无缺陷图像和缺陷图像,使用G1、G2对应生成轨面缺陷图像和无缺陷图像。

2 实验研究

2.1 实验方案

为验证本文技术方法的可行性,设计实验方案包含以下两个方面:①制作轨面图像数据集。数据集包含有缺陷数据集和无缺陷数据集两大类,有缺陷数据的缺陷均有对应的黑白Mask标注。②分别搭建Pix2Pix和CycleGAN生成对抗网络,分别输入相应的数据训练各自模型,并考察其生成图像质量。

本次实验平台为Pycharm2020,所用GPUNVIDIARTX6000,开发语言为Python3.6,配置环境为Pytorch1.7。

2.2 轨面缺陷数据集制作

本文采用RSDDsType-Ⅰ数据集[22]作为实验对象。相比其他故障检测数据集,RSDDs检测难度更大,体现如下: ①识别功能有限。该数据集视觉特征的局限性使得大多数基于复杂纹理和形状特征的目标识别方法无法成功训练[23]。②缺乏可识别的先验信息。轨面缺陷往往以随机或任意的方式出现,其形状和大小各不相同。③较低的灰度值先验[24]往往被光照不均一、反射特性变化和噪声引起的不规则物所抑制。

RSDDsType-Ⅰ数据集共有67张真实缺陷轨面图片和其缺陷对应的黑白Mask标注。数据集部分图片及其黑白Mask标注见图8。RSDDsType-Ⅰ数据集如何判断缺陷说明见文献[22]。

图8 RSDDs Type-Ⅰ数据集部分样本

原始的RSDDsType-Ⅰ数据集不能满足本实验方案要求,必须再次加工获得满足本次实验的数据集样本。首先,利用矩阵运算将原始样本图像裁剪成256×256大小的多个新样本;然后,根据对应Mask标签将裁剪后的图片分为有缺陷轨面数据集和无缺陷轨面数据集;最终,制作的有缺陷轨面数据集包含样本141张及对应黑白Mask标注,部分样本见图9(a),无缺陷轨面数据集包含样本465张,部分样本见图9(b)。

图9 有、无缺陷轨面数据集部分样本

2.3 实验结果与分析

2.3.1 基于Pix2Pix的轨道图像生成实验

将成对的有缺陷轨面数据集样本与其黑白Mask标注图像输入到Pix2Pix模型进行训练,训练共计800轮。为防止过拟合,在Pix2Pix架构中还设置了Dropout层。

训练完成后,随机抽取若干样本测试网络性能,测试结果见图10。由图10可知,通过控制输入到Pix2Pix网络的黑白Mask形状,可以实现对生成缺陷形态的精确控制。

图10 Pix2Pix模型随机抽样测试结果

为进一步验证Pix2Pix的轨面图像生成效果,需要更多的Mask图像,达到丰富数据集的目的。Mask图像采用如下两种方法生成: ①对RSDDsType-I数据集自带的Mask进行随机变换,包括缩放、镜像、移位等操作; ②利用规则图形生成Mask图像。方法②的具体步骤如下:

Step1生成256×256大小的黑色背景。

Step2生成一个随机数n作为缺陷数量。

Step3随机生成一个坐标(x,y)和半径r,以此为参数画圆并填充为白色。

Step4确定横向压缩率dx、纵向压缩率dy、旋转角度θ,并对白色目标进行对应的压缩、旋转操作。

Step5不断重复Step3、Step4,直至完成所有的n个缺陷目标的生成,然后转入下一个Mask的制作。

Step6利用上述两种方法获得了理想的轨面缺陷Mask图像,部分样本见图11。其中,图11(a)由方法①得到;图11(b)由方法②得到。

图11 两种方法生成的部分黑白Mask图像

将Mask图像输入到训练好的Pix2Pix模型,生成有缺陷轨面模拟图像,结果见图12(因篇幅有限,仅展示部分图像)。由图12可知,随机生成的Mask图像均能得到对应的模拟缺陷图像,表明Pix2Pix模型能够生成特定类别的缺陷图像。

图12 设定Mask生成的模拟缺陷图像

2.3.2 基于CycleGAN的轨面图像生成实验

将465张无缺陷轨面数据集样本作为真实X域数据,将141张有缺陷轨面数据集样本作为真实Y域数据,输入CycleGAN进行训练。训练共计200轮。

训练完成后,随机确定轨面图像进行CycleGAN模型效果验证,结果见图13。由图13可知,CycleGAN模型成功地将有缺陷轨面图像转换成无缺陷轨面图像,将无缺陷轨面图像转换成有缺陷轨面图像,且缺陷样式不受定式约束。

图13 CycleGAN生成的图像

3 性能测试

本文方法生成的轨面缺陷模拟图像,从图像内容和缺陷形态而言,十分接近真实的轨面缺陷图像。但对于小样本数据集而言,存在数据分布不平衡、数据缺乏多样性等问题。由于生成器固有特性,无法利用划分数据集的方式来评估生成对抗网络模型生成图像的性能[25]。因此,采用本文方法生成轨面缺陷模拟图像进行数据集增强后,是否解决了数据分布平衡、样本多样性等问题,当前并没有恰当的评价指标。如果增强后的数据集在计算机视觉任务中能更好地发挥各自网络模型的性能,则能间接反映增强数据集解决了原数据集小样本学习难题。因此,本文利用图像分类、目标检测、语义分割这3大典型的计算机视觉任务来测试本文方法生成图像性能,进而验证该方法解决轨面缺陷检测技术小样本学习难题的有效性。在训练过程中,所有实验均同时采用了在线数据增强,最终实验结果仅突出了数据集扩充后的有效性。

3.1 基于VGG19的图像分类性能测试

VGG19是牛津大学视觉几何小组提出的一种深层卷积网络结构,在图像分类任务中性能卓越、广受认可[26]。VGG19训练不需要Mask标注,只需将有、无缺陷轨道图像分别赋予不同的分类标签输入网络即可。

首先,构造原始的轨面缺陷小样本数据集。为模拟数据分布失衡,用于训练的有、无缺陷轨面图像样本分别为20、311张;用于测试的有、无缺陷轨面图像样本分别为121、134张,见表1。表1中,def、no_def分别表示有、无缺陷轨面样本。

表1 原始数据集训练和测试划分 张

然后,利用Pix2Pix和CycleGAN对有缺陷轨面数据集进行增强,结果见表2。表2中,ori为原始数据集;aug(Pix2Pix2)、aug(CycleGAN)和aug(Mix)分别为Pix2Pix2、CycleGAN和两者混合增强的数据集。

表2 增强的有缺陷数据集结果 张

将增强后的有缺陷轨面数据均加入到训练集中,这样新的增强数据集中有缺陷和无缺陷样本数量相等。

最后,分别用原始数据训练集和增强数据训练集训练具有同一结构的VGG19模型,并利用测试集(两数据集的测试集完全相同)测试各自训练模型的性能,结果见表3。表3中,准确率为正确分类的样本数与总样本数之比。由表3可知,采用Pix2Pix和CycleGAN增强数据集均大幅度提高了VGG19分类模型的性能,而联合CycleGAN和Pix2Pix增强方法提升幅度最大,达81.177%。这表明本文方法在图像分类任务中,可显著改善小样本特征带来的数据分布不平衡、缺乏多样性等问题。

表3 基于VGG19的图像分类性能测试结果 %

3.2 基于YOLOv5的目标检测性能测试

YOLOv5是近年来性能最卓越的目标检测算法之一。本文利用YOLOv5来测试增强数据集对目标检测算法性能的提高能力。

使用与3.1节所述中相同的有缺陷轨面图像样本作为训练集和测试集,且在不同数据集的实验中均设置相同的超参数以对比检测结果。以准确率、召回率以及均值平均精度作为检测结果的评价指标,数值越高表明对缺陷的定位和识别越加准确。对3种增强数据集进行实验,结果见表4。表4中,召回率为正样本被预测正确数量与总正样本个数之比;均值平均精度为计算所有类别的准确率和召回率曲线下面积的平均值。由表4可知,联合Pix2Pix和CycleGAN增强数据集方法使得均值平均精度达到最高,相比原数据集提升了18.30%。这表明,联合两种方法丰富数据集对模型泛化能力的提升效果最为显著,部分检测结果见图14。

图14 原始与3种增强数据集的部分缺陷检测结果

3.3 基于UNet的语义分割性能测试

UNet是一种经典的语义分割深度学习模型。本文利用UNet来测试增强数据集对语义分割算法性能的提高能力。该测试试验所用的训练集和测试集同样为3.1节所述数据集。

由于UNet是有监督学习,需要为训练集的各个样本制作相应的标签图像。对于原始训练集,各样本对应的黑白Mask图像即可作为标签;对于Pix2Pix生成的有缺陷轨面数据,其输入的Mask图像可以直接作为标签;对于CycleGAN生成的有缺陷轨面数据,由于是在无缺陷数据上生成的,将生成数据和原数据进行减法运算,并采用全局阈值处理方法[27]转换为二值图像,然后使用形态学处理滤除噪声并填补空隙,即可得到相应的Mask标签图像,见图15。

图15 CycleGAN生成缺陷图像标签

分别用原始数据训练集和增强数据训练集训练具有同一结构的UNet模型,并利用测试集测试各自训练模型的性能。UNet模型的评估指标为Dice Score,该指标主要用于计算两集合间的相似性,即两集合间相似性越高,Dice Score越高。UNet模型在各数据集上的性能见表5。表5中,Dice Score为预测结果与真实标签的重叠区域与两者并集之比。由表5可知,采用Pix2Pix增强数据集降低了UNet分割性能,这是因为只用Pix2Pix生成数据没有有效的增强缺陷的多样性;只采用CycleGAN方法生成的数据,在多样性上有较大增强,因而提升了UNet分割性能;而联合CycleGAN和Pix2Pix增强方法在增强数据分布平衡和缺陷多样性方面更优,因而取得了最好的分割性能,其UNet分割Dice Score达到71.015,相比原始数据集提高了6%。

表5 基于UNet的图像分割性能测试结果

4 结论

1)采用CycleGAN和Pix2Pix相结合的方法,实现小样本数据集语义特征学习和轨面缺陷数据的智能生成,生成数据缺陷类别不变、而形态各异,从而达到数据集增强的目的。

2)利用VGG19、YOLOv5和UNet进行的轨道图像小样本数据集增强性能测试实验表明,本文方法在图像分类任务中的准确率为81.177%,较原数据集提高23.138%;在目标检测任务中,准确率为91.90%,提升26.60%,召回率为87.20%,提升16.00%,均值平均精度为93.50%,提升18.30%;在语义分割任务中Dice Score为71.015,较原数据集提高6%。