CN117649351A - 一种基于扩散模型的工业缺陷图像仿真方法及装置 - Google Patents

Abstract

本发明提出了一种基于扩散模型的工业缺陷图像仿真方法及装置。使用带标签的缺陷图像训练构建的条件扩散模型得到预训练缺陷模型G1；用G1进行采样，生成不同种类的局部缺陷图；将局部缺陷与真实背景图拼接，得到缺陷前后景拼接图；利用缺陷背景和真实背景训练用于特征提取的扩散模型，得到预训练背景模型G2；用G2提取拼接图的特征，将不同时间步长的特征输入生成对抗网络对网络进行训练；输入拼接图，利用训练好的生成对抗网络生成缺陷图像。创新性地结合了扩散模型在不同时间步长中提取的特征层，并使用生成对抗网络，同时结合了扩散模型提取的特征，解决了生成算法训练速度慢、不易训练以及采样速度过低的问题。

Description

一种基于扩散模型的工业缺陷图像仿真方法及装置

技术领域

本发明涉及人工智能领域，具体涉及一种基于扩散模型的工业缺陷图像仿真方法及装置。

背景技术

工业缺陷图像仿真方法是一种通过模拟或生成工业制造过程中可能出现的缺陷图像的技术，可以用于包括测试和验证缺陷检测算法、培训人工智能模型、优化生产过程以减少缺陷率等多种用途。基于上述作用，工业缺陷图像仿真方法在快速原型开发、算法验证测试、产品改进开发等领域有着重要意义和应用价值。

基于深度学习的图像缺陷检测方法的一个共同点是，它们通常需要大量的训练样本，而这在实际环境中是难以获取的。并且一般来说，缺陷图像的背景是复杂的，并不简洁单一，而缺陷相对于整体来说又是细微的，如何不突兀地有效融合缺陷和背景是现有缺陷样本生成亟需解决的问题。目前关于图像生成，生成对抗网络和去噪扩散概率模型都能够生成高质量的样本，极大地提高了生成图像的质量和多样性。但目前的图像生成模型存在以下问题。

生成对抗网络能够有效捕捉数据分布，并对高维、复杂的真实数据进行建模，经过训练的网络可以直接从潜空间采样出逼真的数据样本，理论上可以收敛到全局最优，但是在实际训练过程中，往往难以得到一个优质的生成模型，更多时候模型只能产生没有实际价值的噪声图像。这是因为生成对抗网络存在模型难以收敛模式崩溃、模型不可控的问题，如若没有合适的超参数和正则化子，网络容易崩溃。

扩散模型通过正向加噪和逆向去噪可以使训练变得稳定易行，同时基于数据分布的采样方式也注定了扩散模型生成图像的多样性更加丰富。但是扩散模型采样速度慢，需要经过一个长的时间步长才能加噪到高斯噪声，因此在采用过程中同样需要多次的迭代才能完成采样。其次，扩散模型数据类型单一，且无法降维。

发明内容

本发明提出了一种基于扩散模型特征的两阶段缺陷图像生成算法，创新性地结合了对抗网络和扩散模型，解决了训练速度慢、不易训练以及采样速度过低的问题。

本发明所设计的基于扩散模型的工业缺陷图像仿真方法包含以下步骤：

使用带标签的缺陷图像训练构建的条件扩散模型得到预训练缺陷模型G1；

用G1进行采样，生成不同种类的局部缺陷图；

将局部缺陷与真实背景图拼接，得到缺陷前后景拼接图；

利用缺陷背景和真实背景训练用于特征提取的扩散模型，得到预训练背景模型G2；

用G2提取拼接图的特征，将不同时间步长的特征输入生成对抗网络对网络进行训练；

输入拼接图，利用训练好的生成对抗网络生成缺陷图像。

进一步地，条件扩散模型包括过程对称的下采样模块和上采样模块，每个下采样模块包含两个残差连接模块，在进行上采样时，还需要将上一层的输出与对应的下采样模块的输出求和后再输入该上采样模块，同时上采样层使用最近邻插值法完成。

更进一步地，条件扩散模型在图像下采样到最小时添加自注意力模块。

进一步地，缺陷模型G1输入经过随机时间步长加噪，在训练扩散模型时，依照去噪扩散理论，对每次去噪步骤中的噪声进行预测，并且使用均方差损失函数进行优化，如下所示：

其中，代表噪声；/>是神经网络结构，用于预测从/>到/>的噪声。

进一步地，用于特征提取的扩散模型包括扩散模型网络、图像特征编码网络以及融合网络，其中，扩散模型使用带时间编码输入的U-Net作为基础网络，图像特征编码网络对拼接的图像直接进行特征编码，融合网络通过将扩散模型以及图像特征编码网络提取得到的特征进行融合。

进一步地，生成对抗网络采用WGAN-GP的思想进行定义和训练，判别器采用PatchGAN架构。

优选地，在模型训练过程中，首先使用原始缺陷数据集对扩散模型进行预训练，在扩散模型训练完成之后，固定预训练好的扩散模型权重参数用于特征提取；

抽取扩散模型两次不同去噪步骤和/>的中间模块输出/>和/>以及上采样过程中的每一个残差模块输出/>和/>输入融合网络，其中/>表示上采样过程中的第i个残差模块的输出，同时选取图像编码网络的输出/>以及每个下采样过程中的残差模块输出/>，其中/>表示下采样过程中的倒数第i个残差模块的输出，因此融合网络的各个残差模块的输入如下：

其中代表concat拼接操作,/>代表残差模块内部的各项运算的组合，/>代表第/>个残差模块的输出。

更进一步地，在训练生成对抗网络部分时，使用生成器之前需要人为将缺陷前景x裁剪拼接至缺陷背景y中形成未融合图z，再对拼接图z进行加噪至第和/>的时间步长分别输入扩散模型中用于两个时间步骤的特征提取，计算公式如下：

其中，代表抽取的第/>个时间步长，/>代表抽取的第/>个时间步长，m代表人为选择的前后背景掩码图，/>代表点乘操作，/>为随机高斯噪声，/>表明了任意时间步长的噪声方差。生成器利用扩散模型提取的特征和原始拼接图片提取的特征完成进一步的学习，最终得到融合图片。使用GAN的方法进行训练，生成器将产生更加优质的效果。该方法使用GAN的方式进行训练以及采样，所以网络的采样速度与寻常GAN网络相比保持着竞争力。

基于同一发明构思，本发明还设计了一种电子设备，包括：

一个或多个处理器；

存储装置，用于存储一个或多个程序；

当一个或多个程序被所述一个或多个处理器执行，使得所述一个或多个处理器实现基于脑电解译的疲劳驾驶检测方法。

基于同一发明构思，本发明还设计了一种计算机可读介质，其上存储有计算机程序，所述程序被处理器执行时实现基于脑电解译的疲劳驾驶检测方法。

本发明的优点：

1.本发明创新性地提出了一种基于扩散模型特征的两阶段缺陷图像生成网络框架。在本模型中，第一阶段通过条件控制生成各类缺陷的局部图像，为后续奠定缺陷样本输入基础；第二阶段则将生成的缺陷前景与已有的缺陷背景拼接，独特地设计了不同时间步长的特征提取融合模块，并且创新地结合了扩散模型和生成对抗网络，通过融合算法可以得到真实、肉眼难以辨别的完整缺陷图。本发明解决了工业样本紧缺、图像生成背景复杂、采样速度慢或生成图像质量较低的问题。该模型即使在复杂背景下也能产生高质量的缺陷图像，亦能捕捉微小的缺陷区域。

2.本发明创新性地设计了不同时间步长的特征提取融合模块。该模块使得扩散模型中不同时刻t中所蕴含的图像特征级别可以结合，即在语义特征上能得到补充，实验证明不同时刻t提取得到的特征层进行结合可以更有效的作用于下游任务。

3.本发明创新性地结合了扩散模型和生成对抗网络。考虑到直接利用扩散模型进行融合将面临采样速度慢，模型难以修改等问题，该方法结合了生成对抗网络，旨在利用生成对抗网络的采样速度优势。通过利用预训练的扩散模型进行特征提取，再将提取的特征输入生成对抗网络的生成器中对缺陷的前后景进行充分融合。同时由于结合了扩散模型提取的特征，该网络可以生成更高质量的图像。

附图说明

图1是本发明的流程结构示意图。

图2是本发明的基于扩散模型特征的两阶段缺陷前后景融合网络整体框架图。

图3是本发明的采样层、中间层和上采样层结构示意图。

图4是本发明的残差模块架构图。

图5是本发明的自注意力模块示意图。

图6是本发明的判别器网络架构示意图。

具体实施方式

为了便于本领域普通技术人员理解和实施本发明，下面结合附图及实施例对本发明作进一步的详细描述，应当理解，此处所描述的实施示例仅用于说明和解释本发明，并不用于限定本发明。

实施例一

如附图1所示，本发明设计了一种基于扩散模型特征的两阶段缺陷图像生成方法，在第一阶段中使用条件控制生成各种类别的缺陷局部图像：

首先输入带标签的缺陷图像作为初始数据，经过随机时间步长加噪，训练得到预训练缺陷模型G1；

然后输入高斯分布噪声图，用G1进行采样，生成去噪不同种类的局部缺陷图；

最后将局部缺陷与真实背景图拼接，得到缺陷前后景拼接图。

在第二阶段中将生成的缺陷前景与已有的缺陷背景进行裁剪拼接，通过融合算法后即可得到真实的、肉眼难以区分的完整缺陷图片：

输入有缺陷背景图和真实背景图作为初始数据，经过随机时间步长加噪，对一个新扩散模型进行训练，得到预训练背景模型G2

输入拼接图作为初始数据，经过随机时间步长加噪，用G2提取拼接图像的特征，选取不同时间步长的特征输入生成对抗网络中；

输入第一阶段生成的拼接图，用生成对抗网络进行前后景的融合，得到融合后的缺陷图像，如附图2所示。

本发明提供的方法具体操作过程包含以下步骤：

步骤S1，构建条件扩散模型，即包含上采样及下采样层的扩散框架模型。采样层具体结构如附图3所示。使用带标签的缺陷图像训练构建的条件扩散模型得到预训练缺陷模型G1；使用扩散模型作为一阶段缺陷生成的基础模型，使用G1进行采样，并且通过添加条件控制使得模型在去噪恢复时产生特定种类的缺陷样本，用于生成新的缺陷前景图。具体过程如下：

S1-1，给定带类别标签的大小为3x64x64的缺陷图像作为U-Net的输入，设置去噪时间步长为1000，首先通过一个3x3卷积层将通道数转换为128，之后每次下采样通道数翻倍。在每个下采样模块中使用两次残差连接模块，最后连接一个卷积核大小为3，卷积步长为2的卷积层作为下采样层，一共下采样4次。

S1-2，上采样过程和下采样对称，只不过在进行上采样时，需要将上一层的输出与对应的下采样模块的输出求和后在输入该上采样模块，同时上采样层使用最近邻插值法完成，每个上采样模块同样使用两个残差模块，残差模块如附图4所示。为了减少内存占用，本发明只在图像下采样到最小时添加了自注意力模块，即在中间模块的第一个残差模块后面使用了自注意力机制，自注意力模块如附图5所示。原本的U-net网络仅仅包含上采样和下采样层对图像提取特征，而本方法在扩散模型的U-Net网络通道特征变换的过程中，添加了残差连接以及自注意力机制，用于进一步提升模型的效果。

S1-3，残差模块接受输入特征x，使用组归一化和swish激活函数进行激活后通过一个卷积核大小为3，步长为1的卷积层，然后将其与经过编码的时间信号和类别信息相加，其中时间信号和类别信号通过全连接层转换为相应通道数大小并广播到相应维度，组归一化层将特征按通道分为32组分别进行归一化。将相加得到的输出再次经过组归一化和swish激活层以及一个3x3卷积层，并与原始特征x进行求和，实现跳层链接，最后输入自注意力模块。

S1-4，注意力模块旨在对特征图内部进行注意力机制，对于输入大小为B×C×H×W大小的特征图，模型使用三个1×1卷积层得到同样大小的q、k、v，再将q变换为B×(H×W)×C，k变换为B×C×(H×W)大小。通过矩阵乘法得到B×C×(H×W) ×(H×W)的注意力权重图，它表示了特征图中每个坐标的特征向量与其他坐标特征向量的相似度。通过softmax之后与v进行矩阵乘法，得到的结果再经过一次1x1的卷积层便得到了经过自注意力机制后的输出特征图。

步骤S2，同步采集缺陷图像数据集，利用数据集对步骤S1构建的深度学习网络进行训练，得到训练好的扩散模型。在训练扩散模型时，依照去噪扩散理论，对每次去噪步骤中的噪声进行预测，并且使用均方差损失函数进行优化：

步骤S3，使用步骤S2中训练好的扩散模型生成不同种类的局部缺陷图，即缺陷前景图。

步骤S4，构建用于特征提取的扩散模型，用于对缺陷特征进行识别提取，利用缺陷背景和真实背景训练用于特征提取的扩散模型，得到预训练背景模型G2。使用带时间编码输入的U-Net作为基础网络。在步骤S4及之后，通过将一阶段生成的缺陷图与工业采集的真实背景图进行融合，从而得到完整的缺陷图像。该阶段模型整体分为三个部分：扩散模型网络DDPM、图像特征编码网络Base以及融合网络Fusion，图像特征编码网络即下采样网络提取图像特征，融合网络即加入缺陷特征后的上采样网络。在扩散模型网络中，本发明使用带时间编码输入的U-Net作为基础网络。

步骤S5，同步采集缺陷背景图像数据集和无缺陷背景图像数据集，利用数据集对步骤S4构建的深度学习网络进行训练，得到训练好的扩散模型。

步骤S6，将一阶段生成的缺陷图与工业采集的真实背景图进行融合，从而得到完整的缺陷图像，将生成的缺陷前景图与无缺陷背景图进行拼接，得到拼接数据集。

步骤S7，选取步骤S6中的拼接数据，输入S5中的用于特征提取的扩散模型，提取不同时间步长的特征，并且使用基础网络提取图像特征。融合网络通过将扩散模型以及图像特征编码网络提取得到的特征进行融合，并最终转换为融合图像。

步骤S8，扩散模型网络DDPM对于不同特征需要经过预训练，之后在生成对抗网络训练过程中将其参数固定，每次同时抽取两个不同时间步长t1和t2的特征输入融合网络中。将步骤S7中的不同时间步长的特征以及使用基础网络提取的图像特征，输入融合网络中，并使用梯度下降算法进行训练更新，具体步骤如下：

S8-1，在模型训练过程中，首先使用原始缺陷数据集对扩散模型进行预训练，从而使U-Net能够精确的捕捉到原始数据的数据分布。在扩散模型训练完成之后，固定预训练好的扩散模型权重参数用于特征提取。

S8-2，抽取扩散模型两次不同去噪步骤和/>的中间模块输出/>和/>以及上采样过程中的每一个残差模块输出/>和/>输入融合网络，其中/>表示上采样过程中的第i个残差模块的输出。同时选取图像编码网络的输出/>以及每个下采样过程中的残差模块输出/>，其中/>表示下采样过程中的倒数第i个残差模块的输出。因此融合网络的各个残差模块的输入如下：

其中代表concat拼接操作,/>代表残差模块内部的各项运算的组合。具体而言，本发明选择抽取扩散模型的第10和第100个时间步骤的特征输入融合网络，即。然后对生成器和判别器进行交替训练。

S8-3，生成对抗网络采用WGAN-GP的思想进行定义和训练，其生成器架构与扩散模型的架构基本一致，不同点在于上采样过程中的每个残差模块的输入通道数都因为特征层数的拼接而进行了相应扩大。判别器采用了PatchGAN的架构，如附图6所示。首先通过一个卷积核大小为4x4，步长为2的卷积层进行一次下采样，同时调整通道数为64。然后再经过5次下采样，每次下采样同样使用卷积核大小为4x4，步长为2的卷积层完成，并将通道数翻倍，每次卷积后使用一个LeakyReLU激活函数进行激活。最后使用一个3x3的卷积层将输出转换为1x2x2大小的结果图。

在训练生成对抗网络部分时，使用生成器之前需要人为将缺陷前景x裁剪拼接至缺陷背景y中形成未融合图z，再对拼接图z进行加噪至第和/>的时间步长分别输入扩散模型中用于两个时间步骤的特征提取。计算公式如下：

其中，代表抽取的第/>个时间步长，/>代表抽取的第/>个时间步长，m代表人为选择的前后背景掩码图，/>代表点乘操作，/>为随机高斯噪声。生成器利用扩散模型提取的特征和原始拼接图片提取的特征完成进一步的学习，最终得到融合图片。

S8-4，在训练对抗生成网络时，因为使用了WGAN-GP的架构，因此对抗损失如下：

除此之外，本发明采用Liu等人的想法，额外添加了重建损失和风格损失。添加重建损失旨在希望融合的图片在背景纹理上和原始图像保持一致，但是生成的融合缺陷图像和原始图像保持完全一致是不合理的，因此通过添加松弛尺度将其改进为铰链重建损失，如下所示：

其中z代表缺陷拼接图，生成器，/>生成器，y代表真实缺陷背景图，/>为松弛尺度，控制不一致区域的强度，本发明实验取/>=0.01，/>代表梯度惩罚权重，本发明实验取/>=5。

对于风格损失，使用预训练的VGG-Net提取融合图像与原始图像的特征图，通过在特征层面对其施加惩罚，迫使生成器在融合缺陷的同时保持和原始图像一样的特征纹理和风格样式，具体如下：

本发明提取了VGG-Net下采样过程中的5层激活，因此i=1、2、3、4、5。其中表示从VGG-Net中提取的第i层特征激活，大小为/>。/>代表第i层的权重，分别赋予权重为/>。根据以上，生成对抗网络的整体损失函数如下：

其中和/>分别表示铰链重建损失和风格损失的权重，分别取值为1和3。

步骤S9，重复步骤S6至S8，直至网络收敛。

本发明构建并训练基于扩散模型特征的工业缺陷图像仿真方法，利用生成对抗网络进行快速采样，扩散模型生成缺陷样本，改善了生成对抗网络不易训练及扩散模型采样过慢的单一模型缺陷，提高了模型的语义表征能力。

实施例二

基于同一发明构思，本发明还提供了一种电子设备，包括一个或多个处理器；存储装置，用于存储一个或多个程序；当一个或多个程序被所述一个或多个处理器执行，使得所述一个或多个处理器实现实施例一中所述的方法。

由于本发明实施例二所介绍的设备为实施本发明实施例一基于扩散模型的工业缺陷图像仿真方法所采用的电子设备，故而基于本发明实施例一介绍的方法，本领域所属技术人员能够了解该电子设备的具体结构及变形，故而在此不再赘述。凡是本发明实施例一种方法所采用的电子设备都属于本发明所欲保护的范围。

实施例三

基于同一发明构思，本发明还提供了一种计算机可读介质，其上存储有计算机程序，所述程序被处理器执行时实现实施例一中所述的方法。

由于本发明实施例三所介绍的设备为实施本发明实施例一基于扩散模型的工业缺陷图像仿真方法采用的计算机可读介质，故而基于本发明实施例一介绍的方法，本领域所属技术人员能够了解该电子设备的具体结构及变形，故而在此不再赘述。凡是本发明实施例一种方法所采用的电子设备都属于本发明所欲保护的范围。

本文中所描述的具体实施例仅仅是对本发明精神作举例说明。本发明所属技术领域的技术人员可以对所描述的具体实施例做各种各样的修改或补充或采用类似的方式替代，但并不会偏离本发明的精神或者超越所附权利要求所定义的范围。

Claims (10)

1.一种基于扩散模型的工业缺陷图像仿真方法，其特征在于：

使用带标签的缺陷图像训练构建的条件扩散模型得到预训练缺陷模型G1；

用G1进行采样，生成不同种类的局部缺陷图；

将局部缺陷与真实背景图拼接，得到缺陷前后景拼接图；

利用缺陷背景和真实背景训练用于特征提取的扩散模型，得到预训练背景模型G2；

用G2提取拼接图的特征，将不同时间步长的特征输入生成对抗网络对网络进行训练；

输入拼接图，利用训练好的生成对抗网络生成缺陷图像。

2.根据权利要求1所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：条件扩散模型包括过程对称的下采样模块和上采样模块，每个下采样模块包含两个残差连接模块，在进行上采样时，还需要将上一层的输出与对应的下采样模块的输出求和后再输入该上采样模块，同时上采样层使用最近邻插值法完成。

3.根据权利要求2所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：条件扩散模型在图像下采样到最小时添加自注意力模块。

4.根据权利要求1所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：缺陷模型G1输入经过随机时间步长加噪，在训练扩散模型时，依照去噪扩散理论，对每次去噪步骤中的噪声进行预测，并且使用均方差损失函数进行优化，如下所示：

其中，代表噪声；/>是神经网络结构，用于预测从/>到/>的噪声。

5.根据权利要求1所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：用于特征提取的扩散模型包括扩散模型网络、图像特征编码网络以及融合网络，其中，扩散模型使用带时间编码输入的U-Net作为基础网络，图像特征编码网络对拼接的图像直接进行特征编码，融合网络通过将扩散模型以及图像特征编码网络提取得到的特征进行融合。

6.根据权利要求1所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：生成对抗网络采用WGAN-GP的思想进行定义和训练，判别器采用PatchGAN架构。

7.根据权利要求5所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：在模型训练过程中，首先使用原始缺陷数据集对扩散模型进行预训练，在扩散模型训练完成之后，固定预训练好的扩散模型权重参数用于特征提取；

抽取扩散模型两次不同去噪步骤和/>的中间模块输出/>和/>以及上采样过程中的每一个残差模块输出/>和/>输入融合网络，其中/>表示上采样过程中的第i个残差模块的输出，同时选取图像编码网络的输出/>以及每个下采样过程中的残差模块输出/>，其中表示下采样过程中的倒数第i个残差模块的输出，因此融合网络的各个残差模块的输入如下：

其中代表concat拼接操作, />代表残差模块内部的各项运算的组合，/>代表第/>个残差模块的输出。

8.根据权利要求7所述的基于扩散模型的工业缺陷图像仿真方法，其特征在于：在训练生成对抗网络部分时，使用生成器之前需要人为将缺陷前景x裁剪拼接至缺陷背景y中形成未融合图z，再对拼接图z进行加噪至第和/>的时间步长分别输入扩散模型中用于两个时间步骤的特征提取，计算公式如下：

其中，代表抽取的第/>个时间步长，/>代表抽取的第/>个时间步长，m代表人为选择的前后背景掩码图，/>代表点乘操作，/>为随机高斯噪声，/>表明了任意时间步长的噪声方差；生成器利用扩散模型提取的特征和原始拼接图片提取的特征完成进一步的学习，最终得到融合图片。

9.一种电子设备，其特征在于，包括：

一个或多个处理器；

存储装置，用于存储一个或多个程序；

当一个或多个程序被所述一个或多个处理器执行，使得所述一个或多个处理器实现如权利要求1-8中任一所述的方法。

10.一种计算机可读介质，其上存储有计算机程序，其特征在于：所述程序被处理器执行时实现如权利要求1-8中任一所述的方法。