MedSegDiff:基于擴(kuò)散概率模型的醫(yī)學(xué)圖像分割

摘要：

擴(kuò)散概率模型(Diffusion probabilistic model, DPM)是近年來(lái)計(jì)算機(jī)視覺研究的熱點(diǎn)之一。它在Imagen、Latent Diffusion Models和Stable Diffusion等圖像生成應(yīng)用中表現(xiàn)出了令人印象深刻的生成能力，引起了社區(qū)的廣泛討論。最近的許多研究還發(fā)現(xiàn)，它在許多其他視覺任務(wù)中也很有用，比如圖像去模糊、超分辨率和異常檢測(cè)。受DPM成功的啟發(fā)，我們提出了第一個(gè)基于DPM的一般醫(yī)學(xué)圖像分割模型，我們將其命名為MedSegDiff。為了增強(qiáng)DPM在醫(yī)學(xué)圖像分割中的分步區(qū)域注意力，我們提出了動(dòng)態(tài)條件編碼，該編碼為每個(gè)采樣步建立狀態(tài)自適應(yīng)條件。我們進(jìn)一步提出Feature Frequency Parser (FF-Parser)來(lái)消除高頻噪聲分量在此過(guò)程中的負(fù)面影響。我們?cè)谌N不同圖像模式的醫(yī)學(xué)分割任務(wù)上驗(yàn)證了MedSegDiff，即眼底圖像的視杯分割、MRI圖像的腦腫瘤分割和超聲圖像的甲狀腺結(jié)節(jié)分割。實(shí)驗(yàn)結(jié)果表明，MedSegDiff算法在性能上明顯優(yōu)于SOTA算法，表明了該模型的泛化性和有效性。我們的代碼發(fā)布在https://github.com/WuJunde/MedSegDiff

1 介紹

醫(yī)學(xué)圖像分割是將醫(yī)學(xué)圖像分割成有意義區(qū)域的過(guò)程。分割是許多醫(yī)學(xué)圖像分析應(yīng)用的基本步驟，如診斷、手術(shù)計(jì)劃和圖像引導(dǎo)手術(shù)。這很重要，因?yàn)樗梢宰屷t(yī)生和其他醫(yī)療專業(yè)人員更好地了解他們所看到的東西。它還可以更容易地比較圖像和跟蹤隨時(shí)間的變化。近年來(lái)，人們對(duì)醫(yī)學(xué)圖像的自動(dòng)分割方法越來(lái)越感興趣。這些方法有可能減少人工分割所需的時(shí)間和精力，并提高結(jié)果的一致性和準(zhǔn)確性。隨著深度學(xué)習(xí)技術(shù)的發(fā)展，越來(lái)越多的研究成功地將基于神經(jīng)網(wǎng)絡(luò)(NN)的模型應(yīng)用到醫(yī)學(xué)圖像分割任務(wù)中，從流行的卷積神經(jīng)網(wǎng)絡(luò)(CNN)[11]到最近的視覺變壓器(ViT)[3,22,12,28]。

最近，擴(kuò)散概率模型(diffusion probistic model, DPM)[9]作為一類強(qiáng)大的生成模型[27]得到了廣泛的應(yīng)用，它能夠生成具有高多樣性和高合成質(zhì)量的圖像。最近的大型擴(kuò)散模型，如DALLE2[17]、Imagen[19]和Stable diffusion[18]已經(jīng)顯示出令人難以置信的生成能力。擴(kuò)散模型最初應(yīng)用于不存在絕對(duì)真值的領(lǐng)域。然而，最近的研究表明，它也有效地解決了地面真相是唯一的問題，如超分辨率[20]和去模糊[24]。

受DPM最近成功的啟發(fā)，我們?cè)O(shè)計(jì)了一種獨(dú)特的基于DPM的醫(yī)學(xué)圖像分割模型。據(jù)我們所知，在一般醫(yī)學(xué)圖像分割的背景下，我們首次提出了基于dpm的不同圖像模態(tài)分割模型。我們注意到，在醫(yī)學(xué)圖像分割任務(wù)中，病變/器官往往是模糊的，很難從背景中區(qū)分出來(lái)。在這種情況下，自適應(yīng)校準(zhǔn)過(guò)程是獲得精確結(jié)果的關(guān)鍵。按照這種思路，我們提出了基于普通DPM的動(dòng)態(tài)條件編碼來(lái)設(shè)計(jì)所提出的模型，命名為MedSegDiff。需要注意的是，在迭代采樣過(guò)程中，MedSegDiff對(duì)每一步都設(shè)置了圖像先驗(yàn)條件，以便從中學(xué)習(xí)分割映射。針對(duì)自適應(yīng)區(qū)域關(guān)注，我們將當(dāng)前步驟的分割圖整合到每一步的圖像先驗(yàn)編碼中。具體實(shí)現(xiàn)是將當(dāng)前步分割掩碼與特征級(jí)先驗(yàn)圖像以多尺度方式融合。這樣，損壞的電流階掩模有助于動(dòng)態(tài)增強(qiáng)條件特征，從而提高重建精度。為了消除在此過(guò)程中損壞的給定掩模中的高頻噪聲，我們進(jìn)一步提出了特征頻率解析器(FF-Parser)來(lái)過(guò)濾傅里葉空間中的特征。在每個(gè)跳躍連接路徑上采用ff解析器進(jìn)行多尺度集成。我們?cè)谌N不同的醫(yī)學(xué)分割任務(wù)中驗(yàn)證了MedSegDiff，即光學(xué)杯分割、腦腫瘤分割和甲狀腺結(jié)節(jié)分割。這些任務(wù)的圖像有不同的模態(tài)，分別是眼底圖像、腦部CT圖像、超聲圖像。MedSegDiff在不同模態(tài)下的三種任務(wù)上均優(yōu)于先前的SOTA，表明了所提方法的泛化和有效性。簡(jiǎn)而言之，本文的貢獻(xiàn)是:

--首先提出了基于dpm的一般醫(yī)學(xué)圖像分割模型。

--提出了分步關(guān)注的動(dòng)態(tài)條件編碼策略。

--為了消除高頻分量的負(fù)面影響，提出了高頻解析器。

--SOTA在三種不同圖像模態(tài)的醫(yī)學(xué)分割任務(wù)中的性能。

2? 方法

我們基于文獻(xiàn)[9]中的擴(kuò)散模型來(lái)設(shè)計(jì)模型。擴(kuò)散模型是由正向擴(kuò)散和反向擴(kuò)散兩個(gè)階段組成的生成模型。在正演過(guò)程中，分割標(biāo)簽x0通過(guò)一系列步驟T逐漸加入高斯噪聲。在反向過(guò)程中，通過(guò)對(duì)噪聲過(guò)程進(jìn)行反向，訓(xùn)練神經(jīng)網(wǎng)絡(luò)來(lái)恢復(fù)原始數(shù)據(jù)，可以表示為:

式中θ為反工藝參數(shù)。從高斯噪聲出發(fā)，pθ(xT) = N (xT;0, In×n)，其中I為原始圖像，反向過(guò)程將潛在變量分布pθ(xT)轉(zhuǎn)換為數(shù)據(jù)分布pθ(x0)。為了與正演過(guò)程保持對(duì)稱，反向過(guò)程逐步恢復(fù)噪聲圖像，以獲得最終清晰的分割。

遵循DPM的標(biāo)準(zhǔn)實(shí)現(xiàn)，我們采用UNet作為學(xué)習(xí)網(wǎng)絡(luò)。

圖1:MedSegDiff的圖解。為清晰起見，圖中省略了時(shí)間步長(zhǎng)編碼。

圖1顯示了一個(gè)示例。為了實(shí)現(xiàn)分割，我們對(duì)步長(zhǎng)估計(jì)函數(shù)?通過(guò)原始圖像先驗(yàn)，可以表示為:

其中EIt為條件特征嵌入，在本例中為原始圖像嵌入，Ext為當(dāng)前步驟的分割映射特征嵌入。將這兩個(gè)分量相加并發(fā)送到UNet解碼器D進(jìn)行重構(gòu)。步驟索引t與添加的嵌入和解碼器功能集成在一起。

在每一個(gè)模型中，它都是使用共享學(xué)習(xí)查找表嵌入的，如下[9]

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

在大多數(shù)條件DPM中，條件先驗(yàn)將是一個(gè)唯一的給定信息。然而，醫(yī)學(xué)圖像分割是出了名的模糊對(duì)象。病變或組織通常很難從其背景中區(qū)分出來(lái)。

低對(duì)比度的圖像模式，如核磁共振成像或超聲圖像，使其更糟。只給定靜態(tài)圖像I作為每一步的條件將很難學(xué)習(xí)。為了解決這個(gè)問題，我們提出了每個(gè)步驟的動(dòng)態(tài)條件編碼。我們注意到，一方面原始圖像包含準(zhǔn)確的分割目標(biāo)信息，但難以與背景區(qū)分，另一方面，當(dāng)前步分割圖包含增強(qiáng)的目標(biāo)區(qū)域，但不準(zhǔn)確。這促使我們?nèi)フ?/p>

當(dāng)前步分割信息xt轉(zhuǎn)化為有條件的原始圖像編碼進(jìn)行相互補(bǔ)全。具體地說(shuō)，我們?cè)诠δ芗?jí)別上實(shí)現(xiàn)集成。在原始圖像編碼器中，我們利用當(dāng)前步長(zhǎng)編碼特征來(lái)增強(qiáng)其中間特征。條件特征圖mkI的每個(gè)尺度與形狀相同的xt編碼特征mkx融合，k為層的指數(shù)。這種融合是通過(guò)一種類似于注意力的機(jī)制a來(lái)實(shí)現(xiàn)的。特別是，首先將兩個(gè)特征映射應(yīng)用層歸一化并相乘以獲得親和映射。然后將親和映射與條件編碼特征相乘，增強(qiáng)關(guān)注區(qū)域，即:

其中?表示逐元素乘法，LN表示層歸一化。

該操作應(yīng)用于中間兩個(gè)階段，其中每個(gè)階段都是在ResNet34之后實(shí)現(xiàn)的卷積階段。這種策略有助于MedSegDiff動(dòng)態(tài)定位和校準(zhǔn)分割。雖然該策略是有效的，但另一個(gè)具體的問題是，積分xt嵌入會(huì)產(chǎn)生額外的高頻噪聲。為了解決這個(gè)問題，我們提出了FF-Parser來(lái)約束特征中的高頻成分。

?

2.2? FF-Parser

我們以特征集成的路徑方式連接ff解析器。它的功能是約束xt特性中與噪聲相關(guān)的組件。我們的主要思想是學(xué)習(xí)一個(gè)參數(shù)化的關(guān)注(權(quán)重)映射應(yīng)用于傅里葉空間特征。給定一個(gè)解碼器特征映射m∈RH×W ×C，我們首先沿著空間維度執(zhí)行二維FFT(快速傅立葉變換)，我們可以表示為:

式中F[·]為二維FFT。然后，我們通過(guò)將一個(gè)參數(shù)化的關(guān)注映射a∈CH×W ×C乘以m來(lái)調(diào)制m的譜:m 0 = a?m，(5)，其中?表示元素積。最后，我們采用逆FFT將m0逆回空間域:m0 = F?1[m0]。(6) FF-Parser可以看作是頻率濾波器的一種可學(xué)習(xí)版本，頻率濾波器廣泛應(yīng)用于數(shù)字圖像處理[16]。與空間關(guān)注不同，它對(duì)特定頻率的分量進(jìn)行全局調(diào)整。從而可以學(xué)會(huì)約束高頻分量進(jìn)行自適應(yīng)積分。

圖2:FF-Parser的示例。FFT表示快速傅里葉變換。

2.3? 訓(xùn)練與結(jié)構(gòu)

MedSegDiff按照DPM的標(biāo)準(zhǔn)流程進(jìn)行訓(xùn)練[9]。具體來(lái)說(shuō)，損失可以表示為:

在每次迭代中，將隨機(jī)抽取一對(duì)原始圖像Ii和分割標(biāo)簽Si進(jìn)行訓(xùn)練。迭代次數(shù)從均勻分布和?來(lái)自高斯分布。

MedSegDiff的主要架構(gòu)是一個(gè)修改后的ResUNet[26]，我們用一個(gè)ResNet編碼器和一個(gè)UNet解碼器來(lái)實(shí)現(xiàn)它。詳細(xì)的網(wǎng)絡(luò)設(shè)置如下[14]。I和xt用兩個(gè)單獨(dú)的編碼器進(jìn)行編碼。

該編碼器由三個(gè)卷積階段組成。每個(gè)階段包含幾個(gè)剩余塊。每個(gè)階段的剩余塊數(shù)量遵循ResNet34。每個(gè)殘差塊由兩個(gè)卷積塊組成，每個(gè)卷積塊由群范數(shù)和SiLU[5]活動(dòng)層和一個(gè)卷積層組成。

殘差塊通過(guò)一個(gè)線性層、SiLU激活和另一個(gè)線性層接收時(shí)間嵌入。然后將結(jié)果添加到第一個(gè)卷積塊的輸出中。將得到的EI和Ext加在一起，發(fā)送到最后的編碼階段。連接一個(gè)標(biāo)準(zhǔn)的卷積解碼器來(lái)預(yù)測(cè)最終結(jié)果。

3 實(shí)驗(yàn)

3.1數(shù)據(jù)集

我們對(duì)眼底圖像的光學(xué)杯分割、MRI圖像的腦腫瘤分割和超聲圖像的甲狀腺結(jié)節(jié)分割三種不同圖像方式的醫(yī)學(xué)任務(wù)進(jìn)行了實(shí)驗(yàn)。在REFUGE-2數(shù)據(jù)集[6]、brts -2021數(shù)據(jù)集[2]和DDTI數(shù)據(jù)集[15]上進(jìn)行青光眼、甲狀腺癌和黑色素瘤的診斷實(shí)驗(yàn)，這三個(gè)數(shù)據(jù)集分別包含1200、2000、8046個(gè)樣本。這些數(shù)據(jù)集是公開的，帶有分割和診斷標(biāo)簽。

訓(xùn)練/驗(yàn)證/測(cè)試集按照數(shù)據(jù)集的默認(rèn)設(shè)置進(jìn)行分割。

3.2實(shí)現(xiàn)細(xì)節(jié)

我們分別對(duì)MedSegDiff++、MedSegDiff- l、MedSegDiff- b和MedSegDiff- s模型的大型、大型、基本和小型變體進(jìn)行了實(shí)驗(yàn)。

在MedSegDiff- s、MedSegDiff- b、MedSegDiff- l、MedSegDiff++中，我們分別使用UNet進(jìn)行4倍、5倍、6倍、6倍的下采樣。在實(shí)驗(yàn)中，我們采用100個(gè)擴(kuò)散步驟進(jìn)行推理，這比之前的大多數(shù)研究要小得多[9,14]。除MedSegDiff++和MedSegDiff- l外，所有實(shí)驗(yàn)均在PyTorch平臺(tái)上實(shí)現(xiàn)，并在4臺(tái)Tesla P40 GPU上進(jìn)行了訓(xùn)練/測(cè)試，內(nèi)存為24GB。所有圖像都統(tǒng)一調(diào)整為256×256像素的尺寸。使用AdamW[13]優(yōu)化器以端到端方式訓(xùn)練網(wǎng)絡(luò)。MedSegDiff- b和MedSegDiff- s以32批大小進(jìn)行訓(xùn)練，MedSegDiff- l和MedSegDiff++以64批大小進(jìn)行訓(xùn)練。學(xué)習(xí)率初始設(shè)置為1 ×10?4。在推理中，所有模型都被設(shè)置為集成的25倍。我們使用STAPLE[23]算法對(duì)不同的樣本進(jìn)行融合。為了公平比較，在相同的設(shè)置下再現(xiàn)了基于擴(kuò)散的競(jìng)爭(zhēng)對(duì)手EnsemDiff[25]。

3.3 Main Results

我們將針對(duì)這三個(gè)具體任務(wù)提出的SOTA分割方法與一般醫(yī)學(xué)圖像分割方法進(jìn)行比較。主要結(jié)果如下

如圖1所示。表中，視盤/杯分割采用ResUnet[26]和BEAL[21]，腦腫瘤分割采用TransBTS[22]和EnsemDiff[25]，甲狀腺結(jié)節(jié)分割采用MTSeg[7]和UltraUNet[4]，一般醫(yī)學(xué)圖像分割采用CENet[8]、MRNet[11]、SegNet[1]、nnUNet[10]和TransUNet[3]。我們通過(guò)Dice分?jǐn)?shù)和IoU來(lái)評(píng)估分割性能。

在表1中，我們比較了各種網(wǎng)絡(luò)架構(gòu)實(shí)現(xiàn)的方法，包括CNN (ResUNet, BEAL, nnUNet, SegNet)，視覺轉(zhuǎn)換器(TransBTS, TransUNet)和DPM (EnsemDiff)。我們可以看到，先進(jìn)的網(wǎng)絡(luò)架構(gòu)通常會(huì)獲得更好的結(jié)果。例如，在光學(xué)杯分割中，基于viti的通用分割方法:TransUNet在方法:BEAL上甚至優(yōu)于基于cnn的任務(wù)。在腦腫瘤分割方面，最近提出的基于dpm的分割方法EnsemDiff優(yōu)于之前基于vita的競(jìng)爭(zhēng)對(duì)手TransBTS和TransUNet。MedSegDiff不僅采用了最近成功的DPM，而且針對(duì)一般的醫(yī)學(xué)圖像分割任務(wù)設(shè)計(jì)了相應(yīng)的策略。我們可以看到MedSegDiff在三個(gè)不同的任務(wù)上都優(yōu)于所有其他方法，這表明了MedSegDiff在不同的醫(yī)學(xué)分割任務(wù)和不同的圖像模態(tài)上的泛化。與專門用于腦腫瘤分割的基于dpm的模型(即EnsemDiff)相比，該模型在Dice上提高了2.3%，在IoU上提高了2.4%，這表明了我們獨(dú)特的技術(shù)(即動(dòng)態(tài)調(diào)節(jié)和FF-Parser)的有效性。

圖3:表1中Top-4種常用醫(yī)學(xué)圖像分割方法的視覺對(duì)比。從上到下分別為腦腫瘤分割、光學(xué)杯分割和甲狀腺結(jié)節(jié)分割。

在圖像上，使它們難以被人眼識(shí)別。通過(guò)與這些計(jì)算機(jī)輔助方法的比較，可以明顯地看出，該方法生成的分割圖比其他方法更準(zhǔn)確，特別是對(duì)于模糊區(qū)域。將DPM與所提出的動(dòng)態(tài)調(diào)節(jié)和FF-Parser相結(jié)合，可以更好地定位和校準(zhǔn)低對(duì)比度或模糊圖像的分割。

表1:MedSegDiff與SOTA分割方法的比較。最好的結(jié)果用粗體表示?；疑尘氨硎踞槍?duì)該/這些特定任務(wù)提出的方法。

3.4 消融研究

我們進(jìn)行了全面的消融研究，以驗(yàn)證所提出的動(dòng)態(tài)調(diào)節(jié)和FF-Parser的有效性。結(jié)果如表2所示，其中Dy-Cond表示動(dòng)態(tài)調(diào)節(jié)。我們通過(guò)Dice分?jǐn)?shù)(%)來(lái)評(píng)估這三個(gè)任務(wù)的表現(xiàn)。從表格中，我們可以看到Dy-Cond比普通DPM獲得了相當(dāng)大的改進(jìn)。在區(qū)域定位很重要的情況下，即光學(xué)杯分割，提高了2.1%。在圖像對(duì)比度較低的情況下，如腦腫瘤和甲狀腺結(jié)節(jié)分割，分別提高1.6%和1.8%。它表明，對(duì)于這兩種情況，Dy-Cond通常是DPM的有效策略。在Dy-Cond上建立的FF-Parser減輕了高頻噪聲，從而進(jìn)一步優(yōu)化了分割結(jié)果。它幫助MedSegDiff進(jìn)一步提高了近1%的性能，并在所有三個(gè)任務(wù)上都取得了最佳成績(jī)。

4 結(jié)論

本文提出了一種基于dpm的通用醫(yī)學(xué)圖像分割方案MedSegDiff。我們提出了兩種新的技術(shù)來(lái)保證其性能，即動(dòng)態(tài)條件編碼和ff解析器。

對(duì)三種不同圖像模態(tài)的醫(yī)學(xué)圖像分割任務(wù)進(jìn)行了對(duì)比實(shí)驗(yàn)，結(jié)果表明我們的模型優(yōu)于以往的SOTA。作為DPM在一般醫(yī)學(xué)圖像分割中的第一個(gè)應(yīng)用，我們相信MedSegDiff將成為未來(lái)研究的重要基準(zhǔn)。文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-784445.html

到了這里，關(guān)于MedSegDiff: Medical Image Segmentation withDiffusion Probabilistic Model的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！