基于双动态头Sparse R-CNN的表面缺陷检测算法

郑亚睿,蒋三新

(上海电力大学电子与信息工程学院,上海 201306)

0 引言

在工业生产中,产品的表面缺陷检测至关重要。如何高效找出缺陷,保证产品的产量和良品率,一直是研究热点。目前,传统人工筛选缺陷的方法已逐步被机器视觉自动检测器替代,然而,机器视觉算法[1-4]聚焦在构建模式识别器,没有从缺陷图像中学习到有效特征,复用性不大,很难对各类缺陷进行迁移。

近年来,基于深度学习的算法由于强大的特征提取能力在各领域得到广泛应用,给缺陷检测带来了新思路[5]。D. Weimer等[6]提出了用于工业故障检测的深度CNN架构,可以区分重纹理背景下的表面划痕,并建立了具有少量先验知识的产品交付检测系统;J. Shi等[7]提出了一种多通道掩模区域的卷积神经网络来解决木材表面缺陷的问题,并使用快速浏览网格来确定图像是否适合进一步的检测,减少了检测时间;B. Su等[8]采用Faster R-CNN+FPN作为基本探测器,将一种新的双向注意特征金字塔网络(BAFPN)集成到RPN中,检测光伏电池的缺陷;T. Nakazawa 等[9]通过使用28 600幅合成的22类晶圆图对CNN进行预训练,实现对真实晶圆图像的高精度识别,并且设计的模型能够在没有真实数据的情况下实现对罕见事件的检测,这可以帮助技术阶段的开发。

表面缺陷检测任务是目标检测的一个分支。当前主要的目标检测模型,无论是基于一阶段的YOLO[10]、RetinaNet[11],还是两阶段的Faster R-CNN[12]、Cascade R-CNN[13]等都设置了密集的候选预测框,需要非极大值抑制(Non-Maximum Suppression,NMS)[14]来处理这些冗余,增加了计算成本。并且虽然NMS过程中处理了重复识别问题,但是对于重叠目标和相近目标,会因为置信度较低没有正确识别出来。DETR[15]是去除NMS,使用标签比配进行端到端预测的模型,但是由于其基于Transformer架构,模型体量较大,收敛非常缓慢。对此,根据以上问题,提出双动态头Sparse-RCNN算法,去掉了区域提议网络(region proposal network,RPN)[12]、NMS等人工设计的组件,进行缺陷目标的稀疏预测。有如下贡献:

(1)改进了缺陷图像的特征提取。体现在2个部分:在网络颈部特征金子塔网格和特征提取动态头配合对缺陷进行尺度和空间的感知;在检测头部分针对多头注意力后期注意力图相似度高的问题,提出了交流注意力模块,混合各头部信息,提高缺陷特征的细节提取。

(2)提出多分支的动态实例交互头,并用“双查询”方式取代RPN的密集预测方式,去除了NMS操作,实现端到端的预测。

(3)改进了边界框回归损失函数,使用Alpha-CIoU[16]加速边界框收敛,提高检测精度。

1 算法介绍

本文算法的总体结构见图1,主要分为3个阶段,第1阶段首先使用ResNet骨干网络进行图像的基本特征提取,然后输入网络颈部,由特征金字塔网格与特征提取动态头进行不同尺度和空间的特征感知。第2阶段是网络的检测头部分,包含动态实例交互头和2个特殊的提议,提议框和提议特征,是本算法的核心部分,称为“双查询”操作,具体定义如下:

图1 双动态头Sparse R-CNN模型总体结构

(1)可学习的提议框。网络中使用一组固定的可学习提议框进行区域提议,替代了用区域提议网络RPN进行大量预测。提议框由四维参数表示,分别为中心横纵坐标、高度以及宽度,参数在训练中通过反向传播进行更新。这些可学习的提议框可看作是训练集潜在目标位置的统计数据,是图像中最可能包含目标的区域的初始猜测。

(2)可学习的提议特征。虽然提议框的四维信息提供了物体的大致定位,但是缺失了大量细节信息,如姿态、形状。故引入提议特征(N×d,N是提议特征的数量,d是维数),数量与提议框相同,是一个高维潜在向量,包含丰富的实例特征。最后一阶段是标签匹配算法,进行预测集和标注集信息的一一匹配。本文首先详细介绍特征金字塔网格与特征提取动态头模块,然后是提出的交流注意力模块和动态实例交互头的细节,最后说明标签匹配算法与改进的损失函数部分。

1.1 改进的网络颈部架构

原Sparse R-CNN网络颈部仅使用特征金字塔(feature pyramid networks,FPN)进行多尺度特征融合,本文改进为更深层的特征金字塔网格(feature pyramid grids,FPG)[17]和特征提取动态头结合的模块。

1.1.1 特征金字塔网格

特征金字塔网格是一个多路深层改进的特征金字塔,将特征尺度空间表示为规则网格,并由多路的横向连接融合,从而实现各层次的信息交换。如图2所示,存在两种主干通路以及4 种横向连接,具体连接方式如下:

图2 特征金子塔网格结构图

主干通路相同尺度的特征张量属于一个网络阶段,特征张量的空间步长从前到后逐渐增大。金字塔通路代表跨尺度的信息。从最高分辨率阶段到最低分辨率阶段,以自底向上的方式构建它们。首先,该通路的第一特征图由对应的高分辨率主干或金字塔级通过 1×1 横向卷积连接实现。其次,下采样后使用步长2的3×3卷积,在金字塔路径中创建各个较低层次的特征图。因此在各金字塔通路中,特征层次由多尺度的特征图构成,且各阶段的空间分辨率与主干通路一致。

Across-same:将同一水平特征通过主干通路连接起来,通过在特征图上使用1×1横向卷积来投影特征,并使之与相邻通路中相应的特征融合。

Across-up:直接连接构建了跨水平的自底向上路径,通过步长为2的3×3卷积将低层特征图降采样至一半大小,再与高层特征图融合。缩短了低层特征到高层特征的路径。

Across-down:通过整合向下的连接来实现自上而下的信息流。首先用2倍放缩比例的最近邻插值对高层特征图上采样,然后利用 3×3 卷积使之变为可学习的,上采样特征与下采样特征融合。

Across-skip:为简化如此宽泛网络的训练,在相同层级的首个通路及其后一个通路之间使用1×1卷积添加了跳跃连接。

1.1.2 特征提取动态头

缺陷图像中多个不同尺度的物体经常共存于一幅图像中,并且物体可能在不同的视点下以不同的形状出现。受文献[18]的启发,本文将尺度和空间2种感知结合起来形成轻量化的特征提取动态头,增强特征的提取,网络结构见图3。如果把特征金字塔网格的输出看作是Level×Space×Channel的三维张量,可以将两种感知分别部署在每个特定维度上。尺度感知的模块部署在维度上,学习不同的语义层次,根据单个目标的规模在适当的层次上增强该特征。空间感知模块部署在维度上(即高度×宽度)。它学习空间位置上的连贯区别表征。这样,就实现了动态头的统一感知机制。

图3 特征提取动态头结构图

首先,将特征金字塔网格输出的特征图通过不同倍数的上采样统一到统一尺寸,表示为一个4维的张量F∈RLHWC,L表示金字塔层数,H、W、C分别表示特征的高度、宽度和通道数。定义S=HW,张量重塑为F∈RLSC。

特征动态头的总感知公式为

W(F)=πL(πS(F)·F)·F

(1)

式中:πL、πS为适用于维度L、S的不同感知函数。

πL、πS具体公式为:

(2)

(3)

式中:σ(·)为1×1卷积层近似的线性函数;ϑ(x)为hard sigmoid激活函数,ϑ(x)=max(0,min(1,(x+1/2)))。

πS首先使用可变形注意学习稀疏化,然后在相同的空间位置跨层聚合特征。其中,k是稀疏采样点的数量,Δpk是空间位置偏移量,通过自学习空间补偿来关注有区别的区域,Δmk是pk位置处的自学习重要标量。这两种感知方式可以任意组合,实现不同效果的特征提取,本文动态头由6个动态块组成,采集图像不同角度内在信息。

1.2 改进的检测头架构

原检测头中,使用多头自注意力和单分支动态头进行提议特征和提议框的交互,本文改进为交流注意力和多分支动态头,增强模型特征匹配能力。

1.2.1 交流注意力

提议特征在输入动态实例交互头前,先经过交流注意力模块进行特征信息的梳理。

目前多头注意力[19]中,不同的注意力头进行单独的计算,最后求和。随着深度的增加,注意力头每个块之间学习到的注意力图相似度逐渐升高,表达能力不足以完整建模原图细节[20]。但是不同的注意力头相似度较低,学习到了图像不同的特征信息。

本文为了增强模型对图像特征的提取能力,减少注意图随着深度增加的相似度。提出交流注意力(speaking-head attention,SHA),将多头自注意力用一个参数矩阵重新融合成多个混合注意力,每个新的得到的混合注意力都融合了原先的各头部的注意力,交流注意力的结构如图4所示。

图4 交流注意力结构

本文的SHA模块输入为Q(Query)、K(Key)、V(Value)。使各注意力头部联系起来的主要操作是在Query与Key矩阵运算之后,softmax操作之前插入一个可学习的参数矩阵,混合各注意力头学习到的信息后,再与Value进行注意力计算。

具体计算公式如下:

(4)

(5)

SHA(Q,K,V)=concat(head0,…,headN)W1

(6)

式(4)为交流注意力计算。其中λ为可学习的参数矩阵,使用它将多头注意力映射混合成重新生成的注意力映射,然后再进行softmax计算。

式(5)为各个注意力头部的计算,式(6)为拼接各混合后头部操作。

1.2.2 动态实例交互头

动态实例交互头是本算法的核心,它通过输入特征映射、提议框和提议特征来生成预测对象的位置和特征,其分别和提议框、提议特征进行交互运算。

如图5所示,本文改进了动态实例交互头,通过GMP(全局最大池化)操作和一个全连接建立了一个分支网络,自适应地校准特征。本文在整个模型中堆叠6个级联动态头,每个动态头输出的对象特征和提议框将被用作下一个阶段的输入,具体交互过程如下。

图5 动态实例交互头

首先给定N个提议框,使用RoIrEAligns操作提取每个提议框的特征,提取到的区域池化(ROI)特征的大小为(S×C×C)(设置S=7,C=256)。每个区域池化特征Pi(C)与其提议特征进行一一对应,通过“动态实例层”进行融合。首先,将区域池化特征调整为一个尺寸为(S2,C)的二维向量,作为“动态实例层”的输入之一,并将提议特征直接输入一个全连接中,得到“动态实例层”的长度参数(2×C′×C),其中C′是“动态通道”的超参数。然后将向量等分为(C′×C),重塑为二维向量(C,C′)和(C′,C),分别作为“参数P1”和“参数P2”。动态实例层分别对以上参数进行矩阵运算,每次运算后进行BN和ReLU激活输出,最终得到融合特征,输入进前馈网络进行分类和定位。

1.3 标签匹配算法与改进的损失函数

检测网络的最后一部分,本文使用了三层感知机和线性计算层组成两个前馈网络来预测缺陷的类别和边界框。分类信息由softmax函数输出,边界框信息存于一个预测集合中,由N个固定大小的元素组成,分别为每个查询(Query)学习到绑定的目标位置信息。这并不意味着每个图像需要检测到N个目标,而是在学习中,使用匈牙利二分图匹配算法将输出的N个预测信息与真实的标注信息进行一一匹配,将标注中不存在的预测自动归为背景处理。凭借特殊的匹配算法,不需要在输出中使用非极大值抑制算法进行冗余框的处理。匈牙利二分图匹配算法如下:

(7)

LHung是基于匹配结果计算损失的公式。当预测的N个预测目标与标注信息逐个匹配时,使用损失函数优化。损失函数包括分类损失和边界框损失,表示如下:

(8)

Lbox包括IoU损失和L1损失,它计算预测框和标注框之间的中心点和高度差。计算公式如下:

(9)

式中Liou为预测框和真实框之间的联合交集(IoU),用于评估定位的精度。

原网络使用了GIoU,但是当检测框和真实框之间出现包含的现象时,GIoU就和IoU loss是同样的效果。缺陷数据集缺陷跨越范围大,并且非缺陷位置交错分布,缺陷样本形状各异。为了提高缺陷定位的精度。本文引入Alpha-CIoU作为IoU loss,Alpha是将Box-Cox变换应用于基础的IoU中,公式为

(10)

式中:ρ2α(b,bgt)为预测框与标注框中心点坐标的欧式距离;c为包住它们的最小方框的对角线距离;v为度量长宽比的相似性函数;β为权重函数。

Alpha的幂变换保留了作为性能度量的关键性质,包括非负性、不可分辨恒等式、对称性和三角形不等式。使用Alpha幂指数加成的CIoU有如下优点:

(1)可以梯度自适应加权的边界框回归精度。

(2)α可作为调节α-CloU损失的超参数以满足不同水平的边界框回归精度,其中α>1,通过更多地关注高IoU目标来获得高的回归精度。

(3)改进检测器的效果,在干净或嘈杂的环境下,不会引入额外参数,也不增加训练/推理时间。

2 实验结果与分析

为了评估本文算法的性能,本文将在多种表面缺陷数据集上与主流两阶段目标检测算法Faster R-CNN、Cascade R-CNN、Sparse R-CNN[21]进行对比实验。研究实验在配备了NVIDIA Tesla P100 GPU的Ubuntu操作系统下,采用Mmdetection深度学习框架构建网络。对所有训练情况,使用图像翻转和旋转的增强来增加多样性,并使用AdamW优化器进行优化,学习率为0.0001,设置动量为0.9。

2.1 数据集

采用晶圆表面缺陷数据集(wafer defect dataset,WDD)和热轧钢表面缺陷公开数据集(NEU-DET)训练和测试网络性能。 WDD和NEU-DET的样本分布情况分别见表1、表2。其中,WDD包含6 000张样本(分辨率为512像素×512像素),11类缺陷,分别为:外部异物(Foreign)、金属残余(Gold)、缺损(Incomp)、凸起肿块(Lump)、液体残余(Residual)、划伤(Scratch)、漏刻蚀(UBM)和4种原料不良(Raw I、Raw II、Raw III、Raw IV)。NEU-DET包含1 800张样本(分辨率为200像素×200像素),6类缺陷,分别为:氧化铁皮(Rolled in scale),色斑(Patch),裂纹(Cracking),坑面(Pitted surface),杂质(Inclusion)和划伤(Scratch)。

表1 WDD中各缺陷样本分布

表2 NEU-DET中各缺陷样本分布

2.2 评估方法

为了评估检测器的性能,使用了4个指标,分别为平均精度(average precision,AP)、平均不确定识别率(average dibious,AD)、平均遗漏率(average missed,AM)和平均误检率(average error,AE),其计算方法如下:

(11)

式中:Nc、Nw、Nd、Nm分别为这些缺陷对应的正确判断、错误判断、不确定识别、遗漏的样本数;Nall为样本总数。

为了评价该模型对所有类型缺陷的性能,提出了4个综合指标:综合平均精度(mean AP,mAP)、综合平均误检率(mean AE,mAE)、综合平均不确定识别率(mean AD,mAD)和综合平均漏检率(mean AM,mAM)。

(12)

(13)

式中:c为缺陷类别的数量,αi为i型缺陷的权重,即型缺陷样本在总样本中的比例。

2.3 性能评估

2.3.1 网络收敛

如图6(a)所示,本文算法与同样使用标签匹配方法的端到端检测器DETR和两阶段网络Faster R-CNN进行了收敛对比,相比DETR的150epoch才能达到收敛,本文算法收敛速度显著提升,36epoch时即达到收敛,并且对于同样能快速收敛的Faster R-CNN来说,本文算法在同样的epoch有着更高的检测精度。图6(b)为边界框的收敛情况,本文方法得益于Alpha-CIoU得到改善,在14epoch时已达到收敛,优于原Sparse R-CNN的24epoch。

(a)AP曲线收敛对比

2.3.2 缺陷检测实验

表3为Faster R-CNN、Cascade R-CNN、Sparse R-CNN和本文算法在晶圆缺陷数据集上的实验结果,分别测试了以ResNet-50和ResNet-101为骨干网络的情况。相比其他网络,本文方法在ResNet-101为基准时检测精度最高,mAP为94.3%,并且误检率、不确定识别率和漏检率显著减少,分别为2.7%、2.1%和1.3%。相较于其他检测方法,本文算法计算复杂度有所增加,有更多的参数需要学习训练,所以FPS较低。

表3 基于WDD数据集

如表4所示,本文也在NEU-DET的数据集验证了网络的检测以及泛化能力。对于非晶圆类缺陷也能很好定位和分类,热轧钢带上细微的划痕和裂缝也能识别。本文算法的各项指标也优于其他方法,mAP达到88.1%。但是误检率和漏检率高于晶圆数据集验证结果,分析是由于此数据集的背景更复杂,缺陷细小,对于网络来说缺陷的特征不易学习。

表4 基于NEU-DET数据集

为了进一步分析算法的性能,图7比较了采用ResNet-50为骨干网络的Sparse R-CNN和本文方法检测器在各类缺陷中的检测结果。从图中可以看出,本文算法对大多数缺陷类别的检测精度均有提升,误检、漏检情况有所下降。如图7(b)所示,不确定识别情况改进明显,“Incomplete”和“Gold”类最突出。但对于某些类别检测困难的样本,仍存在一些问题,分析了某些类别难以学习的一些原因。“Raw_Ⅱ”类与“Foreign”类精度较低,“Raw_Ⅱ”类漏检情况较多,达到8.2%。主要原因是此类缺陷与背景缺陷相似,若缺陷部位较小会被网络判定为背景处理。“Foreign”类误检率较高,为10%。此类缺陷大多为黑、灰亮色,形状各不相同。若分布面积较大,易被正确识别,否则易被识别为其他缺陷的部分区域。“UBM”与“Lump”类检测精度最高,均达到0.96以上,此类缺陷颜色对比度高,并且形状较完整简单,相比其他类缺陷易于学习训练。

图8展示了部分两类表面缺陷的检测结果,可以看到Faster R-CNN冗余框去除不够充分,定位不够准确,需要提高NMS阈值修改结果。相比之下,本文算法只保留了一个预测框,定位准确,并且跨非缺陷位置完整地包围住了主要缺陷,置信度也高于Faster R-CNN网络。

图8 冗余框消除结果

本文方法在边框提议阶段能够有效控制边框数量,并且由于特殊的标签预测匹配模式,也会减少冗余的产生,更好定位目标位置。

图9为交流注意力的可视化结果,可视化了特征提议阶段交流注意力对目标的关注情况,其中较深的颜色表示高度关注,而较浅的颜色表示模型对它的关注较少。可以看到一开始交流注意力模块关注了大范围区域,从全局进行特征信息的提取。后面逐渐缩小关注区域,定位到缺陷的具体区域,正确提取缺陷部位图像特征信息。

图9 交流注意力可视化

2.4 消融实验

为了验证双动态头Sparse R-CNN中每个模块的有效性,采用ResNet-50作为骨干网,在WDD数据集上进行了消融试验,实验中均使用相同的训练设置进行训练,结果见表5。可以看到,实验未使用任何改进模块的基准识别率为89.3%,当单独添加FPG和Dyhead时,精度分别从89.3%提升到90.1%和91.8%,结合两者使用时精度从89.3%提升到92.9%,提取特征能力显著增强。SHA模块使精度从89.3%提升到91.2%,Alpha-CIoU改进边框收敛速度的同时,检测精度也从93.7%提升到94.1%。本文方法能取得更好的表现,在于加强了对缺陷图像信息的细节关注以及使用了更精准回归边界框的损失函数,这两种方法不仅有助于关注图像的全局特征,而且可以很好处理缺陷复杂变化的定位问题。

以上结果均是在动态实例交互头添加了改进分支网络的基础上,表6单独验证了分支网络的效果,可以看到精度的提升,说明在分支网络对动态实例交互中的特征校准有所改善。

表6 动态实力交互头分支网络消融

对于本文1.2.1小节的交流注意力模块(SHA),其中注意力头部的混合方式为在softmax操作之前插入混淆矩阵λ。对于混淆注意力头部信息的方式,本部分测试另一种方式并展示在表7中,其中第2行采用了在softmax操作后直接插入混淆矩阵λ的方式,之后再进行注意力计算。从准确率上这样的方式也提高了检测结果,但是不如本文采用的方法效果好,相差0.9%。

表7 SHA添加位置消融实验

3 结论

本文针对目前缺陷检测算法对目标冗余识别的问题,提出双动态头Sparse R-CNN端到端的表面缺陷检测算法,融合了特征金字塔网格和特征提取动态头,对图像进行不同尺度和空间方向的特征提取,并提出交流注意力打通不同注意力头部的信息,保留提议阶段缺陷的细节特征。用动态实例交互头和双查询方法取代了RPN,并去除NMS模块,实现目标的稀疏预测,减少了大量冗余计算。此外,使用Alpha-CIoU损失改进了边框的回归速度。在晶圆表面缺陷数据集WDD和热轧钢表面缺陷NEU-DET数据集上的实验表明,该算法使检测精度分别提高到94.3%和88.1%,可以有效实现对各类缺陷的识别和定位,具有应用价值。同时,本研究仍存在算法模型大、参数量多的不足。在后续的研究中,将研究改进为轻量化网络,提高实时的检测效率。