当前位置：首页 > news >正文

【2023工业图像异常检测文献】DiAD: 基于扩散模型的多类异常检测方法

news 2024/10/30 10:04:18

DiAD: A Diffusion-based Framework for Multi-class Anomaly Detection

1、Background

当前主流的三种异常检测方法：

基于合成的方法：在正常图像上合成异常，通过训练模型识别这些合成的异常来提高检测和定位能力。
基于嵌入的方法：将图像信息编码到特征空间中，使用预训练的网络提取特征，并计算特征的相似性作为异常分数。
基于重建的方法：训练模型在正常数据集上学习模式和特征，然后在测试阶段重建异常图像，通过比较重建图像和输入图像来定位异常。

最近，扩散模型展示了它们强大的图像生成能力。然而，直接使用当前主流的扩散模型无法有效解决多类别异常检测问题。

在这里插入图片描述

1）对于图1（a）中的去噪扩散概率模型（DDPM）（Ho、Jain和Abbeel，2020年），在执行多类别设置时，这种方法可能会遇到生成图像类别错误分类的问题。原因是在向输入图像添加T个时间步的噪声后，图像失去了其原始类别信息。在推理过程中，基于这种高斯噪声类分布进行去噪，可能会生成属于不同类别的样本。

2）潜在扩散模型（LDM）（Rombach等人，2022年）具有作为类别条件的嵌入器，如图1-（b）所示，它不存在DDPM中发现的类别错误分类问题。然而，LDM仍然无法解决生成图像中语义信息丢失的问题。LDM无法同时保留输入图像的语义信息和重建异常区域。例如，它们可能无法在螺丝和榛子等对象的方向上保持与输入图像的一致性，并且在纹理类别图像上与原始图像有显著差异。

为了解决上述问题，我们提出了一个基于扩散的框架DiAD，用于多类别异常检测和定位，如图2所示，包括三个组件：像素空间自动编码器、潜在空间去噪网络和特征空间ImageNet预训练模型。

为了在重建异常区域的同时有效地保持与原始图像一致的语义信息，我们提出了与稳定扩散（SD）去噪网络在LDM中连接的语义引导（SG）网络。

为了进一步提高保留原始图像细节和重建大规模缺陷的能力，我们提出了空间感知特征融合（SFF）块，以整合不同尺度的特征。

最后，通过预训练的模型传递重建和输入图像，以提取不同尺度的特征并计算异常分数。

supplements

去噪扩散概率模型（Denoising Diffusion Probabilistic Model, DDPM）、潜在扩散模型（Latent Diffusion Model, LDM）DiAD框架区别：

去噪扩散概率模型（DDPM）：
- 包含两个过程：正向扩散过程和逆向去噪过程。
- 正向扩散过程通过逐步添加高斯噪声来生成噪声样本。
- 逆向去噪过程从添加了噪声的样本开始，通过模型预测的噪声分布来重建原始数据。
- 模型通过最小化预测噪声和实际噪声之间的差异来训练。
潜在扩散模型（LDM）：
- 专注于低维潜在空间，并引入条件机制。
- 由预训练的自动编码器模型和去噪U-Net网络组成，后者基于注意力机制。
- 在潜在表示空间中进行扩散和去噪操作，然后通过解码器重建图像。
- 训练目标是最小化预测噪声和实际噪声之间的差异，同时考虑条件机制，如文本或图像。
DiAD框架的方法：
- 使用预训练的编码器将输入图像编码为潜在空间表示。
- 向潜在表示中添加噪声，然后使用与SG网络连接的SD去噪网络进行去噪处理。
- 去噪过程重复与扩散过程相同的时间步。
- 重建的潜在表示通过预训练的解码器恢复到原始图像水平。
- 在异常检测和定位方面，输入和重建图像被送入同一个预训练模型以提取不同尺度的特征并计算特征差异。

2、Method

DiAD框架通过结合像素空间的自动编码器、潜在空间的语义引导网络和特征空间的预训练特征提取器，实现了对异常区域的有效重建和检测。SFF块的引入进一步提高了模型在不同尺度上的特征融合能力，从而提高了异常检测的准确性。

在这里插入图片描述

DiAD框架的设计：

DiAD框架的组成：
- 像素空间自动编码器：用于将输入图像编码和解码。
- 潜在空间语义引导（SG）网络：与稳定扩散（SD）去噪网络相连，用于重建异常区域并保留原始图像的语义信息。
- 特征空间预训练特征提取器：用于从输入和重建图像中提取特征并生成异常图。
SG网络的设计：
- 为了解决多类别异常检测中的问题，SG网络被提出以改善LDM在重建异常区域时保留输入图像语义信息的能力。
- 输入图像首先被编码为潜在空间表示，然后通过添加噪声进行正向扩散过程。
- 在逆向去噪过程中，SG网络和SD去噪网络共同工作，以重建图像并恢复其原始结构。
空间感知特征融合（SFF）块：
- 为了同时保留原始样本的正常信息和重建大规模异常区域，提出了SFF块。
- SFF块通过整合不同尺度的特征来增强模型在不同场景下的重建能力，无论是细节纹理还是大面积缺陷。
异常检测和定位：
- 在推理阶段，通过扩散和去噪过程获得重建图像。
- 使用预训练的特征提取器从输入图像和重建图像中提取特征，并计算不同尺度特征图上的差异，以生成异常图。
- 通过计算特征图之间的余弦相似性来得到异常分数，从而实现异常检测和定位。

pseudo-code

import torch
import torch.nn as nn
from torchvision import models# 假设我们有一个预训练的自动编码器，它由编码器E和解码器D组成
class AutoEncoder(nn.Module):def __init__(self):super(AutoEncoder, self).__init__()self.encoder = ...self.decoder = ...def forward(self, x):encoded = self.encoder(x)decoded = self.decoder(encoded)return decoded# 假设我们有一个预训练的特征提取器
class FeatureExtractor(nn.Module):def __init__(self):super(FeatureExtractor, self).__init__()# 通常使用预训练的模型，如ResNetself.model = models.resnet50(pretrained=True)# 冻结所有层，只使用特征提取部分for param in self.model.parameters():param.requires_grad = Falsedef forward(self, x):return self.model(x)# 假设我们有一个语义引导网络SG和稳定扩散网络SD
class SGNetwork(nn.Module):def __init__(self):super(SGNetwork, self).__init__()# 定义SG网络的结构...def forward(self, z):# SG网络的前向传播...class SDNetwork(nn.Module):def __init__(self):super(SDNetwork, self).__init__()# 定义SD网络的结构...def forward(self, z_noisy):# SD网络的前向传播...# 异常检测的主要流程
def anomaly_detection(input_image):# 步骤1: 使用自动编码器的编码器部分将输入图像编码为潜在表示latent_representation = autoencoder.encoder(input_image)# 步骤2: 使用SG网络和SD网络重建潜在表示reconstructed_latent = sg_network(latent_representation)reconstructed_latent = sd_network(reconstructed_latent)# 步骤3: 使用自动编码器的解码器部分重建图像reconstructed_image = autoencoder.decoder(reconstructed_latent)# 步骤4: 使用预训练的特征提取器提取输入图像和重建图像的特征features_input = feature_extractor(input_image)features_reconstructed = feature_extractor(reconstructed_image)# 步骤5: 计算特征之间的差异，得到异常分数anomaly_score = calculate_anomaly_score(features_input, features_reconstructed)return anomaly_score# 辅助函数：计算异常分数
def calculate_anomaly_score(features_input, features_reconstructed):# 这里简化为使用余弦相似度作为异常分数的计算方式anomaly_score = 1 - nn.CosineSimilarity(dim=1)(features_input, features_reconstructed)return anomaly_score# 初始化模型
autoencoder = AutoEncoder()
feature_extractor = FeatureExtractor()
sg_network = SGNetwork()
sd_network = SDNetwork()# 假设input_image是已经预处理的输入图像张量
input_image = ...# 执行异常检测
anomaly_score = anomaly_detection(input_image)