基于OCR-UNet的金属表面缺陷分割*

陈德阳,唐 智,何牧耕

(重庆大学机械与运载工程学院,重庆 400044)

0 引言

金属表面缺陷检测对于产品质量控制和铁路等日常维护至关重要[1]。传统的人工表面缺陷检测主要通过人眼进行,具有主观、耗时耗力且准确性不高等缺点。近年来,机器视觉和图像处理技术因其快速、准确、无损和低成本的特点,广泛的应用于表面缺陷检测系统[2]。传统图像处理法利用局部异常反映的原始属性来检测和分割缺陷,可进一步分为结构法[3]、阈值法[4]、滤波器法[5]和基于模型[6]的方法等。传统方法极度依赖缺陷特征的选择,当光照、背景和摄像机角度发生变化时,缺陷检测的性能将会受到很大影响,同时,表面缺陷的复杂性与多样性,一直是表面缺陷检测待解决的难点。

随着深度学习的兴起,越来越多的学者开始使用深度学习方法进行缺陷检测。目前应用于缺陷检测的深度学习算法以监督学习方法为主,它需要使用带标签的数据集对深度学习网络进行训练。NIE等[7]通过不同的主干网络与DeepLabV3+网络结合对钢铁图像上的缺陷区域获得了较好的分割性能,但其对低级特征的利用率较低,难以准确分割出较小的缺陷。DU等[8]提出一种双流卷积神经网络,在X射线分割领域表现优异,但在卷积下采样过程中有一定的分辨率信息的损失,从而降低了模型对小缺陷的分割性能。LING等[9]提出一种深度连体语义分割网络,实现PCB焊接小缺陷的语义分割,不过该方法依赖于标准的模板图像,不适用于缺少模板图像的缺陷检测场景。WANG等[10]提出一种基于FCN的轮胎X射线图像缺陷检测方法,实现了轮胎图像中缺陷的定位与分割,但其主要适用于分割单个缺陷,在存在多缺陷的场景下难以满足缺陷分割要求。许国良等[11]提出一种基于协同注意力的小样本手机屏幕缺陷分割网络,使用少量缺陷样本完成手机屏幕缺陷分割。

深度学习法提取的特征代表性强,模型具有较高的识别准确率和泛化能力,能够端对端的进行缺陷检测,因此,本文以UNet[12]网络为基础,提出一种基于OCR-UNet的分割网络用于金属表面缺陷检测,其具有强大的全局上下文建模能力,且能够克服金属表面缺陷样本与背景极度不均衡的问题,最终实现对金属表面缺陷尤其是小型缺陷的准确分割。

1 UNet网络

UNet网络是一种广泛的应用于医学图像分割领域的语义分割网络,它适用于训练边缘不规范的小目标数据集,由编码网络和解码网络组成,形成U型网络结构,如图1所示。编码网络有5个特征提取阶段,包括4次下采样,每次下采样都将缺陷特征图的大小减半,通道数加倍,特征图从扁平状变得小而密,更加立体;解码网络有4个特征融合阶段,包括4次上采样,每次上采样都将缺陷特征图的大小加倍,通道数减半,特征图最终恢复到输入时的大小,能保留更多的缺陷位置等细节信息,对小缺陷的分割效果显著。特征融合时采用跳跃连接来复用低级特征,保证缺陷高级语义信息与低级位置信息的融合,多级多尺度特征的提取与融合对尺寸大小变化的缺陷识别更有利。

图1 UNet网络结构

2 基于OCR-UNet的表面缺陷分割网络

虽然UNet网络在多尺度的缺陷分割方面能表现出较好的性能,但是它仅融合了缺陷周围局部范围内的信息,忽略了全局范围内各个缺陷之间的相互关系,难以同时识别出同一图像中形状、大小变化的缺陷,提取的特征对多缺陷共存的场景表示能力不足,因此,有必要增强全局范围内的像素类别信息,以更准确的识别出存在的缺陷。为此,本文引入全局目标上下文卷积神经网络(object-contextual representations CNN,OCRNet)[13]以增强特征对分割的表达能力,记作OCR-UNet,它能够弥补上下文信息的不足,包含粗分割分支和精分割分支,利用粗分割分支的类别概率增强像素特征表达,显式的增强像素的类别表达能力并用于精分割分支获得更好的分割结果。

2.1 网络架构

图2 OCR-UNet网络结构

(1)

式中:g(·)为转换函数,表示首先利用1×1卷积层—BN批归一化层—ReLU激活函数层的结构调整各级特征的维度至一致;然后,通过双线性插值的方法将低分辨率的特征图上采样到与最大尺度的特征图大小一致;∪(·)表示特征拼接与融合函数,表示将所有尺度一致的特征图以维度拼接的方式拼接在一起并使用1×1卷积融合,得到用于粗略分割的最终特征图Fcoarse∈Rh×w×c5,h、w、c5分别表示粗分割特征图的长、宽和通道数。

(2)

式中:h(·)为转换网络,其网络结构为3×3卷积—BN批归一化层—ReLU激活函数层。

增强特征由像素特征与粗分割模块的类别概率分布图相互聚合生成,能有效提高像素的类别表达能力,加强特征对缺陷的描述,提升缺陷分割的性能,特征聚合模块的聚合过程[15]如式(3)所示。

(3)

式中:Softmax(Fcoarse)表示由粗分割特征获取缺陷类别概率分布图,像素特征Fp和粗分割特征Fcoarse通过3个转换矩阵Wq、Wk、Wv分别生成3个关系矩阵Q、K和V,dk表示关系矩阵K的通道维度,Fz∈Rh×w×c7代表增强了像素类别的表达能力的增强特征,该特征以计算自注意力的方式融合了全局范围内的上下文信息。

为了保证网络保留基本的初始像素特征信息,将像素特征生成模块得到的像素特征与增强特征以维度拼接的方式拼接在一起,并利用特征融合模块将特征融合,获得精分割特征图,特征融合模块的结构为1×1卷积—BN批归一化层—ReLU激活函数层,此过程如式(4)所示,Ffine∈Rh×w×c8表示精分割特征图。

Ffine=∪(Fp,Fz)

(4)

2.2 损失函数

缺陷分割问题可以归结为判断像素是背景或缺陷的二值分类问题,图像分割神经网络最常使用的损失函数是二元交叉熵损失(binary cross entropy loss,BCE Loss)函数,如式(5)所示。

(5)

式中:tij代表某像素真实类别标签,pij为某像素预测类别标签,W、H分别代表图像的宽度和高度。

可以看出交叉熵损失采用的是类间竞争机制,重点学习类别占比高的样本,无法解决数据不均衡问题。由于本文背景负样本的数量远大于缺陷正样本的数量,如果仅采用BCE Loss,网络训练时将陷入局部最小值。为了降低前景与背景比例极度不均衡带来的影响,本文引入交并比损失(intersection over union loss,IoU Loss)函数,它是一种衡量预测值与真实值的匹配程度的损失,它更注重前景图像分类的准确性,而对大量的背景像素不敏感,它的计算过程如式(6)所示。

(6)

由上式可以看出,IoU Loss更关注不容易分割的缺陷样本,可以有效缓解缺陷前景和背景的类别不平衡,但是如果仅用IoU Loss,在部分缺陷预测错误时将导致梯度太大而使训练过程不稳定,产生振荡现象。因此,本文设计一种BCE Loss与IoU Loss共同作用的加权损失函数,记作WLoss,如式(7)所示。

Lw=wt·Lbce+(1-wt)·Liou

(7)

式中:wt表示两种损失之间的权重。

由于缺陷的边缘与背景具有较高的相似度,造成缺陷边缘等细节分割不准确,影响表面缺陷分割的整体准确性,因此本文添加细节聚合模块[14],用于增强低级的边缘等空间细节信息的表达能力,提升网络对缺陷边缘的识别能力。使用缺陷边缘标签生成模块对缺陷标签进行处理,如图3a所示,首先对缺陷标签进行步长为1、2、4的Laplacian卷积,获得3个尺度的缺陷边缘信息;然后将小尺度的标签图用双线性插值的方法上采样到原始大小并通过1×1卷积融合3个通道的特征,获得最终的缺陷边缘的标签信息,图3b为本文使用的Laplacian卷积核。

(a) 缺陷边缘标签生成模块(b) Laplacian卷积核

获得缺陷边缘的标签信息之后,由UNet骨干网络生成的基本特征图通过细节分割模块获得缺陷的边缘细节分割图,进而计算边缘分割损失Ledge,本文的网络总损失Ltotal由粗分割损失Lcoarse、精分割损失Lfine和边缘分割损失Ledge共同组成,如式(8)所示。

Ltotal=Lcoarse+Lfine+Ledge

(8)

3 实验验证结果分析

为了对本文所提算法做出客观公正的评价,所有实验均在相同的实验环境下进行。本实验在Linux操作系统下完成,使用Pytorch版本为1.7.0的深度学习框架,后端使用CUDA 11.0软件开发工具包来使用GPU加速,GPU加速器使用的是Tesla P100,在训练过程中采用梯度裁剪方法,防止模型训练时发生梯度爆炸或消失的现象。

3.1 实验数据集

本文选用BSD[15]和RSDD[16]数据集来验证所提出的算法的有效性。BSD数据集为从某机床滚珠丝杠主轴上采集的表面点蚀缺陷图像,RSDD数据集为从快速轨道获取的铁路轨道表面缺陷数据集。为便于网络训练,使用随机翻转、旋转、亮度变换的数据增强扩增数据集样本并裁剪为特定大小。各数据集的图像尺寸和训练集(Train set)与验证集(Test set)图像数以及图像尺寸(Image Size)如表1所示。

图4展示了数据集的部分图像,从中可以观察出,有的缺陷与背景相比占比很小,肉眼几乎难以识别;有的图像除缺陷之外还包含大量污染物,对缺陷的识别造成干扰;光照不均匀也是部分缺陷难以同时被识别出的一个重要因素。

(a) BSD数据集部分图像

(b) RSDD数据集部分图像

3.2 模型训练与结果分析

本文在上述两个数据集下分别训练了本文提出的OCR-UNet网络模型,共计迭代次数设为200 epoch,加权损失函数的权重wt设为0.3,优化函数使用AdamW,初始学习率设为1e-4,批大小为8。

图5为各数据集训练过程的训练集损失(train loss)、测试集损失(test loss)和缺陷分割准确率(Acc)变化曲线,由图中可以看出,BSD数据集最终的缺陷分割准确率保持在0.9左右,RSDD数据集的缺陷分割准确率稳定在0.85左右。

(a) BSD数据集训练过程 (b) RSDD数据集训练过程

在本文的实验验证过程中,除了实现本文所提出的方法,还对比了传统的分割网络模型FCN[17]、DeepLabV3+[18]以及UNet。表2和表3为本文方法与几种对比方法在两个数据集上缺陷分割的定量评价结果,其中LI(loss improving)表示在OCR-UNet网络基础上以本文损失函数改进常用的二元交叉熵损失函数。使用交并比(intersection over union,IoU)和马修斯相关系数(matthews correlation coefficient,Mcc)[19]作为本文的像素级分割结果的评价指标。同时,为了增强小缺陷图像在评价结果中的贡献程度,使用图像级的平均交并比评价指标,计算公式如式(9)所示。

表2 BSD数据集实验结果

表3 RSDD数据集实验结果

(9)

式中:DIoU表示图像级的平均缺陷分割交并比,IoUi表示某缺陷图像的缺陷分割交并比,m代表缺陷图像总数。

从表2可以看出,本文方法在BSD数据集上的缺陷分割效果较对比方法更优,达到了0.824的IoU值。由于UNet提取的特征是多级多尺度且相互融合的,故能实现比FCN和DeepLabV3+更好的效果,而本文提出的OCR-UNet网络增强了特征对全局范围内的上下文信息的表达能力,能更准确的分割出金属表面的缺陷。使用本文提出的损失改进以往图像分割常用的交叉熵损失,能有效避免局部最优结果,突出缺陷分割结果对网络损失的贡献,实现更好的缺陷分割性能。

对于RSDD数据集,由于其图像尺寸高宽比例极大,导致缺陷分割任务的难度巨增。由表3可以看出,FCN和DeepLabV3+网络的分割结果较差,UNet相较于前者有较大的分割效果提升,本文的OCR-UNet方法相较于其他几种方法仍有较大提升,其IoU值依然能达到0.806。

对于BSD数据集而言,其图像级的平均交并比(DIoU)与像素级的交并比(IoU)相比差距较大,而RSDD数据集的这两个指标差别较小,这是因为BSD数据集相较于RSDD数据集而言,小缺陷图像占数据集的比例更大,而对比方法对小缺陷的识别性能不佳,但无论是哪一个数据集,本文方法都能保证比其他对比方法更高的DIoU指标值,证明本文方法在保证缺陷分割总体性能的同时,对小缺陷的分割效果更好。

3.3 缺陷分割主观视觉效果分析

图6为各数据集缺陷分割的视觉效果对比示例,由图6a可以看出,FCN没有检测出小缺陷样本,DeepLabV3+的分割结果存在碎片化的情况,UNet受环境干扰而将不是缺陷的区域分割为缺陷,存在错检的问题,本文方法不仅分割结果与标签最接近,而且能够将非常微小的缺陷分割出来;图6b也可以看出,FCN和DeepLabV3+在分割小缺陷时存在严重的漏检情况,UNet依然有部分错检的问题,本文方法相较于其他几种方法能更好的将可能存在的缺陷分割出来,且分割结果边缘清晰,与真实情况基本一致。从定性和定量结果可以看出,本文方法在UNet网络基础上,通过融合多级多尺度的输出特征,添加全局目标上下文注意力模块,改进损失函数的方法,使本文模型更适应于形态多变的金属表面缺陷,相较于几种传统分割网络模型,其准确性与可靠性更高。

(a) BSD数据集缺陷分割结果对比

(b) RSDD数据集缺陷分割结果对比

4 结束语

针对金属材料表面缺陷难以准确识别的问题,本文基于深度学习领域的UNet网络模型,提出OCR-UNet网络模型,通过对网络结构和损失函数等改进,实现了比以往更高的识别率。文中的主要研究工作为:

(1)基于传统UNet网络进行改进,融合解码的多级多尺度特征以适应形状、大小变化的缺陷,引入全局目标上下文网络结构,增强特征的类别表达能力,实现全局范围内同目标之间的信息交互,有效克服卷积所提取特征的局部性问题,最后,通过引入交并比损失与边缘分割损失来改进网络训练时的损失函数,改善缺陷前景与背景极度不均衡的问题,实现准确的金属表面缺陷分割;

(2)与FCN、DeepLabV3+和UNet网络模型进行对比,本文改进之后的模型在BSD和RSDD数据集上的IoU分别达到0.824、0.806,Mcc分别达到0.903、0.892,在分割指标上均有明显的提升,验证了其可行性和有效性。