基于Diffusion Probabilistic Model的醫(yī)學(xué)圖像分割

前言

在前面的一篇文章ICLR 2023：基于 diffusion adversarial representation learning 的血管分割中，我們已經(jīng)介紹過了 diffusion model 在醫(yī)學(xué)圖像分割上的一個(gè)應(yīng)用，推薦對 diffusion model 不了解的同學(xué)優(yōu)先閱讀，其中講了一些基本概念。上一篇文章是將 diffusion 應(yīng)用到自監(jiān)督學(xué)習(xí)中，而 MedSegDiff 是一個(gè)有監(jiān)督的框架，現(xiàn)在已更新到 V2 版本， V2 版本區(qū)別于 V1 使用了 Transformer，且適用于多分類。MedSegDiff-V1 已被接收在 MIDL 2023。

MedSegDiff

MedSegDiff 在原版 DPM 的基礎(chǔ)上引入了動(dòng)態(tài)條件編碼，增強(qiáng) DPM 在醫(yī)學(xué)圖像分割上的分步注意力能力。特征頻率解析器（FF-Parser）可以消除分割過程中損壞的給定掩碼中的高頻噪聲。DPM 是一種生成模型，由兩個(gè)階段組成，正向擴(kuò)散階段和反向擴(kuò)散階段。在正向過程中，通過一系列步驟 T，將高斯噪聲逐漸添加到分割標(biāo)簽 x0 中。在反向過程中，訓(xùn)練神經(jīng)網(wǎng)絡(luò)通過反向噪聲過程來恢復(fù)原始數(shù)據(jù)：

其中 theta 是反向過程參數(shù)。從高斯噪聲開始，p(xT) 表示原始圖像，反向過程將潛在變量分布 p(xT) 轉(zhuǎn)換為數(shù)據(jù)分布 p(x0)。反向過程逐步恢復(fù)噪聲圖像，以獲得最終的清晰分割。該模型使用 U-Net 作為學(xué)習(xí)網(wǎng)絡(luò)，步長估計(jì)函數(shù)由原始圖像先驗(yàn)條件確定：

其中 EI 是條件特征嵌入，即原始圖像嵌入，Ex 是當(dāng)前步驟的分割映射特征嵌入。這兩個(gè)組件被添加并發(fā)送到 U-Net 的解碼器進(jìn)行重建。步長索引 t 與新增的嵌入和解碼器功能集成在一起，使用共享的 look-up table 進(jìn)行嵌入，這在 DDPM 的論文中有介紹。總而言之，MedSegDiff 模型基于 DPM，使用 U-Net 進(jìn)行學(xué)習(xí)。步長估計(jì)函數(shù)由原始圖像先驗(yàn)得到，步長索引與新增的嵌入和解碼器功能集成在一起。使 MedSegDiff 在三項(xiàng)具有不同圖像模式的醫(yī)學(xué)分割任務(wù)中表現(xiàn)不錯(cuò)。先看下 MedSegDiff 整體流程圖：

下面我們將逐一介紹動(dòng)態(tài)條件編碼和 FF-Parser。

動(dòng)態(tài)條件編碼

對于 MRI 或超聲之類的低對比度圖像，很難將感興趣的對象與背景分開。所以使用動(dòng)態(tài)條件編碼方法來解決這個(gè)問題?？梢宰⒁獾?，原始圖像包含準(zhǔn)確的目標(biāo)分割信息，不過很難與背景區(qū)分開，而當(dāng)前步驟的 grand truth 包含增強(qiáng)的目標(biāo)區(qū)域，但不準(zhǔn)確。

為了整合這兩個(gè)信息來源，使用類似注意力的機(jī)制將條件特征圖的每個(gè)尺度與當(dāng)前步驟的編碼特征融合。這種融合是首先對兩個(gè)特征圖分別應(yīng)用層歸一化，然后將它們相乘以獲得 affine map 來實(shí)現(xiàn)的，再將 affine map 與條件編碼特征相乘以增強(qiáng)注意力區(qū)域。如 MedSegDiff 流程圖所示，此操作應(yīng)用于中間兩個(gè)階段，其中每個(gè)階段都是在 Resnet34 之后實(shí)現(xiàn)的卷積階段。但是，集成當(dāng)前條件編碼功能可能會(huì)產(chǎn)生額外的高頻噪聲。為了解決這個(gè)問題，使用 FF-Parser 來限制特征中的高頻分量。

FF-Parser

FF-Parser 的流程如下圖所示：

FF-Parser 是頻率濾波器的可學(xué)習(xí)版本，它可以全局調(diào)整特定頻率的分量，以限制高頻分量進(jìn)行自適應(yīng)集成。首先使用二維 FFT（快速傅立葉變換）沿空間維度對解碼器特征圖 m 進(jìn)行變換，生成頻譜 M。然后，將參數(shù)化的注意力地圖 A 與 M 相乘以調(diào)整頻譜，得出 M'。最后，使用逆向 FFT 將 M' 反向回空間域，以獲得修改后的特征圖 m'。使用 FF-Parser 可以學(xué)習(xí)適用于傅里葉空間特征的權(quán)重圖，該權(quán)重圖可用于全局調(diào)整特定頻率的分量。這種技術(shù)不同于空間注意力，后者調(diào)整特定空間位置的組成部分。

實(shí)驗(yàn)

下圖分別是腦部 MRI、眼底視盤和甲狀腺結(jié)節(jié)的超聲圖像分割結(jié)果的可視化，可以看出 MedSegDiff 在簡單解刨結(jié)構(gòu)的二分類上效果還是不錯(cuò)的。

對比其他 SOTA 方法的結(jié)果如下表：

MedSegDiff-V2

MedsegDiff-v2 通過將 Transformer 機(jī)制整合到原始的U-Net骨干中，增強(qiáng)了基于擴(kuò)散的 MedSegDiff-v1。具體來說，MedsegDiff-v2 引入了新的 Spectrum-Space Transformer（SS-former），對噪聲和語義特征之間的相互作用進(jìn)行建模。驗(yàn)證了 Medsegdiff-v2 對具有不同模態(tài)圖像的五個(gè)分割數(shù)據(jù)集的十八個(gè)器官的有效性。

概述

如下圖所示，MedsegDiff-v2 結(jié)合了錨點(diǎn)條件和語義條件兩種不同的條件方式，以提高擴(kuò)散模型的性能。錨點(diǎn)條件將錨分割特征（條件模型的解碼分割特征）集成到擴(kuò)散模型的編碼特征中。即允許使用粗略但靜態(tài)的參照來初始化擴(kuò)散模型，有助于減少擴(kuò)散方差。

然后將語義條件強(qiáng)加于擴(kuò)散模型的 embedding，理解為將條件模型的語義 embedding 集成到擴(kuò)散模型的 embedding 中。這種條件集成由 SS-former 實(shí)現(xiàn)，它彌合了噪聲和語義嵌入之間的鴻溝，并利用 Transformer 的全局和動(dòng)態(tài)特性抽象出更強(qiáng)的特征表達(dá)形式。

Medsegiff-v2 是使用 DPM 的標(biāo)準(zhǔn)噪聲預(yù)測損失 Lnoise 和錨損失 Lanchor 進(jìn)行訓(xùn)練的。Lanchor 是 Dice loss 和 CE loss 的組合。總損失函數(shù)表示為：

其中 t ≡ 0 (mod α) 通過超參數(shù) α 控制監(jiān)督條件模型的時(shí)間，β 是另一個(gè)用于加權(quán)交叉熵?fù)p失的經(jīng)驗(yàn)超參數(shù)?？偠灾琈edsegdiff-v2 顯著提高了 MedsegDiff 的性能。該方法采用了新的基于 Transformer 的條件 U-Net 框架和兩種不同的條件方式，以提高擴(kuò)散模型的性能。

Anchor Condition with Gaussian Spatial Attention

與卷積層相比，Transformer 具有更強(qiáng)的表示性，但對輸入方差更敏感。為了克服這種負(fù)面影響，所以使用了錨條件運(yùn)算，如上面概述中的介紹，該運(yùn)算將條件模型的解碼分割特征（錨點(diǎn)）集成到擴(kuò)散模型的編碼器特征中。此外，還使用了高斯空間注意力來表示條件模型中給定分割特征的不確定性（概率）。

在第一個(gè)公式中，表示在錨點(diǎn)特征上應(yīng)用高斯核以進(jìn)行平滑激活，因?yàn)殄^點(diǎn)可能不完全準(zhǔn)確，且高斯核的均值和方差是可以學(xué)習(xí)的。選擇平滑 ground truth 和原始圖之間的最大值以保留最相關(guān)的信息，從而生成平滑的錨特征。在第二個(gè)公式中，將平滑錨點(diǎn)特征集成到擴(kuò)散模型中以獲得增強(qiáng)特征。首先應(yīng)用 1x1 卷積將錨特征中的通道數(shù)減少到 1（經(jīng)常作用于解碼器的最后一層）。最后，在錨點(diǎn)特征上使用 sigmoid 激活函數(shù)，將其添加到擴(kuò)散模型的每個(gè)通道中，類似于空間注意力的實(shí)現(xiàn)。

Semantic Condition with SS-Former

關(guān)于 SS-Former 的作用，我們只做簡單的總結(jié)。對比 MedSegDiff-v1，是一種將條件模型分割 embedding 集成到擴(kuò)散模型 embedding 中的新架構(gòu)，其使用頻譜空間注意力機(jī)制來解決擴(kuò)散和分割 embedding 之間的域差距。此外，注意力機(jī)制在傅里葉空間中合并語義和噪聲信息，和 MedSegDiff-v1 是類似的。

實(shí)驗(yàn)

Medsegiff-v2 是可以進(jìn)行多分類任務(wù)的，只是目前 Github 還沒有維護(hù)多分類的代碼，不過 Medsegiff-v2 的實(shí)現(xiàn)代碼已經(jīng)提交了。下圖為在腹部 CT 圖像中的多器官分割上的表現(xiàn)。

對比其他 SOTA 方法的結(jié)果如下表：

總結(jié)

關(guān)于 diffusion model 在醫(yī)學(xué)圖像分割上的應(yīng)用，其是否能適應(yīng)多分類且復(fù)雜的解刨結(jié)構(gòu)還需要進(jìn)一步探索，訓(xùn)練和推理時(shí)的效率也低于常規(guī)的有監(jiān)督神經(jīng)網(wǎng)絡(luò)，可以進(jìn)一步優(yōu)化。

審核編輯：劉清