基于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。

審核編輯：黃飛