基于渐进式特征注入和全局-局部判别的图像生成网络，用于提升铁路表面缺陷检测效果

《Pattern Recognition Letters》：Progressive feature injection and global-local discrimination-based image generation network for improving rail surface defect inspection

铁路表面缺陷检测数据稀缺，提出两步生成网络（RSD-GAN）解决：1）背景预训练利用充足的无缺陷样本；2）渐进式缺陷特征注入（PDFIB）结合联合判别器（全局-局部分支），提升生成图像质量（FID 63.91/58.39）与多样性（IS 2.2803/2.5525），并增强局部缺陷真实性。实验验证下游检测mIoU达0.7019（刮痕）/0.7375（裂纹），优于基线方法。

　　

作者：王德洲 | 苏世祥 | 陆晓波

东南大学自动化学院，南京，210096，中国

摘要

铁路表面缺陷检测是铁路系统中的重要任务。尽管现有方法在一定程度上取得了成功，但由于缺陷样本的稀缺，检测精度仍然有限。扩大训练数据集是解决这一问题的一个直观方法。然而，现有的图像生成方法难以同时生成高质量和多样化的样本。本文提出了一种两步铁路表面缺陷生成方法，包括背景特征学习和缺陷前景特征学习。具体来说，我们首先利用丰富的无缺陷样本预训练背景生成网络。然后，我们设计了一个渐进式缺陷特征注入分支，在生成的多样化掩码区域内学习缺陷特征，并将其注入无缺陷特征中。同时，我们开发了一个具有全局和局部分支的联合判别器，使网络能够关注整个图像和局部不规则缺陷区域，并引入了一个轻量级的匹配判别器以确保缺陷掩码的合理性。实验表明，我们生成的缺陷图像质量更高（划痕的FID最低为63.91，裂纹的FID最低为58.39），多样性也更丰富（划痕的IS最高为2.2803，裂纹的IS最高为2.5525）。此外，在下游检测任务中，我们的方法取得了最佳性能，划痕和裂纹缺陷的mIoU值分别为0.7019和0.7375。

引言

由于长期处于高维护压力下，铁轨通常会出现各种类型的表面缺陷，如划痕和裂纹。划痕是不规则且呈块状的缺陷，而裂纹则是铁轨侧面鱼鳞状剥落的痕迹。这些缺陷会影响铁路系统的安全性和稳定性，甚至可能导致严重的交通事故和经济损失[1]。因此，需要定期进行精确的铁路检测。

铁路图像数据集的特点是缺陷样本不足，而无缺陷样本较多，这在缺陷检测领域很常见。最近，基于深度学习的方法因其卓越的特征提取能力而成为缺陷检测的主流。然而，缺陷样本的稀缺带来了两个新问题：数据不平衡和模型泛化能力差[2]。一方面，缺陷样本和无缺陷样本的数量差异使得模型难以很好地学习缺陷样本的特征表示，从而降低了检测精度。另一方面，缺陷特征的缺乏阻碍了模型学习多样化的特征，导致模型过拟合。

为了解决上述问题，人们在三个方面做出了努力：小样本学习、无监督学习和数据增强。小样本学习方法旨在利用有限的数据获得良好的检测器[3]，通常通过任务之间的知识转移[4]、[5]或领域泛化框架[6]、[7]来实现，但由于铁轨表面特征的独特性，很难获得类似的领域数据集，当领域差异较大时往往失败。无监督学习方法仅使用无缺陷样本进行训练[8]、[9]；在测试时，通过将重建的无缺陷图像与输入进行比较来定位缺陷，但它们只能处理单一类型的缺陷，这不适合具有多种缺陷类型的铁轨表面。数据增强可以生成更多样本：传统方法（例如翻转、平移、旋转）使用简单的线性变换，不会引入新的特征分布，从而限制了精度的提升。最近，生成对抗网络（GANs）已经为铁路紧固件[10]、[11]、医学疾病[12]、[13]和其他一些工业缺陷[14]、[15]生成了缺陷图像，在一定程度上缓解了样本稀缺的问题。值得注意的是，数据驱动策略在工业场景中已被证明很有价值，例如增强焊膏缺陷检测[16]、优化轮胎挤出参数[17]和基于钢铁行业数据的决策[18]。然而，本研究中的铁路表面缺陷生成仍面临三个关键挑战。首先，过度依赖稀缺的缺陷样本。大多数现有的基于GAN的方法需要相对大量的真实缺陷样本，这与铁轨表面缺陷的稀有性相矛盾。其次，缺陷多样性不足且不可控。以往的研究通过对比学习或多重判别器[10]、[19]来提高多样性，但缺陷区域的多样性仍然不可控，并且是由网络被动决定的。第三，局部缺陷区域的真实性较差。现有方法关注整个图像，忽略了不规则的局部缺陷区域，导致细节模糊或与背景纹理不一致，从而降低了下游检测的实用性。

在本文中，为了解决铁路表面缺陷生成中的上述挑战，我们在样本有限的情况下开发了一种新颖的两步生成网络（RSD-GAN）。我们的贡献可以总结如下：

• •
  我们提出了一种两步训练范式，分别学习背景和缺陷前景特征，通过利用丰富的无缺陷样本进行背景学习，减少了对稀缺缺陷样本的依赖。设计了一个具有掩码约束损失的渐进式缺陷特征注入分支（PDFIB），以学习和注入缺陷特征到无缺陷特征中，提高了生成缺陷的多样性和质量。
• •
  我们开发了一个联合判别器，同时考虑整个图像和不规则的局部缺陷区域，提高了局部缺陷区域的真实性。
• •
  广泛的实验表明，生成的缺陷图像质量高且多样性丰富，提高了检测模型的性能。

小样本图像生成

小样本图像生成旨在利用有限的数据生成高质量的图像。这是一个具有挑战性的任务，但对于数据集稀少或难以收集的情况更为实用。最近的小样本图像生成方法可以根据是否进行预训练分为两类。基于预训练的方法通过利用具有足够图像的类似领域的知识来提高性能[20]、[21]、[22]。

提出的方法

我们采用两步训练范式，将生成过程分为背景特征学习和缺陷前景特征学习，以在样本有限的情况下平衡多样性和真实性。如图1所示，该框架包括一个生成器（集成背景生成网络和渐进式缺陷特征注入分支PDFIB）以及两个判别器。首先，背景特征学习利用丰富的无缺陷样本来训练背景生成

数据集和实验设置

共获得了2395张无缺陷图像、148张划痕图像和100张裂纹图像。无缺陷图像用于预训练背景生成网络。此外，59张划痕图像和40张裂纹图像被用于训练PDFIB及其对应的判别器。剩余的缺陷图像用于测试不同方法的性能。

实验在Ubuntu 20.04操作系统上使用PyTorch 1.7.1、Intel i9-11900KF CPU和NVIDIA RTX 3060 GPU上进行。

结论

在本文中，我们提出了一种结合渐进式缺陷特征注入和全局-局部判别的两步铁路表面缺陷生成网络。这种两步范式确保了在样本有限的情况下背景和缺陷前景的多样性和真实性。具体来说，PDFIB学习丰富的、多样的缺陷特征；DFM捕获上下文相关的缺陷信息，并使网络关注局部缺陷区域；ToMask模块结合匹配判别器和掩码

CRediT作者贡献声明

王德洲：撰写——原始草稿、可视化、验证、软件、方法论、数据整理。苏世祥：撰写——审阅与编辑、方法论、形式分析、数据整理。陆晓波：撰写——审阅与编辑、监督、资金获取、概念化。

利益冲突声明

作者声明他们没有已知的可能影响本文工作的竞争性财务利益或个人关系。

致谢

本工作部分得到了国家自然科学基金（项目编号：62271143）的支持，部分得到了江苏省前沿技术研发计划（项目编号：BF2024060）的支持，还得到了东南大学大数据计算中心的支持。