基于Swin-Transformer的磁瓦缺陷检测
2023-06-27陈荣演杨创富陈宇琪宁洪龙
陈荣演,邱 天*,杨创富,张 昕,陈宇琪,符 晓,宁洪龙
(1. 五邑大学智能制造学部,江门 529020;2. 阜阳师范大学物理与电子工程学院,阜阳 236037;3. 华南理工大学高分子光电材料与器件研究所,广州 523648)
0 引言
磁瓦因其良好的物理性质被广泛应用于各种电子产品的线圈、电动机以及变压器当中,然而在磁瓦生产加工过程中产生的缺陷不仅影响外观而且还会影响相关产品质量,因此磁瓦缺陷检测是不可或缺的工序。近些年来磁瓦检测技术依赖传统视觉检测方式,也取得一定的研究成果[1-5],但是传统视觉检测方式依赖人工设定提取特征,而且提取特征难度大、检测的鲁棒性差。随着深度学习在视觉领域的广泛应用,这种由数据驱动获取特征的方式使卷积神经网络在磁瓦缺陷检测中也有很多的研究成果[6-10]。尽管卷积神经网络在缺陷检测取得很好的检测效果,但是其缺乏全局信息,对长远距离的特征提取有局限性,而近年来提出的视觉Transformer[11]模型因其抛弃卷积网络结构,使用完全多头自注意力机制学习特征,因而拥有更大的感受野并对全局特征的表述更强,并且能够捕获更多丰富的特征,对小物体的检测也有更好的表现,从而表现出很强的竞争力,其中最具代表的为Liu等[12]提出的Swin-Transformer模型,其通过窗口滑动和分层策略使计算量显著减少并提高建模能力。由于注意力机制的表现能力较好,因此近些年来使用注意力模型也出现越来越多的研究成果[13-14]。
本文提出基于Swin-Transformer 的磁瓦检测算法,为了使网络学习更多有用的特征,首先对采集到的图像进行增强,让磁瓦缺陷图像的对比度更高和缺陷边缘更清晰,然后使用Swin-Transformer 作为特征提取网络提取特征,并结合YOLOv3 模型的neck 模块构建磁瓦缺陷检测算法,同时为了使网络获取更丰富的梯度组合,在neck部分进行优化设计CSPA(cross stage partical attention)模块,最终实现了较好的缺陷检测结果。
1 图像增强
由于磁瓦材料特性,采集到的磁瓦图像整体上都会呈现颜色灰暗以及缺陷对比度低等特点,为了加强特征的表述和使提取特征更加容易,因此对收集到的磁瓦图片进行增强操作,增强的步骤如下:
步骤1:对输入图像进行高斯滤波去除噪声,然后使用拉普拉斯算子对图像进行锐化处理,以加强边缘梯度,拉普拉斯锐化计算如公式(1):
其中:L(x,y)为锐化后图像,P(x,y)为原图像,∇2f为拉普拉斯滤波结果,其定义如公式(2):
步骤2:将采集的磁瓦图像先通过对数转换,再通过傅里叶变换转换到频域上,其中频域的高频对应原图的细节部分,图像的低频对应原图的近似部分。
步骤3:设计合适的滤波器传递函数,使频域的高频细节部分得到增强,并且压缩低频动态范围,本文选用的滤波器传递函数如公式(3),其二维表示图如图1所示。
步骤4:用设计好的滤波器对频域图片进行滤波。
步骤5:将滤波后图像变换回空间域得到增强后图片。
图1 滤波器传递函数二维显示
图像增强步骤流程如图2 所示,经过上述步骤后,原图和增强图像的直方图如图3 所示,从图3可以看出原图经过增强后缺陷对比度更加明显,其直方图也更加均衡。
图2 图像增强步骤流程
图3 原图和增强后图像对比
2 磁瓦缺陷检测网络原理及其改进
2.1 Swin-Transformer原理
Swin-Transformer 是Transformer 应用在视觉领域的一项重要成果,并且在目标检测、分割和分类等视觉任务中取得很好的成绩。Transformer舍去了传统卷积神经网络结构,利用多头自注意力机制提取特征,可以更好地捕获上下文信息,并削弱了网络对外部信息的依赖,因此Transformer 善于捕获特征或数据的内在相关性和提取表述性更强的特征。但是Transformer应用在视觉任务中受到两个重要的挑战:①视觉任务的多尺度问题;②视觉图像分辨率高,计算参数庞大。而Swin-Transformer 通过窗口滑动限制在局部窗口计算自注意力可以很大程度降低计算成本和复杂度;使用分层策略使网络可以提取多个尺度的特征,对小物体的识别也有很好的效果,因此Swin-Transformer 可以作为磁瓦目标检测的特征提取网络部分。
Swin-Transformer 结构包括编码模块(Patch Partition+Linear Embedding)、自注意力计算模块(Swin-Transfomer Block)和移动拼接模块(Patch Merging)。编码模块将输入数据进行分块编码和线性变换,然后输入到自注意力计算模块(如图4 所示)。自注意力计算模块由两个串联注意力计算模块组成,每个注意力计算模块都由一个多头自注意力模块和多层感知机(MLP)在进行层归一化(LN)后进行残差连接得到。在W-MSA模块中将输入特征图尺寸的8 × 8 划分为2 × 2个窗口,划分方式如图5所示,并对每个窗口计算多头自注意力;而SW-MSA 通过循环移位的方式将窗口进行移动,可以获取相邻窗口的交互注意力信息,然后对每个窗口计算多头自注意力,其划分方式如图6所示;多头注意力的计算见公式(4)和公式(5)。
其中:Q,K,V分别为查询、键、值向量;B是相对位置偏置;W是线性变换的权重。
在进行自注意力计算后,通过移动拼接模块合并相邻的分块实现不同尺度特征图的获取,从而构建特征的层级结构。在Swin-Transformer中通过移动拼接模块后特征图的下采样倍数分别为8、16、32,和YOLOv3 中的骨干网络(Backbone)采样倍数刚好一致,因此可以直接用Swin-Transformer作为特征提取的骨干网络结合YOLOv3的neck构成目标检测网络。
图4 自注意力计算模块(Swin-Transfomer Block)
图5 W-MSA
图6 SW-MSA
2.2 CSPA模块
为了使网络能够获取更加丰富的梯度组合从而增强网络学习能力,本文在CSPNet[15](cross stage partial network)基础上设计了CSPA(cross stage partical attention)模块用于提取特征。其中改进方法有:①在CSP 层嵌入CA 模块;②通过增加分支获取更多梯度流。其中嵌入CA模块[16]是为了使网络能够捕获跨通道的信息和方向感知,帮助网络更加精确定位到感兴趣的目标特征。CA 模块是通过水平和垂直方向生成相应注意力图来强调感兴趣的特征区域,其相应的方法流程如图7所示,其输入特征首先经垂直和水平方向进行一维全局池化并分别对这两个方向特征进行编码,然后每个特征图分别沿着一个方向去捕获特征的长距离相关性,最后将包含距离相关性的特征注意力图和输入特征相乘从而获得更受关注的特征信息。
图7 CA注意力模块
2.3 改进磁瓦缺陷检测算法
改进磁瓦缺陷检测网络如图8 所示,其由Swin-Transformer 作为特征提取部分,在经过移动拼接模块(patch merging)对输入图像进行下采样,把下采样8、16 和32 倍下采样的输出特征输入到YOLOv3 的neck 模块,并且使用设计的CSPA 模块对neck 部分进行优化。最后通过输出层输出共3 个尺度、27 个通道的输出特征。对每个尺度的输出特征计算预测损失,模型损失可以分成边框回归损失、置信度损失和类别损失三项之和,见公式(6),其中边框回归损失选用GIoU损失进行计算,见公式(7),回归损失和置信度损失都采用BCEWithLogitsLoss 损失,见公式(8)和公式(9),最后通过损失反向传播优化参数,从而获得磁瓦的检测模型。
其中:C是真实框A和预测框B的最小外接框。
其中:σ(g)为通过sigmoid函数的输出。
图8 改进磁瓦缺陷检测网络
3 实验
3.1 实验数据集
实验所选用的图片数据均来自现场采集并对采集后的图像进行图像增强处理,数据集共332 张,分为四种缺陷类型,分别是表面缺陷(surface)、内部损伤(intra)、崩块(break)和孔洞(hole),并用labelimg 为每张图片进行缺陷标注。在进行网络训练时将数据集随机按照9∶1的比例划分训练验证集和测试集,另外训练集和验证集按照8∶2 的比例划分。同时对训练的数据集进行随机翻转、平移和mosaic 增强等操作来对数据增强。
3.2 评价指标
为了客观地对磁瓦缺陷检测效果进行评价,本文选取精确率(Precision)、召回率(Recall)以及全类平均精度(mAP)作为评价指标,评价指标的定义见公式(10)~公式(12)。
其中:ri是按升序排列的Precision的插值段第一个插值对应的Recall,k是类别数量,Pinter是排序后当前点以及后面点所有Precision中最大的值。
3.3 实验结果
所有实验均在PyTorch 1.7、 Python3.8、Cuda10.2开发环境下进行,并且训练的批处理数量大小均为16、迭代训练周期都为2000 次,学习率初始值均为0.001。表1 对比了几种主流的目标检测算法和本文改进算法应用在磁瓦缺陷检测中的检测结果,可以看出本文改进算法的mAP比YOLOv3高出5.3个百分点,比YOLOv5高出1.5 个百分点。其中一些检测效果如图9 所示,从图9可以看出本文改进算法能够检测出不明显的缺陷,并且缺陷目标检测的置信度和预测结果都有很好的效果。
表1 磁瓦缺陷检测算法对比
图9 各种算法检测效果
在测试集中,改进算法的各个类别预测RP曲线如图10 所示,图例中的数值为AP值(平均精度),从图10 可以看出孔洞(hole)的检测性能相对较弱,这是因为训练数据集中孔洞样本的数量较少使训练不充分造成的,另外几类的AP值都较高,说明该改进算法能较好地完成磁瓦缺陷检测任务。
图10 改进算法各个类别PR曲线
3.4 消融实验
为了证明CSPA 模块的有效性,在同样的训练环境和相同的参数下分别使用Swin-Transformer+CSPA 模块改进的neck(编号1)、Swin-Transformer+YOLOv3 的neck(编号2)进行实验,评价指标为Precision、Recall以及mAP。对照结果如表2 所示,可以看出使用CSPA 模块的neck 对磁瓦缺陷检测的mAP有1.4 个百分点的提升。
表2 消融实验
4 结语
针对磁瓦缺陷检测特征提取困难和基于传统卷积神经网络对全局特征表述较差、上下文不交互等缺点,本文通过设计增强滤波器对磁瓦图像中的缺陷进行增强,以及提出基于Swin-Transformer 和结合YOLOv3 的neck 模块的磁瓦缺陷检测算法,且为了获取更丰富梯度信息,在neck部分优化设计了CSPA模块。根据实验结果可知对采集到的磁瓦图像的缺陷增强后再运用Swin-Transformer 相比于主流目标检测算法具有更高的检测精度,另外Swin-Transformer 相比于卷积神经网络能更好地提取全局特征和上下文的交互,从而提高特征提取的能力,对于一些不太明显的特征也能很好地提取出来。