基于DiAD扩散模型的多类异常检测工作

背景和动机

在工业生产过程中，工业异常检测是一个重要的环节，旨在及时发现和排除产品的制造异常，确保产品质量和安全性。近年来，随着计算机视觉和图像处理技术的快速发展，视觉技术被广泛应用于工业异常检测中。一般的单类异常检测算法分别为不同类别的物体训练不同的模型，非常消耗内存，并且也不适用于正常样本表现出大的类内多样性的情况，本文作者致力于用统一的框架从不同的物体类别中检测异常。

现有的基于计算机视觉的工业异常检测技术包括基于特征的、基于重构的和基于合成的技术。最近，扩散模型因其强大的生成能力而闻名，因此本文作者希望通过扩散模型将异常区域重构成正常。然而如下图1所示，直接将扩散模型应用于多类异常检测会存在几点问题：对于DDPM来说，应用于多类异常检测会出现类别错误的情况，因为在加入1000步噪声后，图像已变成随机高斯噪声，因为没有其他限制条件，DDPM在去噪过程中从随机高斯噪声中去噪最终获得随机类别的图像。LDM通过交叉注意力引入了条件约束，在多类异常检测应用场景中，LDM可以通过引入类别条件解决了DDPM在多类异常检测任务中类别错误的问题；然而LDM仍然无法解决在随机高斯噪声下去噪并保持图像语义信息一致性的问题，比如钉子的方向、齿轮的方位等。因此，现有的扩散模型尽管展现了其强大的生成能力，但是无法很好的解决多类异常检测的任务。   

图1：上图展示了DDPM、LDM和Ours扩散模型去噪网络的框架，下图展示了三种模型对于同样的输入异常图像的重构效果

所以本文作者提出了DiAD来解决多类异常检测任务。在MVTec-AD、VisA等数据集上在图像和像素级别的AUROC、AP、F1max和PRO共七个指标上实现了新的SoTA，超越了UniAD、RD4AD等模型。

网络结构

本文作者提出一种基于扩散模型框架的多类异常检测方法。本方法包含三个空间：Pixel Space、Latent Space和Feature Space。首先输入图像在Pixel Space中经过预训练的Auto-encoder得到Latent Variable ，随后Latent Variable 进入Latent Space经过向一步步加入随机高斯噪最终得到近随机高斯噪声，加噪后的Latent Variable 输入至SD Denoising Network的同时输入图像输入至语义引导网络(Semantic-Guided Network)，经过Semantic-Guided Network的输入图像将会加入到SD Denoising Network的Decoder Blocks中，经过大量的去噪过程后得到重构的Latent Variable ，此时再进入Pixel Space中的Auto-decoder得到重构图像，最终输入图像和重构图像将同时输入到Feature Space中，通过同一个预训练的特征提取网络提取不同尺度的特征，比较不同特征图上的余弦相似度用于异常得分的计算与定位并将不同尺度得到的结果合起来得到最终像素级别的异常定位得分和异常分类得分。本文作者改进了扩散模型的Denoising Network，在Denoising Network的基础上添加了与其结构相似的Semantic-Guided Network来保持输入图像与重构图像语义信息的一致性，使扩散模型能够在高步数加噪下仍保持与原输入图像语义信息的高度一致性并将异常区域重构为正常区域。   

图2：DiAD模型框架

语义引导网络(Semantic-Guided Network)

语义引导网络首先通过卷积神经网络将输入图像下采样到与加噪后的隐变量经过去噪网络第一层编码块得到的特征图同维度同尺度特征，其次复制预训练去噪网络的编码块和解码块的参数用于之后模型的微调，为了保持重构图像与输入图像语义信息的一致性，将语义引导网络的中间层和一层解码块与去噪网络对应模块相连，最终去噪网络的输出为    

其中表示去噪网络解码块、表示去噪网络中间块、表示去噪网络编码块、表示语义引导网络中间块、表示卷积神经网络层、表示语义引导网络解码层。

空间感知特征融合模块(Spatial-aware Feature Fusion Block)

为了将不同空间尺度的信息结合起来，使用空间感知特征融合模块融合在语义引导网络中不同尺度编码块的输出结果， 表示编码块四中的低尺度的输出特征图，表示编码块三中的高尺度的输出特征图，编码块四中的最终的输出特征图为：

表示编码块三中有三层在本次实验中使用，表示一个基本的包含一个3×3卷积层、归一化层和激活层的卷积模块。

图3：空间感知特征融合模块

异常分数计算    

将重构图像和输入图像共同输入到特征空间中同一个预训练的卷积神经网络resnet50特征提取器中提取不同尺度的特征，并通过余弦相似度计算不同尺度的缺陷得分，计算方式为：

其中代表第层特征。最终的缺陷定位得分为:

其中表示上采样率、表示使用的特征层的数量，缺陷的分类得分为经过全局平均池化后的缺陷定位得分的最大值。

实验结果

MVTec-AD数据集实验结果

图4：MVTec-AD数据集上多类异常检测实验结果，对应的指标分别为图像级别AUROC/AP/F1max。

图5：MVTec-AD数据集上多类异常检测实验结果，对应的指标分别为像素级别AUROC/AP/F1max。

VisA数据集实验结果

图6：VisA数据集上多类异常检测实验结果，cls表示图像级别的结果而seg表示像素级别的结果。

总结

本文作者解决了现有扩散模型框架无法应对多类异常检测任务并提出了DiAD扩散模型框架用于多类异常检测。具体而言，作者在SD去噪网络的基础上引入了语义引导网络(Semantic-Guided Network)保持输入图像和重构图像的语义一致性，并且还提出了空间感知特征融合模块(Spatial-aware Feature Fusion Block)将不同尺度的特征相融合。最终本文实现了在保持与输入图像语义信息一致的前提下将异常区域重构成正常图像，同时实现SoTA。

审核编辑：黄飞