引用本文

岳忠牧, 张喆, 吕武, 赵瑞祥, 马杰. De-DDPM: 可控、可迁移的缺陷图像生成方法. 自动化学报, 2024, 50(8): 1539−1549 doi: 10.16383/j.aas.c230688

Yue Zhong-Mu, Zhang Zhe, Lv Wu, Zhao Rui-Xiang, Ma Jie. De-DDPM: A controllable and transferable defect image generation method. Acta Automatica Sinica, 2024, 50(8): 1539−1549 doi: 10.16383/j.aas.c230688

关键词

数据增强，数据集扩充，缺陷图像生成，深度学习 

摘要

基于深度学习的表面缺陷检测技术是工业上的一项重要应用, 而缺陷图像数据集质量对缺陷检测性能有重要影响. 为解决实际工业生产过程中缺陷样本获取成本高、缺陷数据量少的痛点, 提出了一种基于去噪扩散概率模型(Denoising diffusion probabilistic model, DDPM)的缺陷图像生成方法. 该方法在训练过程中加强了模型对缺陷部位和无缺陷背景的差异化学习. 在生成过程中通过缺陷控制模块对生成缺陷的类别、形态、显著性等特征进行精准控制, 通过背景融合模块, 能将缺陷在不同的无缺陷背景上进行迁移, 大大降低新背景上缺陷样本的获取难度. 实验验证了该模型的缺陷控制和缺陷迁移能力, 其生成结果能有效扩充训练数据集, 提升下游缺陷检测任务的准确率.

文章导读

近年来随着深度学习技术在计算机视觉领域的不断发展, 越来越多的行业开始使用该技术来解决一些譬如表面缺陷检测[1−3]、医学图像检测[4−6]、目标搜索识别[7−10]等的行业内问题. 深度学习摒弃了人工对图像的数据特征进行描述的过程, 转而使用大量的训练数据使模型自主学习到图像不同维度的特征, 进而完成对图像的感知. 但当训练数据较少时, 模型容易出现数据不均衡、过拟合等问题, 严重影响模型性能[11]. 

在工业缺陷检测领域, 缺陷数据集的收集与制作是不可规避的难题. 由于工业生产过程中良品率高, 表面有缺陷的样本数量总体偏少, 收集到的少量缺陷样本难以囊括足够的缺陷特征[12]. 此外相比于缺陷特征明显的样本, 形态较小、显著性较弱的临界缺陷在缺陷样本中占比更低[13], 且在人工判定时容易遗漏[14], 收集更为困难. 总而言之, 缺陷数据产生的不确定性, 缺陷形态、显著性强弱的不确定性, 加剧了缺陷样本的收集难度. 缺陷样本数量不足, 进一步严重制约了下游任务模型的工作性能. 

解决缺陷样本不足问题最直接的方法, 就是生成缺陷样本[12, 14]. 近年来, 以自回归模型、变分自编码器(Variational autoencoders, VAE)、生成对抗网络(Generative adversarial network, GAN)、去噪扩散概率模型(Denoising diffusion probabilistic model, DDPM)为代表的一系列方法在图像生成领域表现出优越的性能, 一些研究人员已经成功使用GAN等方法生成大量缺陷图像用以扩充训练数据[13, 15−19]. DDPM通过加噪和去噪进行图像生成的方式能够更好地覆盖样本分布, 其优异性能使之成为现阶段图像生成领域最先进的技术方案(State-of-the-art, SOTA)模型[20]. 但是将DDPM应用于缺陷样本生成时, 尚存在以下问题: 

1)工业生产过程中, 缺陷样本数量往往仅有几十个[12]. 而DDPM在该数量级的样本上训练后, 其生成的缺陷图像与训练集仅有像素级的微小差别, 模型发生过拟合[21]. 其生成结果显然无法满足生成多样性缺陷的需求, 无法对下游模型性能带来提升. 

2)缺陷图像中往往缺陷部位占比较小, 大部分都是无缺陷的背景区域. DDPM在训练过程中学习到大量背景的分布后, 其生成结果的背景也与训练集有相似的趋势. 工业生产中广泛存在具有相似缺陷, 但是背景不同的样本, 例如不同纹理的织物、瓷砖等. DDPM在单一背景的训练集上训练后, 没有将缺陷迁移至训练集以外背景的能力. 

基于上述问题, 本文提出了De-DDPM (Defect DDPM), 其简要流程如图1所示. De-DDPM在训练阶段通过由缺陷掩码x_mask和缺陷类别class控制的特征Unet网络, 加强对缺陷部位细粒度特征的学习. 生成阶段在特征Unet引导缺陷生成的基础上, 通过调节缺陷控制模块中的引导系数σ来控制缺陷显著性. 通过新背景y和背景融合模块的作用, 将生成的缺陷和不同背景进行融合. 因此De-DDPM模型有以下两个特点:

图 1  De-DDPM简要流程

1)缺陷特征可控. De-DDPM着重学习缺陷部位的细粒度特征, 减少了对训练数据量的需求. 模型通过缺陷控制模块中特征Unet的引导, 可以改变生成缺陷的类别、位置和形态, 通过调节缺陷语义信息引导系数σ, 可以控制生成缺陷的显著性. 模型具有生成多样性缺陷, 尤其是多样性临界缺陷的能力, 其生成结果极大丰富了训练数据集. 

2)缺陷可迁移. 在背景光滑、色彩单一的缺陷数据集进行简单的数据标注并完成训练后, 模型可以将譬如裂纹、擦痕、斑块等缺陷泛化至其他具有复杂背景的无缺陷样本中, 大大降低复杂背景上缺陷数据的获取难度. 

与DDPM相比, De-DDPM模型引入更多的信息来加强对生成缺陷的控制. 值得注意的是, 训练阶段De-DDPM用x_mask区分缺陷和背景, 进而对缺陷和背景的数据分布进行差异化学习. 相比缺陷数据本身, 新的缺陷掩码更容易获取. 通过人工绘制、掩码中加入扰动、训练简单的网络进行生成等方式, 可以得到大量形态各异的缺陷掩码. 通过新掩码的引导, De-DDPM模型能够生成更加丰富的缺陷图像, 进而提升缺陷的多样性. 

本文第1节介绍相关工作; 第2节详细阐述De-DDPM的整体结构; 第3节通过大量实验, 验证De-DDPM的有效性; 第4节对本文进行总结.

图 2  DDPM简要流程

图 3  特征Unet网络结构

本文提出了一种新的缺陷生成方法De-DDPM. 所提模型在训练阶段通过特征Unet差异化学习图像中缺陷部分和无缺陷部分的数据分布; 生成阶段通过缺陷控制模块, 模型能对生成缺陷的形态、位置、显著性等特征进行灵活控制, 通过背景融合模块, 模型能将缺陷泛化到不同风格的背景上. 实验表明, De-DDPM可以在大量无缺陷样本上生成逼真、可控的缺陷, 能够进一步显著提高下游缺陷检测网络的性能. 对不同于训练集的样本背景, 模型成功完成缺陷迁移, 效果良好.