TLU - Net：一种用于钢材表面缺陷自动检测的深度学习方法

摘要： 钢铁表面缺陷检测是钢铁板制造过程中的一个关键步骤。近年来，已经研究了许多基于机器学习的自动化视觉检测 (AVI) 方法。然而，由于 AVI 方法的训练时间和准确性问题，大多数钢铁制造行业仍然使用人工视觉检测。自动钢铁缺陷检测方法可以用于更经济、更快速的质量控制和反馈。但是，为分割和分类准备带注释的训练数据可能是一个昂贵的过程。在这项工作中，我们提出使用基于迁移学习的 U-Net (TLU-Net) 框架进行钢铁表面缺陷检测。我们使用 U-Net 架构作为基础，并探索了两种编码器：ResNet 和 DenseNet。我们使用随机初始化和使用 ImageNet 数据集训练的预训练网络来比较这些网络的性能。实验使用 Severstal 数据进行。结果表明，迁移学习在缺陷分类中的性能比随机初始化提高了 5%（绝对值）。我们发现，迁移学习在缺陷分割中的性能比随机初始化提高了 26%（相对值）。我们还发现，随着训练数据的减少，迁移学习的收益会增加，并且迁移学习的收敛速度比随机初始化更好。
关键词： 自动视觉检测 (AVI)、DenseNet、ResNet、表面缺陷检测、迁移学习、U-Net
1. 引言
钢铁是人类最重要的建筑材料之一。缺陷检测是钢铁板质量控制的关键步骤。这个过程主要涉及使用工业相机捕获钢铁表面的图像，然后识别、定位和分类缺陷，这有助于纠正缺陷的原因。通常，这个过程是手动完成的，这既不可靠又耗时。不可靠的质量控制可能会给制造商造成巨大的经济损失。可以使用计算机视觉方法的自动分类来替代或辅助手动检测。
2. 提出的基于迁移学习的 U-Net
我们提出的联合钢铁缺陷分割和分类的架构如图 2 所示。该架构接受一个尺寸为 H × W 的输入图像，并将每个像素分类为一种或多种类型的缺陷。它主要涉及四个部分：（1）U-Net 架构（2）初始化类型（3）分类（4）目标函数

2.1 U-Net 架构
U-Net 是一个具有跳跃连接的编码器-解码器架构。编码器使用编码器块对图像进行编码，并使用池化降低分辨率。这有助于提取图像的多尺度特征。解码器在每一步中上采样表示。跳跃连接可以使解码器选择不同尺度的特征，从而更准确地预测对象边界。
2.2 迁移学习
我们探索了两种用于迁移学习的编码器块。这两个网络都使用 ImageNet 数据集进行训练。我们在以下小节中简要回顾了这两个网络的特征。
2.3 分类
编码器输出对输入图像进行丰富的抽象表示。因此，我们提出使用空间平均池化来提取图像表示。图像表示通过具有 sigmoid 激活的线性分类器传递，以实现多标签分类。
2.4 目标函数
联合分割和分类问题被表述为两个损失的加权组合，如下所示：
3. 实验和结果
我们使用 Kaggle 竞赛 - “Severstal: 钢铁缺陷检测” 数据集进行所有实验。在每个实验中，输入图像可能包含一种或多种类型的缺陷。训练集包括 12568 张图像，其中 6666 张图像至少包含一个缺陷区域。专家进行了地面真实分类，以提供缺陷类型分类和通过视觉检查注释的缺陷区域。图像的分辨率为 256x1600 像素。我们使用全局平均值和标准差对图像进行归一化。我们应用随机垂直/水平翻转作为数据增强。相同的增强应用于原始图像和相应的地面真实掩码，以与增强的图像配对。
3.1 实验设置
我们使用 75% 的数据用于训练，12.5% 的数据用于验证，12.5% 的数据用于测试。所有实验都使用五个编码器和解码器的 U-Net。网络使用公式 1 中的目标函数进行训练，批量大小为 16，使用 Adam 优化器，学习率为 5 × 10−4，β1=0.99 和 β2=0.99。我们训练网络 10 个 epoch，并使用早期停止。我们将使用随机初始化的 U-Net 作为基线。我们在 PyTorch 中实现了网络，并使用了 PyTorch 分割库。具有随机初始化的 ResNet/DenseNet 分别表示为 ResNet(Random)/DenseNet(Random)。ImageNet 预训练的对应物分别表示为 ResNet(Imagenet)/DenseNet(Imagenet)。为了了解模型的样本复杂性，我们还使用 50% 的训练数据训练这些网络。我们将预训练初始化称为 TLU-Net，将随机初始化称为 U-Net。
3.2 评估指标
我们使用多标签分类准确率 (MLA) 和 4 个类别中平均受试者工作特征曲线下面积 (AUC) 来评估钢铁缺陷分类性能。MLA 定义为预测正确标签数与标签总数的比例。我们将多标签分类视为 4 个单独的二分类器，并计算 4 个分类器的平均 AUC。
我们使用 DICE 和交并比 (IoU) 来评估钢铁缺陷分割的性能。每个类别的 DICE 指标定义如下：
3.3 结果和讨论
图 5 显示了不同网络和初始化的平均 MLA(%) 比较。从图中可以看出，与 100% 训练数据的随机初始化相比，TLU-Net 在 MLA 方面实现了 5%（绝对值）的提高。这表明使用 ImageNet 学习到的特征可以帮助钢铁缺陷分类。当训练数据减少到 50% 时，MLA 差距增加到 8%（绝对值）。TLU-Net 的性能在训练数据减少时不会显着下降。这表明当标注数据点数量有限时，TLU-Net 是有帮助的。在 100% 和 50% 的训练数据的情况下，ResNet(ImageNet) 和 DenseNet(ImageNet) 的性能最好。这可能是因为 DenseNet 的参数比 ResNet 少，因此需要的数据更少，但表示能力也更弱。
图 6 显示了不同网络和初始化的 AUC。从图中可以看出，在所有情况下，所有类别中最好的 AUC 都是在类别 1 中实现的。类别 3 和类别 4 的缺陷表现出最差的性能。这主要是因为每个类别的训练数据样本数量。与 U-Net 相比，使用 TLU-Net 可以提高所有类别的 AUC。DenseNet(ImageNet) 具有最高的 AUC。
图 4 显示了不同网络在 100% 和 50% 训练数据情况下的 DICE 和 IoU 的箱线图。在所有情况下，TLU-Net 的中位数和平均性能都优于 U-Net。在 100% 训练数据的情况下，ResNet(ImageNet) 的 DICE/IoU 优于所有其他模型，并且其 75% 置信区间较小。使用 ResNet 的 TLU-net 的 DICE 比 U-Net 使用 ResNet 的 DICE 提高了约 26%（相对值）。使用 DenseNet 的 TLU-net 的 DICE 比 U-Net 使用 ResNet 的 DICE 提高了约 5%（相对值）。但在 50% 训练数据的情况下，使用迁移学习的 DenseNet 的 DICE 提高了 60%（相对值），而 ResNet 的 DICE 提高了 12%（相对值）。这清楚地表明，随着标注样本数量的减少，使用迁移学习的收益会更高。
图 8 显示了不同网络和初始化在不同 epoch 使用验证数据时的 DICE/MLA 指标。这有助于我们了解不同网络的收敛速度。从图中可以看出，TLU-Net 在 epoch 开始时的 DICE 值高于 U-Net。与 U-Net 相比，TLU-Net 的收敛 DICE 值更高。这清楚地表明，迁移学习有助于模型的快速收敛。对于 MLA 曲线也是如此。有趣的是，预训练模型的起始精度显着高于随机初始化。这主要是因为分类器直接使用编码器输出进行分类，并且编码器使用预训练权重进行初始化。这意味着预训练特征有助于区分不同的钢铁缺陷。
图 7 显示了使用 ResNet 的 TLU-Net 和使用 ResNet 的 U-Net 的四个示例掩码预测。在第一列中，TLU-Net 可以更准确地检测缺陷，而 U-Net 无法检测到缺陷的某些部分。在第二列中，TLU-net 由于一些光照差异而显示一些假阳性检测。在第三列中，两个网络都未能检测到某些缺陷。这主要是因为训练数据中此类缺陷的样本很少。我们假设 TLU-net 也失败了，因为缺陷形状更复杂，而预学习特征可能无法帮助这种情况。在第四列中，TLU-Net 可以检测到细线形式的缺陷，而 U-Net 无法检测到这些线。
图 9 显示了使用 ResNet/DenseNet 的 TLU-Net 和 U-Net 之间 DICE 差异的直方图。从图中可以看出，直方图偏向正值。我们观察到，在 ResNet 的情况下，81% 的图像得到了改进，而在 DenseNet 的情况下，63% 的图像得到了改进。
4. 结论
在这项工作中，我们提出使用迁移学习框架进行钢铁缺陷分类和分割。我们使用 U-Net 架构作为基础架构，并探索了两种编码器：ResNet 和 DenseNet。我们使用随机初始化和使用 ImageNet 数据集训练的预训练网络来比较这些网络的性能。我们发现，迁移学习在缺陷分割和分类方面的性能都更优越。我们还发现，