基于DiAD擴(kuò)散模型的多類異常檢測工作

背景和動機(jī)

在工業(yè)生產(chǎn)過程中，工業(yè)異常檢測是一個重要的環(huán)節(jié)，旨在及時發(fā)現(xiàn)和排除產(chǎn)品的制造異常，確保產(chǎn)品質(zhì)量和安全性。近年來，隨著計(jì)算機(jī)視覺和圖像處理技術(shù)的快速發(fā)展，視覺技術(shù)被廣泛應(yīng)用于工業(yè)異常檢測中。一般的單類異常檢測算法分別為不同類別的物體訓(xùn)練不同的模型，非常消耗內(nèi)存，并且也不適用于正常樣本表現(xiàn)出大的類內(nèi)多樣性的情況，本文作者致力于用統(tǒng)一的框架從不同的物體類別中檢測異常。

現(xiàn)有的基于計(jì)算機(jī)視覺的工業(yè)異常檢測技術(shù)包括基于特征的、基于重構(gòu)的和基于合成的技術(shù)。最近，擴(kuò)散模型因其強(qiáng)大的生成能力而聞名，因此本文作者希望通過擴(kuò)散模型將異常區(qū)域重構(gòu)成正常。然而如下圖1所示，直接將擴(kuò)散模型應(yīng)用于多類異常檢測會存在幾點(diǎn)問題：對于DDPM來說，應(yīng)用于多類異常檢測會出現(xiàn)類別錯誤的情況，因?yàn)樵诩尤?000步噪聲后，圖像已變成隨機(jī)高斯噪聲，因?yàn)闆]有其他限制條件，DDPM在去噪過程中從隨機(jī)高斯噪聲中去噪最終獲得隨機(jī)類別的圖像。LDM通過交叉注意力引入了條件約束，在多類異常檢測應(yīng)用場景中，LDM可以通過引入類別條件解決了DDPM在多類異常檢測任務(wù)中類別錯誤的問題；然而LDM仍然無法解決在隨機(jī)高斯噪聲下去噪并保持圖像語義信息一致性的問題，比如釘子的方向、齒輪的方位等。因此，現(xiàn)有的擴(kuò)散模型盡管展現(xiàn)了其強(qiáng)大的生成能力，但是無法很好的解決多類異常檢測的任務(wù)。

圖1：上圖展示了DDPM、LDM和Ours擴(kuò)散模型去噪網(wǎng)絡(luò)的框架，下圖展示了三種模型對于同樣的輸入異常圖像的重構(gòu)效果

所以本文作者提出了DiAD來解決多類異常檢測任務(wù)。在MVTec-AD、VisA等數(shù)據(jù)集上在圖像和像素級別的AUROC、AP、F1max和PRO共七個指標(biāo)上實(shí)現(xiàn)了新的SoTA，超越了UniAD、RD4AD等模型。

網(wǎng)絡(luò)結(jié)構(gòu)

本文作者提出一種基于擴(kuò)散模型框架的多類異常檢測方法。本方法包含三個空間：Pixel Space、Latent Space和Feature Space。首先輸入圖像在Pixel Space中經(jīng)過預(yù)訓(xùn)練的Auto-encoder得到Latent Variable ，隨后Latent Variable 進(jìn)入Latent Space經(jīng)過向一步步加入隨機(jī)高斯噪最終得到近隨機(jī)高斯噪聲，加噪后的Latent Variable 輸入至SD Denoising Network的同時輸入圖像輸入至語義引導(dǎo)網(wǎng)絡(luò)(Semantic-Guided Network)，經(jīng)過Semantic-Guided Network的輸入圖像將會加入到SD Denoising Network的Decoder Blocks中，經(jīng)過大量的去噪過程后得到重構(gòu)的Latent Variable ，此時再進(jìn)入Pixel Space中的Auto-decoder得到重構(gòu)圖像，最終輸入圖像和重構(gòu)圖像將同時輸入到Feature Space中，通過同一個預(yù)訓(xùn)練的特征提取網(wǎng)絡(luò)提取不同尺度的特征，比較不同特征圖上的余弦相似度用于異常得分的計(jì)算與定位并將不同尺度得到的結(jié)果合起來得到最終像素級別的異常定位得分和異常分類得分。本文作者改進(jìn)了擴(kuò)散模型的Denoising Network，在Denoising Network的基礎(chǔ)上添加了與其結(jié)構(gòu)相似的Semantic-Guided Network來保持輸入圖像與重構(gòu)圖像語義信息的一致性，使擴(kuò)散模型能夠在高步數(shù)加噪下仍保持與原輸入圖像語義信息的高度一致性并將異常區(qū)域重構(gòu)為正常區(qū)域。? ?

圖2：DiAD模型框架

語義引導(dǎo)網(wǎng)絡(luò)(Semantic-Guided Network)

語義引導(dǎo)網(wǎng)絡(luò)首先通過卷積神經(jīng)網(wǎng)絡(luò)將輸入圖像下采樣到與加噪后的隱變量經(jīng)過去噪網(wǎng)絡(luò)第一層編碼塊得到的特征圖同維度同尺度特征，其次復(fù)制預(yù)訓(xùn)練去噪網(wǎng)絡(luò)的編碼塊和解碼塊的參數(shù)用于之后模型的微調(diào)，為了保持重構(gòu)圖像與輸入圖像語義信息的一致性，將語義引導(dǎo)網(wǎng)絡(luò)的中間層和一層解碼塊與去噪網(wǎng)絡(luò)對應(yīng)模塊相連，最終去噪網(wǎng)絡(luò)的輸出為 ? ?

其中表示去噪網(wǎng)絡(luò)解碼塊、表示去噪網(wǎng)絡(luò)中間塊、表示去噪網(wǎng)絡(luò)編碼塊、表示語義引導(dǎo)網(wǎng)絡(luò)中間塊、表示卷積神經(jīng)網(wǎng)絡(luò)層、表示語義引導(dǎo)網(wǎng)絡(luò)解碼層。

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

為了將不同空間尺度的信息結(jié)合起來，使用空間感知特征融合模塊融合在語義引導(dǎo)網(wǎng)絡(luò)中不同尺度編碼塊的輸出結(jié)果， 表示編碼塊四中的低尺度的輸出特征圖，表示編碼塊三中的高尺度的輸出特征圖，編碼塊四中的最終的輸出特征圖為：

表示編碼塊三中有三層在本次實(shí)驗(yàn)中使用，表示一個基本的包含一個3×3卷積層、歸一化層和激活層的卷積模塊。

圖3：空間感知特征融合模塊

異常分?jǐn)?shù)計(jì)算

將重構(gòu)圖像和輸入圖像共同輸入到特征空間中同一個預(yù)訓(xùn)練的卷積神經(jīng)網(wǎng)絡(luò)resnet50特征提取器中提取不同尺度的特征，并通過余弦相似度計(jì)算不同尺度的缺陷得分，計(jì)算方式為：

其中代表第層特征。最終的缺陷定位得分為:

其中表示上采樣率、表示使用的特征層的數(shù)量，缺陷的分類得分為經(jīng)過全局平均池化后的缺陷定位得分的最大值。

實(shí)驗(yàn)結(jié)果

MVTec-AD數(shù)據(jù)集實(shí)驗(yàn)結(jié)果

圖4：MVTec-AD數(shù)據(jù)集上多類異常檢測實(shí)驗(yàn)結(jié)果，對應(yīng)的指標(biāo)分別為圖像級別AUROC/AP/F1max。

圖5：MVTec-AD數(shù)據(jù)集上多類異常檢測實(shí)驗(yàn)結(jié)果，對應(yīng)的指標(biāo)分別為像素級別AUROC/AP/F1max。

VisA數(shù)據(jù)集實(shí)驗(yàn)結(jié)果

圖6：VisA數(shù)據(jù)集上多類異常檢測實(shí)驗(yàn)結(jié)果，cls表示圖像級別的結(jié)果而seg表示像素級別的結(jié)果。

總結(jié)

本文作者解決了現(xiàn)有擴(kuò)散模型框架無法應(yīng)對多類異常檢測任務(wù)并提出了DiAD擴(kuò)散模型框架用于多類異常檢測。具體而言，作者在SD去噪網(wǎng)絡(luò)的基礎(chǔ)上引入了語義引導(dǎo)網(wǎng)絡(luò)(Semantic-Guided Network)保持輸入圖像和重構(gòu)圖像的語義一致性，并且還提出了空間感知特征融合模塊(Spatial-aware Feature Fusion Block)將不同尺度的特征相融合。最終本文實(shí)現(xiàn)了在保持與輸入圖像語義信息一致的前提下將異常區(qū)域重構(gòu)成正常圖像，同時實(shí)現(xiàn)SoTA。

審核編輯：黃飛