Unet图像缺陷检测入门基础
本文总览缺陷检测工作流
前言
本文为目标检测与缺陷检测工作流基础入门文章,属于搬运内容,原作者: Vinithavn, 原文链接在文章末尾。
1. 简介
什么是目标检测?
给定一张图片,人类可以快速识别其中的物体。例如,我们能够判断图片中是否有汽车、树木、人物等。如果我们可以分析图像并检测物体,那么是否也可以教会机器做到这一点?
答案是肯定的。借助深度学习和计算机视觉技术,我们可以实现目标检测的自动化。通过构建深度学习和计算机视觉模型,机器不仅可以检测和定位物体,还可以计算它们之间的距离,预测它们的未来状态等。目标检测技术在计算机视觉和机器学习中有着广泛的应用,例如目标跟踪、监控系统、人类行为识别甚至自动驾驶汽车等。
请参考下图:

图1:道路交通的目标检测
该图显示了一辆车视角下道路交通的目标检测结果。我们可以看到模型检测出了其他车辆、交通信号灯等。如果这是一辆自动驾驶汽车,它需要能够检测驾驶路径、其他车辆、行人以及交通信号灯,以实现平稳和安全的驾驶。
进一步理解目标检测后,我们将介绍一种更高级的技术——图像分割。
目标检测与图像分割的区别
通过分析下图,我们可以轻松理解两者之间的区别:

图2:目标检测与图像分割
两种方法都试图在图像中识别和定位对象。在目标检测中,使用边界框来实现,即模型通过绘制矩形框来定位对象。而在图像分割中,每个像素都会被标注。这意味着分割模型会对图像中的每个像素进行分类,将其归类为有意义的对象类别。由于预测每个像素的意义,该过程也称为密集预测(dense prediction)。
“图像分割的输出形式称为掩膜(mask),即与原始图像大小相同的图像,其中每个像素的值表示对象是否存在。”
在本案例中,我们将使用图像分割技术。如果您想了解更多关于图像分割的知识,可以参考此处。
理解目标检测与图像分割后,让我们进一步了解本案例的具体问题陈述。
2. 问题描述
我们提供了一组产品图像,其中部分产品存在缺陷,而部分没有。给定一张产品图像,我们需要检测它是否存在缺陷,并定位这些缺陷的位置。
3. 机器学习任务定义
问题可以被定义为一个图像分割任务。给定一张产品图像,我们需要绘制其分割掩膜。如果产品有缺陷,分割掩膜应能够标注出缺陷的位置。
4. 性能评估指标
图像分割问题中最常用的评估指标之一是交并比(IoU)得分。如下图所示,IoU得分的计算方式为:

IoU = 预测分割区域和真实分割区域的交集面积 ÷ 预测分割区域和真实分割区域的并集面积.
公式: IoU = TP / (TP + FN + FP)
该指标的取值范围为0到1,其中1表示完全重叠,0表示完全不重叠。
本案例中使用的损失函数是Dice损失,
其公式为:
Dice系数 = 2 × 交集面积 / (预测面积 + 真实面积)
Dice损失 = 1 - Dice系数
5. 理解数据
数据集包含两个文件夹:train和test。
train文件夹包含六类图像,每类图像被分为两个子文件夹,一个包含1000张无缺陷图像,另一个包含130张有缺陷图像(文件夹名称以_def结尾)。
test文件夹包含120张有缺陷图像,需要预测其分割掩膜。

图3:训练数据集结构
6. 数据准备
6.1 准备图像和分割掩膜
现在我们需要为每张图像准备相应的图像数据和分割掩码。这些图像被分为十二个文件夹。以下是一些图像的示例。


第一张图像表示有缺陷的产品,而第二张图像表示没有缺陷的产品。接下来,我们需要为这些图像准备分割掩码。分割掩码应检测图像中的缺陷部分。对于上述图像,期望的分割掩码如下所示:


图4:样本产品图像,图5:对应的分割掩膜
对于无缺陷的图像,其分割掩膜为空白图像。

幸运的是,我们还提供了一个文件,包含绘制掩膜的关键信息,如椭圆的半长轴、半短轴、旋转角度、中心坐标等。利用这些信息,我们可以通过skimage库绘制椭圆分割掩膜。
通过使用以下所示的 get_data 函数,可以获取用于绘制椭圆的数据:

我们可以利用这些信息并使用 skimage 函数绘制椭圆形分割掩码。
需要注意的是,这些掩码信息仅适用于有缺陷的图像。对于无缺陷的图像,我们需要创建空白图像作为分割掩码。

6.2 数据加载
结构化数据以如下形式提供:

最终,图像路径和对应掩膜被组织为结构化数据。然后再使用pytorch Dataloader进行图像加载


7. 建模
我们选择使用广泛应用于图像分割任务的UNet模型。

图7:UNet架构
模型的结构类似于英文字母“U”,因此得名Unet。模型的左侧是收缩路径(也称为编码器),它用于捕捉图像中的上下文信息。编码器由传统的卷积层和最大池化层堆叠而成。池化层的作用是减少图像的高度和宽度,同时增加深度或通道数。在收缩路径的末尾,模型能够理解图像中的形状、模式、边缘等特征,但会丢失这些特征的“位置信息”。
由于我们的目标是生成图像的分割图,单靠从收缩路径获得的信息是不够的。我们需要一个高分辨率的输出图像,其中每个像素都被分类。
“如果我们使用一个普通的卷积网络,它包含池化层和全连接层,我们只会保留‘是什么(WHAT)’的信息,而丢失‘在哪里(WHERE)’的信息,这并不是我们想要的。在分割任务中,我们需要同时保留‘是什么’和‘在哪里’的信息。”
因此,我们需要对图像进行上采样以保留“位置信息”。这一过程在模型右侧的扩展路径(也称为解码器)中完成。扩展路径用于通过上采样技术对捕获的上下文信息进行定位。有多种上采样技术,例如双线性插值、最近邻方法、转置卷积等。如果你有兴趣了解更多关于上采样技术的内容,可以参考相关资料。
8. 训练
现在,我们已经准备好所有用于训练的数据,并选择好了模型。接下来开始训练模型。

由于无缺陷图像的数量远高于缺陷图像,为了取得更好的结果,我们仅从无缺陷图像中取一个样本集。模型使用Adam优化器进行编译,并采用Dice损失作为损失函数。使用的性能指标是IoU得分。
经过10轮训练后,我们实现了0.98的IoU得分以及0.007的Dice损失,这已经相当不错。以下是一些图像的分割结果。

我们可以看到,模型能够预测出与原始分割图相似的分割图。
9. 测试数据分割图的预测
现在,让我们尝试解决实际问题,即预测并绘制测试图像的分割掩码。下图显示了一些测试图像的预测分割结果。




可以看到,模型在测试数据上的表现良好,能够检测出测试图像中的缺陷。
10. 改进措施
正如上文提到的,缺陷图像的数量明显少于无缺陷图像。因此,我们可以通过对缺陷图像应用上采样和数据增强技术来改进训练过程。