??一種新穎的擴散模型雙重采樣策略,DDPM + DDIM ~
論文:Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images
0、摘要
??病理區域的分割掩模在許多醫學應用中很有用,例如腦腫瘤和中風管理。此外,疾病圖像的健康反事實可以用來增強放射科醫生的培訓文件,并提高分割模型的可解釋性。)
??本文提出了一種弱監督方法來生成一個疾病圖像的健康版本,然后使用它來獲得像素異常圖。為此,本文首先使用 ACAT 生成一個大致覆蓋病理區域的顯著性圖,然后,本文提出了一種技術,允許對這些區域進行有針對性的修改,同時保留圖像的其余部分。
??具體而言,本文使用在健康樣本上訓練的擴散模型,并在采樣過程的每一步結合去噪擴散概率模型(DDPM)和去噪擴散隱式模型(DDIM)。DDPM 用于修改顯著性圖中受病變影響的區域,而 DDIM 則確保在這些區域之外重建正常的解剖結構。這兩部分在每一步都進行融合,以確保生成的樣本具有一致的外觀,并在編輯和未編輯部分之間實現無縫過渡。
??當本文方法應用于健康樣本時,輸入圖像在沒有顯著修改的情況下被重建。本文在腦病變分割任務中與替代的弱監督方法進行了比較,在所考慮的模型中獲得了最高的平均 Dice 和 IoU 分數。
??
Figure 1 | 首先通過 DDIMs 的反向采樣技術將異常圖像 x 0 x_0 x0? 轉換為其噪聲版本 x K x_K xK?。隨后,利用 DDPM 采樣來修改通過 ACAT 生成的顯著圖識別出的病理性區域,旨在根據周圍區域的上下文信息恢復正常的解剖結構:與此同時,圖像中不含任何病理元素的區域則通過 DDIM 采樣恢復至其原始外觀。在整個采樣過程中,這兩個組件被融合在一起,以確保編輯部分與未編輯部分之間過渡自然且逼真,最終生成的圖像 x ^ 0 \hat x_0 x^0? 在視覺上呈現出連貫且自然的效果;
1、引言
1.1、當前挑戰
??(1)自動編碼器(Autoencoders)并不總能生成清晰的圖像,也不能保證正確地映射到健康版本;生成對抗網絡(GANs)的訓練有時不穩定,依賴于許多超參數,并且會產生較差的樣本;
??(2)已有研究使用擴散模型和分類器引導來恢復正常的解剖結構。然而,用于指導采樣過程的梯度需要從未經噪聲處理的樣本訓練的分類器計算得出。由于在醫學成像中,樣本的類別通常由微小細節決定,而這些細節可能在僅幾次噪聲處理后就丟失了,因此這種分類器經常產生不可靠的預測。因此,采用這種方法時無法保證保留樣本的原始結構,許多正常組織的細節可能會被修改;
??(3)盡管對抗性反事實注意力(Adversarial Counterfactual Attention,ACAT)側重于生成反事實圖像,但其主要優勢在于能夠準確識別病理區域,這些區域隨后被用于分類流程中。然而,它在生成可信的反事實示例方面存在不足;
??圖 2 展示了這一現象,可以看到 ACAT 能夠生成一個大致識別出病理區域的顯著性圖(下面那行)。然而,在反事實示例中,病變仍然可見(上面那行)。相比之下,本文的方法不僅優化了顯著性圖,還生成了一個病理完全被消除的反事實圖像。
??
Figure 2 | IST-3 的原始圖像(a)以及相應的健康反事實圖像(第二行)和對應的異常圖(底行),這些異常圖是通過去噪自編碼器(DenoisingAE,b)、f-AnoGAN(c)、AnoDDPM(d)、AutoDDPM(e)、分類器引導(f)、無分類器引導(g)、ACAT(h)和 Dif-fuse(i)獲得的:ACAT 生成了一個合理的異常圖,但無法完全移除病變。Dif-fuse 在優化 ACAT 獲得的異常圖的同時,還創建了一個可信的反事實示例。其他方法則引入了偽影和/或對病理區域的識別不夠準確;
1.2、本文貢獻
??利用通過 ACAT 獲得的顯著性圖來指導擴散模型的圖像生成過程。首先,在健康樣本上訓練一個去噪擴散概率模型(DDPM),并結合 DDPM 和 DDIM 采樣,以去除圖像中的病理性區域:
??(1)本文提出了一種新穎的擴散模型雙重采樣策略,這種策略無需病變注釋,即可對由分割掩碼識別出的感興趣區域(ROIs)進行修復,同時保留圖像的其余部分;
??本文的創新之處在于在每個時間步混合兩個組件的方法,這使得編輯部分和未編輯部分之間實現了平滑的融合。這不僅能夠生成醫學圖像的真實反事實示例,還能生成病理區域的異常圖;
??(2)在 WMH 和 BraTS 2021 數據集上,本文的方法與現有的弱監督醫學圖像分割方法進行比較,結果顯示本文方法在兩個數據集中所考慮的方法中均實現了最高的平均 Dice 和 IoU 分數;
??本文方法在圖像質量方面也具有可比性,這通過在 IST-3 上使用核初始距離(Kernel Inception Distance,KID)來衡量,與無約束(無掩碼)的擴散采樣方法相比,本文方法具有更準確的異常分割這一額外優勢;
2、方法
??本文的采樣方法不僅能夠生成高度逼真的反事實圖像,還能優化第一步中通過 ACAT 獲得的初始顯著性圖。這是因為選定的區域可能不會被擴散模型完全修改,從而使得初始顯著性圖中識別出的健康解剖特征得以保留。方法總覽如 圖 1 所示。
2.1、Diffusion Models
??原文略,可參考:【Diffusion綜述】醫學圖像分析中的擴散模型(一)中 2.2 節;
2.2、Dif-fuse
??本文使用了在健康樣本上訓練的 DDPM,并通過 ACAT 從反事實示例中獲取顯著圖,本文選擇 ACAT,因為它在識別腦部和肺部 CT 掃描中的病理性區域方面表現出色。然而,原則上,顯著圖也可以通過其他任何方法生成。
??Dif-fuse 算法流程:
??給定一個病變圖像 x 0 x_0 x0?,首先選擇一個噪聲量 K ∈ [ 0 , T ] K∈[0,T] K∈[0,T],利用逆向 DDIM 采樣方案,將圖像映射到其噪聲版本 x K x_K xK?:
??隨后,使用一個尺寸為 5 × 5 5×5 5×5 的高斯核對顯著性圖進行平滑處理,以獲得一個更加均勻且孤立像素更少的掩碼 m m m。接著,通過 DDPM 采樣對掩碼內的病變區域進行編輯。
??由于擴散模型是在正常樣本上訓練的,這些區域被映射為健康外觀。其余解剖結構需要保留,因此采用 DDIM 采樣法處理掩模外的區域。為了獲得一致的結果,將掩模部分與圖像其余部分混合:
??其中 ⊙ ⊙ ⊙ 是哈達瑪積。這樣,編輯過程集中在顯著圖捕捉到的部分,防止對掃描結構特征進行隨意更改。實際上,DDIM 采樣保證了不需要編輯的部分能夠重建。此外,病理部分的修改由 DDPM 在考慮周圍解剖背景的情況下進行。
??當使用公式(9) 計算 x ^ t ? 1 \hat x_{t?1} x^t?1? 時,兩個組件的總和可能不會產生完全一致的結果。然而,這種不一致性會在下一個擴散步驟中得到解決,該步驟能更好地融合這兩個組件。
??如果只是用 DDPM 計算 x ^ 0 \hat x_0 x^0?,然后在采樣過程結束時才應用掩模,情況就不會如此。圖 3 展示了這種效果,可以看到,僅在采樣過程(b)結束時應用掩模生成的正常圖像,表現出一些偽影,并且編輯區域和未編輯區域之間缺乏平滑過渡。
??
Figure 3 | IST-3 的輸入圖像(a)以及僅在采樣過程結束時應用掩碼生成的正常圖像(b):可以觀察到(b)中存在一些偽影,并且編輯部分與未編輯部分之間沒有平滑的過渡;
??通過這種方式,能夠獲得給定病理圖像的正常版本。為了獲得異常圖,首先計算原始圖像與生成圖像之間的差異,然后對得到的差異圖依次應用腐蝕和膨脹操作,其中腐蝕和膨脹均使用 5 × 5 5×5 5×5 的核,以去除噪聲,最后再次進行膨脹和腐蝕操作,同樣使用 5 × 5 5×5 5×5 的核,以填補差異圖中的小孔洞。
2.3、訓練細節
??(1)擴散模型 60000 次迭代,batch size=10;
??(2)AdamW 優化器,學習率 1e-4, β 1 = 0.9 , β 2 = 0.999 β_1 = 0.9, β_2 = 0.999 β1?=0.9,β2?=0.999,weight decay=0.05;
??(3)Improved DDPM 采用的損失;
??(4)EMA:0.99,線性噪聲調度: [ 1 0 ? 4 , 0.02 ] [10^{-4}, 0.02] [10?4,0.02];
??(5)1000 個采樣步,UNet 通道為128;
3、實驗與結果
3.1、數據集
??(1)IST-3 數據集(CT):3035 名出現中風癥狀的患者的大腦影像數據;本文共考慮了5681 次掃描,其中 46.31% 被分類為陰性(無病變),其余掃描結果為陽性。對每次掃描的 11 個切片進行了處理,并將每個切片調整至 256×256。
??(2)BraTS 2021 數據集(MRI):1251 名患者的掃描,每個掃描有 155 層,本文移除最上面和最下面的 25 層,利用 0 填充到 256×256,最終得到 131,164 個切片,其中 79,113 個為陽性。
??(3)WMH 數據集(MRI):白質高信號強度分割挑戰,使用測試集中 110 次掃描,對每個切片中心裁剪并調整大小為 256×256。
??由于 IST-3 數據集中沒有病變的注釋,利用該數據集來評估生成圖像的質量,而不是分割精度。另一方面,對于 BraTS 2021 和 WMH 數據集,能夠獲取病變的注釋,這能夠對創建的異常圖進行定量分析。IST-3 和 BraTS 2021 被劃分為訓練集、驗證集和測試集,比例為 70%、15%和15%。在 WMH 上,本文在沒有進一步微調的情況下評估在 BraTS 2021 上訓練的模型,以測試它們的跨領域泛化能力。
3.2、反事實示例
??
Figure 4 | 原始圖像來自 BraTS 2021,包含真實分割掩模(a)和健康圖像(上排),以及相應的異常圖(下排),這些圖像是通過 DenoisingAE(b)、f-AnoGan(c)、AnoDDPM(d)、AutoDDPM(e)、分類器引導(f)、無分類器引導(g)和 Dif-fuse(h)獲得的:f-Ano GAN在生成可信的反事實方面表現不佳,而其他方法則產生了更高質量的結果。然而,DenoisingAE、AnoDDPM 和 AutoDDPM 并不能完全去除損傷,而使用 CG 和 CFG 生成的反事實則表現出一些偽影;
??
Figure 6 | 來自 BraTS 2021 的健康輸入圖像(a)、通過 Dif-fuse 生成的圖像(b)以及異常圖(c):可以觀察到,本文的方法能夠很好地重建健康樣本;
3.3、超參數
??
Figure 5 | 在驗證數據集上,使用不同的閾值百分位數來二值化顯著性圖以及不同的噪聲量 K K K 所獲得的 Dice 分數: 在 K = 500 K=500 K=500 以及使用顯著性圖中第 90 百分位的像素時,能夠獲得最佳結果;
??
Figure 7 | 來自 BraTS 2021 的原始圖像及其真實的分割掩碼(a),以及通過 Dif-fuse 生成的健康圖像(頂部行)和對應的異常圖(底部行),其中分別使用了 250(b)、500(c)和 750(d)步的加噪過程: 當使用較低的噪聲量時,病理區域未能被完全移除;而當噪聲水平過高時,則可能會引入顯著的偽影;
3.4、定量分析
??
Table 1 | 在 BraTS 2021 和 WMH 測試數據上的 Dice 和 IoU 分數,以及在 IST-3 上的 KID(核初始距離),這些指標均是基于三次運行的平均值(并附上了標準誤差):Dif-fuse 在 BraTS 2021 和 WMH 上均實現了最佳的異常分割性能。而采用 DDIM 反演后接 DDPM 采樣的消融實驗在 IST-3 上獲得了最佳的 KID 值;
??有點意思,樸素又實惠 ?(???ω???)?