CN115375604A - 一种基于量子化自编码器的无监督缺陷检测方法 - Google Patents

Abstract

一种基于量子化自编码器的无监督缺陷检测方法，首先对输入图像进行预处理，将预处理后的图像输入编码网络中，映射得到中间特征图，并加上位置编码信息；将中间特征图中每个向量与编码表中的向量进行最邻近匹配，用最接近的向量对其进行替换，得到量化的特征图；然后，将量化后的特征图输入到解码网络中，解码网络将特征图重新恢复成图片；最后对于原图和重建图片，一方面逐像素计算均方误差损失，另一方面分块计算SSIM损失，综合两个指标来判别和定位缺陷。本发明方法利用VQ‑VAE的编码表机制，以及提出的向量三元组损失，抑制了模型的泛化能力，提高缺陷的检出效果，同时通过MSE和SSIM两种指标的结合，提升了对缺陷的定位能力。

Description

一种基于量子化自编码器的无监督缺陷检测方法

技术领域

本发明属于机器视觉技术领域，通过量子化自编码器VQ-VAE将图像映射到特征空间再进行重建，利用独特的编码表强化对缺陷的检测能力，为一种基于量子化自编码器的无监督缺陷检测方法。

背景技术

缺陷检测是工业生产中非常重要的一环，基于机器视觉的缺陷检测技术可以极大的提升工业生产的效率。近年来随着深度学习技术在计算机视觉领域的发展，卷积神经网络在诸多的图像任务上都取得了显著的效果，然而这些方法绝大多数都是需要大量标注数据的有监督学习方法。在实际的工业场景中，缺陷样本往往是难以获取的，导致能够用于有监督学习的数据很少，而且标注数据本身的成本也是巨大的。无监督的方法可以只基于能够大量获取的无缺陷样本进行训练，例如使用生成对抗网络或者自编码器来对图像进行重建，但由于只使用了正样本来进行训练，当模型接收到缺陷样本时，就无法很好的对缺陷区域进行重构。无监督学习的方法往往在效果上与有监督学习的方法有着一定的差距，而且由于模型泛化能力过强，即使只使用正样本训练，模型常常仍能够重建出缺陷，影响检测精度。

发明内容

本发明要解决的问题是：基于有监督学习的方法需要大量的缺陷样本，这是现实的工业场景中难以获取到的。而基于无监督学习的方法在效果上往往难以达到有监督学习的效果，其中很大一部分原因是模型泛化能力过强，导致缺陷区域也能被重建出来，从而无法与原图进行对比来判别缺陷。而且，对于缺陷的定位，如果只简单的将原图和重建图像逐像素地进行差分，往往效果并不好。因此，需要一种能够抑制缺陷重构能力的方法来提升缺陷检测效果，并且需要提出更有效的定位缺陷的方法。

本发明的技术方案为：一种基于量子化自编码器的无监督缺陷检测方法，构建一个量子化自编码器用于缺陷检测，量子化自编码器包括编码网络、编码表和解码网络，

量子化自编码器对输入图像进行归一化预处理，将预处理后的图像输入编码网络中，映射得到中间特征图，并加上位置编码信息；然后将中间特征图中每个向量与编码表中的向量进行最邻近匹配，用最接近的向量对其进行替换，得到量化的特征图；将量化后的特征图输入到解码网络中，解码网络将特征图重新恢复成图片；使用无缺陷的训练样本图像训练量子化自编码器；

缺陷检测时，将待检测图像输入训练好的量子化自编码器，对于原图和量子化自编码器重建的图片，一方面逐像素计算均方误差损失，另一方面分块计算SSIM损失，综合两个指标来判别和定位缺陷。

进一步的，量子化自编码器具体为：

1)编码网络encoder：

输入编码网络的图像映射到一个潜在空间中，得到连续的编码向量z′_e：

z′_e＝encoder(x)

在编码器中加入位置编码position encoding，用位置信息来对特征进行约束，采用Transformer中的提出的Sinusoidal Position Encoding：

编码向量z′_e与位置编码相加共同组成最终的特征向量z_e；

2)构建编码表和向量量化：

构建一个容量为K的Embedding层，其中每个向量的维度均为64，将其随机初始化：

E＝[e₁，e₂，e₃...e_K]

对于1)中得到的特征向量z_e，通过最邻近搜索匹配这K个向量中最相似的向量，记为z_q，使连续的编码向量实现离散化；

3)解码网络decoder：

将得到的编码向量z_q输入到解码网络中，解码网络对编码向量进行上采样并重建原图：

x′＝decoder(z_q)

其中，使用直通估计的方法，在前向传播时使用z_q进行计算，在反向传播时使用z_e计算梯度：

decoder(z_q)→decoder(z_e+sg[z_q-z_e])

sg表示不计算梯度，这样在前向传播时该式等价于decoder(z_q)，在反向传播时由于z_q-z_e不提供梯度，该式等价于decoder(z_e)。

采取本发明方法，与现有技术相比本发明的有益效果如下：

(1)本发明对现有的基于自编码器或对抗生成网络的无监督缺陷检测方法进行分析，提出了基于VQ-VAE的缺陷检测方法。利用VQ-VAE的编码表机制，解码器的输入全部来自于编码表的向量，而这些向量都是由正常图片学习而来的，异常区域经过编码器得到的向量被其替换掉，与之前的自编码器或对抗生成网络方案相比，异常区域重建后的差别会加大，从而提高对缺陷的检出效果。本发明在VQ-VAE的基础上结合位置编码、向量三元组损失来优化编码表向量，能够更准确重建图像。

(2)为了抑制网络的泛化能力，本发明提出在编码器中加入位置编码用位置信息来对特征进行约束。原本编码表中的向量是没有位置约束的，只要特征足够接近，来自某一区域的编码表向量也可以用于重构其他区域。本发明加入位置编码后对编码表向量增加了新的约束，我们期望在编码表中检索时，输入数据的编码向量和在编码表检索到的向量不仅在特征空间中足够接近，且在空间位置上也尽可能的接近。

(3)同时，为了抑制模型的泛化能力，防止模型重建能力过强导致异常区域也被重建，本发明还提出了向量三元组损失，进一步提高编码表中向量的区分度，也强迫编码器生成的向量更接近正常样本的向量，从而提高对缺陷的检出效果，同时也压缩了编码表的大小。与现有技术中的三元组损失Triplet-loss有所不同，Triplet-loss是输入一批样本，对于其中的一个样本要和一个同类样本组成正样本对，和一个不同类的样本组成负样本对，实现输入样本与同类样本更接近，不同类的样本则远离，主要用于分类；而本发明的向量三元组损失的两个向量都来自编码表，和训练出来的具有代表性的向量计算损失，也并不考虑同类不同类的概念，只是和编码表中第一近、第二近的向量计算，让编码表中的向量变得“稀疏”，进而抑制模型的泛化能力。

(4)为了更精准的定位缺陷，本发明在逐像素计算MSE的基础上，提出分块计算SSIM的方法，结合两个标准，能够更加精确的定位到缺陷具体的位置。进一步的，还提出结合空间注意力的结构相似性损失SAM-SSIM，SAM-SSIM在SSIM的基础上通过空间注意力来给所有图像块分配权重，让网络动态的调节不同图像块的重要程度，使得VQ-VAE能更准确地重建细节和边缘，也能够更加精确的定位到缺陷具体的位置。

附图说明

图1为本发明的基于量子化自编码器的无监督缺陷检测方法流程图。

图2为DAGM表面缺陷检测数据集检测样例图。图2a)为输入原图，圈出的区域为缺陷；图2b)为使用变分自编码器VAE重构后的图片，可见缺陷也被重构出来了，无法达到缺陷检测的效果；图2c)为本发明中使用VQ-VAE重构的图像，缺陷没有被重构出来，将原图和重构图像比对后，即可检测到缺陷。由此可见以往基于自编码器的方案泛化能力过强，将缺陷区域也重构出来，而本发明提出的基于量子化自编码的无监督缺陷检测方法能够抑制模型的泛化能力，提升检出效果。

图3为使用工业产品数据集对本发明方法进行测试，从左到右依次为(a)原图、(b)重构图像、(c)缺陷置信度可视化图像。

具体实施方式

针对现有技术，本发明提出一种无监督缺陷检测方法，只利用无缺陷样本进行训练，通过量子化自编码器VQ-VAE能够有效抑制模型的泛化能力，让模型只能重建出正常样本，无法很好的重建出缺陷区域，从而实现对缺陷的判别和定位，具有重要的实际应用价值。

本发明一种基于量子化自编码器的无监督缺陷检测方法，具体实施如下：

1、训练图像预处理：

批量读入无缺陷的训练样本图像，缩放到统一的尺寸512*512，并进行数据增强。对增强后的输入图像按以下方式进行归一化：

其中mean和std分别为训练集图像像素值的均值和方差，由全体训练集图像统计得到。

2、编码网络encoder：

将预处理后的图像输入到编码网络，映射到一个潜在空间中，得到连续的编码向量z′_e：

z′_e＝encoder(x)#(1)

在编码器中加入位置编码position encoding，用位置信息来对特征进行约束，编码向量z′_e与位置编码相加共同组成最终的特征向量z_e；此处的z_e是一种简略的写法。z_e实际上是由顶层特征图或底层特征图卷积得到的m×m×64的多维向量。

编码网络是由底层编码器和顶层编码器组成，而编码器是由若干卷积层、BN层和ReLU层组成的残差模块堆叠而成。

图像输入到编码网络后首先经过底层编码器被下采样四倍，得到底部特征图。然后通过顶层编码器进一步下采样两倍，得到顶部特征图。底部特征图长宽为输入图像的四分之一，保留了较多的局部细节，有助于重建出更清晰的图像。顶部特征图长宽为输入图像的八分之一，包含更多的特征语义信息。顶层特征图经过卷积层将通道数压缩为64，之后通过在向量字典中进行最邻近搜索来实现量化。将顶层特征图上采样两倍后与底层特征图合并，同样经过卷积层将通道数压缩为64然后进行量化。

为了进一步抑制网络的泛化能力，本发明提出在编码器中加入位置编码(position encoding)，用位置信息来对特征进行约束。本发明采用Transformer中的提出的Sinusoidal Position Encoding：

编码向量z′_e与位置编码相加共同组成最终的特征向量z_e，也就是最终的特征图。

3、构建编码表和向量量化：

构建一个容量为K的Embedding层，其中每个向量的维度均为64，将其随机初始化：

E＝[e₁，e₂，e₃...e_K]

对于2中得到的向量z_e，通过最邻近搜索匹配这K个向量中最相似的向量，记为z_q，连续的编码向量也就实现了离散化。

具体的，z_e中每个位置的向量都与向量词典(e_k，k∈1...K)中的所有向量计算余弦距离，并与向量词典中最接近的向量的索引进行替换，从而实现离散化。

4、解码网络decoder：

将3中得到的编码向量z_q输入到解码网络中，解码网络对编码向量进行上采样并重建原图：

x′＝decoder(z_q)#(5)

由于量化的过程是不能产生梯度的，为了能够进行反向传播和梯度回传，使用直通估计(Straight-Through Estimator)的方法，即在前向传播时使用z_q进行计算，在反向传播时使用z_e计算梯度。

decoder(z_q)→decoder(z_e+sg[z_q-z_e])

sg表示不计算梯度(stop gradient)，这样在前向传播时该式等价于decoder(z_q)，在反向传播时由于z_q-z_e不提供梯度，该式等价于decoder(z_e)。

5、训练：

编码网络、编码表和解码网络共同构成了量子化自编码器(VQ-VAE)。使用无缺陷的正常样本训练网络。损失函数L由重构损失(reconstruction loss)、编码表损失(codebook loss)和向量三元组损失(embedding triplet loss)三部分构成：

L＝L_r+L_c+L_e#(6)

具体的，重构损失L_r又由均方差损失MSE和SSIM损失组成：

L_r＝MSE+SSIM#(7)

MSE即重建损失函数，也是绝大多数图像重建算法里最常见的损失函数：

MSE(X，Y)＝∑||X_i-Y_i||₂#(8)

X、Y分别为输入原图和重建图像。

但使用MSE存在两大问题：①当原图和重构图像的像素值大小保持大致一致时，即异常区域和正常区域颜色相近但结构不同时，MSE会难以区分；②MSE损失对于轻微的边缘小误差也会有强响应，会干扰真正的异常定位。而基于SSIM的方法可以有效的避免以上两个问题，取得比MSE更好的效果。不过由于SSIM不能对全图使用，需要将原图和重构图像切割为若干个图像块后进行计算。

SSIM从三个维度量比较两个图像块x、y的相似性：亮度l(x，y)，对比度c(x，y)，结构相似度s(x，y)。最终x和y的相似度为这三者的函数：

SSIM＝l(x，y)+c(x，y)+s(x，y)#(9)

亮度通过比较色块的像素平均值μ_x、μ_y来计算，对比度通过比较色块的像素值方差σ_x、σ_y来计算，结构相似度考虑了两个色块之间的协方差σ_xy：

令C₃＝C₂/2，整理后可得到SSIM的公式：

除了MSE直接与SSIM两个指标结合应用来判别和定位缺陷，本发明针对SSIM进一步提出了基于空间注意力的结构相似性损失SAM-SSLM，传统SSIM中N×N个图像块的重要程度是相同，注意力相同，即每个图像块的权重都是1，本发明进行了改进，通过提取空间注意力，为每个图像块分配不同的权重。为了提取出每个图像块对应的空间注意力，本发明在VQ-VAE上新增了空间注意力模块SAM。该模块由一系列的反卷积和激活函数组成，将z_q映射到大小为N×N×1的特征图W，对应每个图像块的空间注意力：

W＝SAM(z_q)#(14)

将空间注意力与对应图像块的SSIM相乘，得到SAM-SSIM：

其中W_i为第i个图像块对应的空间注意力。在逐像素计算MSE的基础上，提出结合空间注意力的结构相似性损失SAM-SSIM，VQ-VAE结合这两个标准重建图像，能够更准确地重建细节和边缘。

编码表损失L_c目的是让编码表中的向量与图像经过编码器后的向量尽可能接近，本发明中使用L2损失函数：

相应的，也需要让图像经过编码器后的中间向量尽可能接近编码表中的向量，同样的使用L2损失函数，得到编码表损失L′_e：：

为了降低模型的泛化能力，降低异常被重建的几率，本发明在以上基础上引入了向量三元组损失，对L′_e进行改造得到向量三元组损失L_e，对于特征图上的每一个向量z_ek，都在编码表中找到最接近它的向量z_p，和第二接近它的向量zn，三个向量组成三元组共同计算损失，使得z_ek和z_p尽可能接近，z_ek和z_n尽可能远离，从而让编码表中的向量变得“稀疏”：

6、缺陷检测

缺陷检测：根据训练后的VQ-VAE模型，实现工业图像的缺陷检测。将待检测的图片经过和训练过程一样的预处理后输入到量子化自编码器中，得到重建图像。由于网络只经过了正常样本的训练，且解码器获得的输入向量全部来自于编码表，而编码表仅有正常样本训练而来，实际上是缺乏关于缺陷的信息的，这样解码器对于异常和缺陷就无法很好的重构，通过对比原图和重构图像的差异可对缺陷进行判别和定位。

具体的，对比原图和重建图像的差异时，也会和5中一样计算两张图的MSE和SSIM，或MSE和SAM-SSIM，先逐像素计算两张图的MSE损失，记为S_mse：

S_mse＝||x-y||₂

再将原图和重建图都切分成11*11的图像块，计算SSIM，记为S_ssim：

S_ssim＝l(x，y)+c(x，y)+s(x，y)

如计算SAM_SSIM，记为

综合考虑这两个指标，得到加权后的缺陷置信度：

S＝S_ssim+S_mse

或

S＝S_ssim+S_{sam_ssim}

当缺陷置信度大于指定的阈值时，此处便被判别为缺陷。

图3为使用工业产品数据集对本发明方法进行测试，可见本发明能够重构出无缺陷的图像，并在缺陷检测中准确检测出原图的缺陷。

Claims (8)

1.一种基于量子化自编码器的无监督缺陷检测方法，其特征是构建一个量子化自编码器用于缺陷检测，量子化自编码器包括编码网络、编码表和解码网络，

量子化自编码器对输入图像进行归一化预处理，将预处理后的图像输入编码网络中，映射得到中间特征图，并加上位置编码信息；然后将中间特征图中每个向量与编码表中的向量进行最邻近匹配，用最接近的向量对其进行替换，得到量化的特征图；将量化后的特征图输入到解码网络中，解码网络将特征图重新恢复成图片；使用无缺陷的训练样本图像训练量子化自编码器；

缺陷检测时，将待检测图像进行归一化预处理后输入训练好的量子化自编码器，对于原图和量子化自编码器重建的图片，判别和定位缺陷。

2.根据权利要求1所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是量子化自编码器具体为：

1)编码网络encoder：

输入编码网络的图像映射到一个潜在空间中，得到连续的编码向量z′_e:

z′_e＝encoder(x)

在编码器中加入位置编码position encoding，用位置信息来对特征进行约束，采用Transformer中的提出的Sinusoidal Position Encoding：

编码向量z′_e与位置编码相加共同组成最终的特征向量z_e；

2)构建编码表和向量量化：

构建一个容量为K的Embedding层，其中每个向量的维度均为64，将其随机初始化：

E＝[e₁,e₂,e₃...e_K]

对于1)中得到的特征向量z_e，通过最邻近搜索匹配这K个向量中最相似的向量，记为z_q，使连续的编码向量实现离散化；

3)解码网络decode：

将得到的编码向量z_q输入到解码网络中，解码网络对编码向量进行上采样并重建原图：

x′＝decoder(z_q)

其中，使用直通估计的方法，在前向传播时使用z_q进行计算，在反向传播时使用z_e计算梯度：

decoder(z_q)→decoder(z_e+sg[z_q-z_e])

sg表示不计算梯度，这样在前向传播时该式等价于decoder(z_q)，在反向传播时由于z_q-z_e不提供梯度，该式等价于decoder(z_e)。

3.根据权利要求2所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是编码网络由底层编码器和顶层编码器组成，编码器由若干卷积层、BN层和ReLU层组成的残差模块堆叠而成；

图像输入到编码网络后首先经过底层编码器被下采样四倍，得到底层特征图，然后通过顶层编码器进一步下采样两倍，得到顶层特征图，底层特征图和顶层特征图分别加上各自对应的位置编码，然后顶层特征图上采样两倍后与底层特征图合并，此处的特征图即编码网络的输出z_e，z_e为一个m×m×64的多维向量：

之后通过1×1卷积层将z_e的维度压缩到64，然后在向量字典中进行最邻近搜索来实现向量量化。

4.根据权利要求2所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是构建编码表和向量量化中，z_e中每个位置的向量都与向量词典e_k中的所有向量计算余弦距离，k∈(1...K)，并与向量词典中最接近的向量的索引进行替换，从而实现离散化:

5.根据权利要求1所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是使用无缺陷的正常样本训练量子化自编码器，损失函数由重构损失、编码表损失和向量三元组损失三部分构成：

L＝L_r+L_c+L_e

其中，重构损失L_r由均方差损失MSE和结构相似度SSIM损失组成:

L_r＝MSE+SSIM

编码表损失L_c目的是让编码表中的向量与图像经过编码器后的向量尽可能接近，使用L2损失函数：

其中n为z_e中向量的总个数，z_ek表示z_e的第k个向量，z_qk是z_ek在编码表中最接近的向量；

相应的，也需要让图像经过编码器后的中间向量尽可能接近编码表中的向量，同样的使用L2损失函数，得到编码表损失L′_e：

进一步的，对L′_e进行改造得到向量三元组损失L_e，对于特征图z_e上的每一个向量z_ek，都在编码表中找到最接近它的向量z_p，和第二接近它的向量z_n，三个向量组成三元组共同计算损失，使得z_ek和z_p尽可能接近，z_ek和z_n尽可能远离，从而让编码表中的向量变得“稀疏”：

6.根据权利要求5所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是重构损失L_r由均方差损失MSE和基于空间注意力的结构相似性损失SAM-SSIM组成:

L_r＝MSE+SAM_SSIM

在量子化自编码器上增加空间注意力模块SAM，该模块由一系列的反卷积和激活函数组成，将z_q映射到大小为N×N×1的特征图W，对应每个图像块的空间注意力：

W＝SAM(z_q)

将空间注意力与对应图像块的SSIM相乘，得到SAM-SSIM：

其中W_i为第i个图像块对应的空间注意力。

7.根据权利要求1所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是缺陷检测时，将待检测的图片经过和训练过程一样的预处理后输入到量子化自编码器中，得到重建图像，对比原图和重建图像的差异，计算两张图的MSE和SSIM，先逐像素计算两张图的MSE损失，记为S_mse，再将原图和重建图都切分成N*N的图像块，计算SSIM，记为S_ssim，综合考虑这两个指标，得到加权后的缺陷置信度：

S＝S_ssim+S_mse

当缺陷置信度大于指定的阈值时，此处便被判别为缺陷。

8.根据权利要求1所述的一种基于量子化自编码器的无监督缺陷检测方法，其特征是缺陷检测时，将待检测的图片经过和训练过程一样的预处理后输入到量子化自编码器中，得到重建图像，对比原图和重建图像的差异，计算两张图的MSE和SAM_SSIM，先逐像素计算两张图的MSE损失，记为S_mse，再将原图和重建图都切分成N*N的图像块，计算SAM_SSIM，记为S_{sam_ssim},综合考虑这两个指标，得到加权后的缺陷置信度：

S＝S_ssim+S_{sam_ssim}

当缺陷置信度大于指定的阈值时，此处便被判别为缺陷。

Images