葉晉豫,李 嬌,鄧紅霞,張瑞欣,李海芳
(太原理工大學(xué) 信息與計(jì)算機(jī)學(xué)院,山西 晉中 030600)
關(guān)于醫(yī)學(xué)圖像分割[1,2]的研究方法層出不窮,最先進(jìn)的醫(yī)學(xué)圖像分割框架大多基于U-Net[3]或其變體[4]。遵循這一技術(shù)路線有許多算法,如U-Net的各種變體[5-8],通過跳躍連接保存了編碼器層的單粒度信息,但忽略了豐富的多尺度空間信息,使得大量邊緣信息丟失,這極大地影響了其在分割任務(wù)中的性能。
最近,受Transformer在自然語言處理方面取得巨大成功的啟發(fā)[9],其也應(yīng)用于圖像處理領(lǐng)域[10,11]。利用Transformer的關(guān)注全局的關(guān)鍵優(yōu)勢(shì)結(jié)合CNN的局部特性來構(gòu)建一個(gè)更優(yōu)的分割網(wǎng)絡(luò)。目前,初步研究已經(jīng)嘗試將Transformer在醫(yī)學(xué)圖像分割領(lǐng)域應(yīng)用[12-14]。一個(gè)常見的策略是將視覺Transformer直接應(yīng)用于醫(yī)學(xué)圖像。ViT[15](vision Transformer)是第一個(gè)純粹基于Transformer的架構(gòu)。Touvron等[16]描述了幾種培訓(xùn)策略,使ViT在ImageNet數(shù)據(jù)集上訓(xùn)練良好。Transformer也被應(yīng)用于醫(yī)學(xué)圖像分割領(lǐng)域[17-22],利用了Transformer和CNN的特性來提高模型的分割能力。例如,Chen等[2]和Hata等[1]將二維或三維解碼器添加到類似ViT的編碼器中。Hu等[11]使用類似Swin Transfomer[23]的編碼器-解碼器結(jié)構(gòu)用于分割?,F(xiàn)在,CNN和Transformer架構(gòu)都展示了各自的獨(dú)特優(yōu)勢(shì),結(jié)合兩種架構(gòu)的優(yōu)勢(shì)進(jìn)行綜合分析是非常有意義的。
所以,本文提出了一種混合模型SwinEA(Swin Transformer fusion edge-aware medical image segmentation network),即上下文金字塔結(jié)構(gòu)結(jié)合Swin Transformer的網(wǎng)絡(luò),基于Swin Transformer的邊緣感知上下文金字塔融合多尺度特征的醫(yī)學(xué)圖像分割網(wǎng)絡(luò)。該網(wǎng)絡(luò)利用上下文金字塔結(jié)構(gòu)融合捕捉豐富的多尺度上下文信息,計(jì)算相鄰像素間的局部相關(guān)性的卷積,用于提取邊緣和角落等局部特征;并且結(jié)合設(shè)計(jì)用于序列到序列預(yù)測(cè)的Swin Transformer,利用其具有的先天性全局自注意力機(jī)制的替代體系結(jié)構(gòu),用來解決CNN的歸納偏置問題。
本文提出的SwinEA的總體架構(gòu)如圖1(a)所示。SwinEA由編碼器、解碼器和邊緣感知模塊組成。其中編碼器、解碼器基本單元是Swin Transformer塊。
圖1 SwinEA模型框架
首先,編碼器要將輸入的醫(yī)學(xué)圖像分割為4×4的非重疊補(bǔ)丁塊,再經(jīng)過一個(gè)線性嵌入層(linear embedding)將特征維度投影到任意維度(表示為C),轉(zhuǎn)換后的補(bǔ)丁塊通過4層Swin Transformer塊和補(bǔ)丁合并層以生成分層不同尺度的特征。具體來說,補(bǔ)丁合并層(patch merging)負(fù)責(zé)下采樣和增加維度,Swin Transformer塊負(fù)責(zé)特征的學(xué)習(xí)。受U-Net[3]的啟發(fā),設(shè)計(jì)了一個(gè)基于Swin Transformer的對(duì)稱解碼器,該解碼器由Swin Transformer塊和補(bǔ)丁擴(kuò)展層(patch expanding)構(gòu)成,與補(bǔ)丁合并層相反,補(bǔ)丁擴(kuò)展層專門設(shè)計(jì)用于執(zhí)行上采樣,補(bǔ)丁擴(kuò)展層將相鄰維度的特征圖重塑為一個(gè)2倍分辨率的上采樣特征圖,再使用最后的補(bǔ)丁擴(kuò)展層進(jìn)行4倍上采樣,恢復(fù)圖像的分辨率并將特征映射到輸入的分辨率大小(W×H),然后再將這些上采樣的特征輸出像素級(jí)分割預(yù)測(cè)。
在編碼器解碼器結(jié)構(gòu)的中間,本文不以跳躍連接進(jìn)行連接,而是引入了一種上下文金字塔模塊進(jìn)行連接,將提取的上下文特征融合多尺度的特征通過邊緣感知(edge-aware module)模塊,以彌補(bǔ)損失多尺度空間信息,更好提取邊緣和角落等局部特征,通過重構(gòu)跳躍連接為解碼器提供不同級(jí)別的全局上下文信息。
Swin Transformer塊與傳統(tǒng)的多頭自注意(MSA)模塊不同,它是基于滑動(dòng)窗口構(gòu)造的,由兩個(gè)連續(xù)的Transformer組成,結(jié)構(gòu)如圖1(b)每一個(gè)Swin Transformer模塊由LN(LayerNorm)層、多頭自我注意模塊、殘差連接和具有GELU非線性的2層MLP(多層感知器)。兩個(gè)連續(xù)的Transformer塊分別采用基于窗口的多頭自注意(W-MSA)模塊和基于移動(dòng)窗口的多頭自注意(SW-MSA)模塊。基于這樣的窗口劃分機(jī)制,Swin Transformer塊可以表示為
(1)
(2)
(3)
(4)
self-Attention的計(jì)算方式如下
(5)
補(bǔ)丁合并層:輸入的補(bǔ)丁分為4部分,并通過補(bǔ)丁合并層連接在一起。通過這樣的處理,特征分辨率將被2倍下采樣。由于連接操作導(dǎo)致特征維度增加了4倍,因此在連接的特征上應(yīng)用線性層將特征維度統(tǒng)一為原始維度的2倍。補(bǔ)丁合并層無需卷積或插值操作即可實(shí)現(xiàn)下采樣和特征維數(shù)的增加。
對(duì)應(yīng)于編碼器構(gòu)建一個(gè)對(duì)稱的解碼器,它也是基于Swin Transformer塊構(gòu)建的。因此,對(duì)應(yīng)于編碼器中使用的補(bǔ)丁合并層,我們?cè)诮獯a器中構(gòu)建了一個(gè)補(bǔ)丁擴(kuò)展層,用于提取深度特征,進(jìn)行上采樣。補(bǔ)丁擴(kuò)展層將相鄰維度的特征圖重塑為更高分辨率的特征圖(2倍上采樣)相應(yīng)地將特征維度減少到原始維度的一半。補(bǔ)丁擴(kuò)展層將相鄰維度的特征地圖重新塑造成一個(gè)2倍上采樣分辨率的特征圖。最后,利用網(wǎng)絡(luò)結(jié)構(gòu)最后一個(gè)補(bǔ)丁擴(kuò)展層進(jìn)行4倍上采樣,恢復(fù)圖像的分辨率到輸入分辨率,然后經(jīng)過線性投影層,將這些上采樣特征輸出像素級(jí)的分割預(yù)測(cè)。
為了解決固定尺度卷積操作在分割中的局限性,本文設(shè)計(jì)了邊緣感知模塊,通過不同尺度的空間信息提高了網(wǎng)絡(luò)結(jié)構(gòu)的魯棒性。此外,U形網(wǎng)絡(luò)中通常使用的跳躍連接,在接收字段不匹配會(huì)引入不相關(guān)的混亂和有語義差距。在本文中,提出全局上下文金字塔結(jié)構(gòu)的邊緣感知(EA)模塊來解決這些問題,如圖1(c)所示。在EA模塊,會(huì)對(duì)跳躍連接進(jìn)行重構(gòu)。在解碼過程中,淺層特征展示詳細(xì)的邊界信息,同時(shí)也會(huì)帶來一些背景噪音。因此,邊緣感知模塊用于提取邊緣特征,并進(jìn)一步引導(dǎo)解碼器,同時(shí)抑制淺層噪聲并細(xì)化物體的輪廓。
在邊緣感知模塊中,利用不同大小的卷積核運(yùn)算在多個(gè)分支中提取不同尺度的特征,通過融合不同粒度的空間信息,使網(wǎng)絡(luò)能夠?qū)W習(xí)更多的上下文信息?;赨形結(jié)構(gòu),在編碼器和解碼器之間設(shè)計(jì)了多個(gè)邊緣感知(EA)模塊,旨在通過重構(gòu)跳轉(zhuǎn)連接為解碼器提供不同級(jí)別的全局上下文多尺度信息。眾所周知,高層特征表達(dá)更多的語義信息,而淺層特征攜帶更多細(xì)節(jié)信息。因此,淺層深度主干的特征用于產(chǎn)生邊緣特征,所以提出的邊緣感知模塊可以盡可能多提取淺層特征用于產(chǎn)生邊緣特征。在EA模塊中,將這一階段的特征圖與所有更高一級(jí)階段的特征圖相結(jié)合,重構(gòu)跳過連接。
如圖1(c)所示,以第二階段的邊緣感知模塊為例,將各階段的不同尺度的特征映射到與第二階段相同的通道空間中的3×3卷積。接下來,生成的特征映射第三階段特征地圖和第四階段特征地圖上采樣到與第二階段特征地圖相同的大小并連接。然后為了從中提取全局上下文信息不同層次的特征圖,利用3種可分離卷積[24](Dsconv@1,Dsconv@2,Dsconv@4),它們具有不同的膨脹率(1,2和4)進(jìn)行平行使用,并且在這里采用可分離卷積可以減少模型參數(shù)。值得注意的是平行路徑和擴(kuò)張的數(shù)量速率隨融合階段的數(shù)目而變化。最后,通過卷積得到最終的特征地圖。
以上,可以總結(jié)出各個(gè)階段的EA模塊(為了簡化公式,忽略正則卷積)
(6)
其中,Gk為插入第k階段,F(xiàn)k為編碼器中第k階段的特征圖,?2i-k為上采樣操作速率為2i-k,c表示串接操作和Dsconv@2i-k為可分離的擴(kuò)張卷積膨脹率為2i-k。
通過引入多個(gè)邊緣感知模塊在編碼器和解碼器之間,全局語義高層次的空間信息流可以逐層引導(dǎo)到不同階段。
醫(yī)學(xué)圖像分割的一個(gè)主要挑戰(zhàn)是分類分布不平衡。為了進(jìn)一步優(yōu)化我們的模型,也為了更好地評(píng)估標(biāo)簽輪廓和預(yù)測(cè)的輪廓之間的相似性,我們將組合一個(gè)新的損失函數(shù),由交叉熵?fù)p失和Dice Loss組成。
交叉熵?fù)p失函數(shù)lce計(jì)算如下
(7)
其中,σθ(x) 是深度神經(jīng)網(wǎng)絡(luò)的輸出,x是網(wǎng)絡(luò)的輸入,而y代表標(biāo)簽。θ表示神經(jīng)網(wǎng)絡(luò)的可訓(xùn)練參數(shù)。
Dice Loss函數(shù)lDice計(jì)算如下
(8)
其中,|X∩Y| 表示X和Y集合的交集,|X| 表示元素個(gè)數(shù),|Y| 表示分割任務(wù)中的標(biāo)簽及預(yù)測(cè)結(jié)果。
我們采用由Dice損失lDice和交叉熵?fù)p失函數(shù)lce計(jì)成的聯(lián)合損失來執(zhí)行所有分割任務(wù)。
公式如下
ltotal=λlDice+(1-λ)lce
(9)
其中,λ是Dice損失和交叉熵之間的權(quán)衡損失。為了公平比較,我們實(shí)驗(yàn)中的每個(gè)單獨(dú)的任務(wù)都使用相同的損失函數(shù)來執(zhí)行。
Synapse多器官分割數(shù)據(jù)集:該數(shù)據(jù)集包括3779張軸向腹部臨床CT圖像。訓(xùn)練集包括18個(gè)樣本,測(cè)試集包括12個(gè)樣本。其中每個(gè)CT圖像樣本由85~198個(gè)512×512像素的切片組成,體素空間分辨率為 ([0.54~0.54]×[0.98~0.98]×[2.5~5.0]) mm3。 本文對(duì)8個(gè)腹部臟器(主動(dòng)脈、膽囊、左腎、右腎、肝、胰、脾、胃)進(jìn)行實(shí)驗(yàn)驗(yàn)證得到其平均DSC和平均Hausdorff距離(HD)。
ACDC(心臟自動(dòng)診斷挑戰(zhàn)):ACDC是一個(gè)公共心臟MRI數(shù)據(jù)集,包括100例樣本。一系列短軸切片,覆蓋心臟從左心室基部到心尖,厚度為5~8 mm。短軸平面內(nèi)空間分辨率從0.83到1.75 mm2/pixel,對(duì)應(yīng)的標(biāo)簽包括左心室(LV)、右心室(RV)和心肌(MYO)。數(shù)據(jù)集分為70個(gè)訓(xùn)練樣本(1930個(gè)軸向切片)、10個(gè)驗(yàn)證樣本和20個(gè)測(cè)試樣本。
2.2.1 環(huán)境
SwinEA網(wǎng)絡(luò)是基于Python 3.6和Pytorch 1.7.0實(shí)現(xiàn)的。對(duì)于所有的訓(xùn)練案例,都使用數(shù)據(jù)增強(qiáng),如翻轉(zhuǎn)和旋轉(zhuǎn),以增加數(shù)據(jù)的多樣性。輸入的圖像大小設(shè)置為224×224,patch大小設(shè)置為4。本文在內(nèi)存為32 GB的Nvidia V100 GPU、Linux操作系統(tǒng)上訓(xùn)練本文的模型。
在訓(xùn)練期間,批量大小為24,使用的是SGD優(yōu)化器,動(dòng)量為0.9,權(quán)值衰減為1e-4,對(duì)本文的反向傳播模型進(jìn)行優(yōu)化。實(shí)驗(yàn)設(shè)置300個(gè)epoch。
2.2.2 評(píng)估指標(biāo)
本文使用兩種類型的指標(biāo)來評(píng)估所有模型。
評(píng)價(jià)預(yù)測(cè)與分割標(biāo)簽重疊程度的Dice相似性系數(shù)(DSC)
(10)
其中,P為預(yù)測(cè)的分割圖,G為標(biāo)簽。
HD(Hausdorff distance),測(cè)量兩個(gè)分割貼圖之間最大的對(duì)稱距離
(11)
式中:d(·) 為歐幾里得距離,sup和inf分別為上極值和下極值。P為預(yù)測(cè)的分割圖,G為標(biāo)簽。本文采用95%HD,以消除一個(gè)非常小的子集的異常值的影響。
2.3.1 在Synapse多器官CT數(shù)據(jù)集上的驗(yàn)證
在Synapse多器官CT數(shù)據(jù)集上,本文提出的SwinEA與之前較新的方法的比較見表1,其中2-3列分別為平均DSC和HD評(píng)價(jià)指標(biāo)結(jié)果,后8列為8個(gè)器官的具體DSC指標(biāo)結(jié)果。實(shí)驗(yàn)結(jié)果表明,本文的基于Swin Transformer融合多尺度上下文金字塔的方法分割精度達(dá)到了79.65%(DSC↑)和20.97%(HD↓)。與最近的TransUnet、Swin-Unet、MTUnet[25]方法相比,本文的算法在DSC評(píng)價(jià)指標(biāo)和HD評(píng)價(jià)指標(biāo)上有著提升,這表明本文的方法可以實(shí)現(xiàn)更好的分割邊緣預(yù)測(cè)。
表1 不同算法在Synapse多器官分割數(shù)據(jù)集的評(píng)價(jià)指標(biāo)對(duì)比結(jié)果
不同方法對(duì)Synapse多器官CT數(shù)據(jù)集的分割結(jié)果如圖2所示。從圖中白色框標(biāo)出的地方可以看出,本文方法在一些分割邊緣上比其它一些先進(jìn)方法均有著分割邊緣清晰,區(qū)域明顯的效果。因此,在本文工作中,驗(yàn)證了基于Swin Transformer的上下文金字塔融合多尺度特征結(jié)合的醫(yī)學(xué)圖像分割網(wǎng)絡(luò)通過邊緣融合模塊捕捉了豐富的多尺度上下文信息,計(jì)算相鄰像素間的局部相關(guān)性的卷積,在提取邊緣和角落等局部特征方面表現(xiàn)得很好,也獲得更好的分割結(jié)果。
圖2 不同算法在Synapse數(shù)據(jù)集上的分割對(duì)比結(jié)果
表1對(duì)不同算法的實(shí)驗(yàn)結(jié)果進(jìn)行了定量評(píng)估。U-Net為原始深度卷積網(wǎng)絡(luò)的方法;Transunet為生成模型的編碼結(jié)構(gòu)使用Transformer,解碼器中使用卷積結(jié)構(gòu);SwinUnet為編解碼器結(jié)構(gòu)均使用純的Transformer方法;SwinEA為本文提出的方法。從表中可以看出,通過使用上下文金字塔結(jié)構(gòu)與Swin Transformer方法共同結(jié)合,DSC提升到79.65%,HD提升到20.97%。數(shù)據(jù)方面也驗(yàn)證,在Swin Transformer網(wǎng)絡(luò)中引入上下文金字塔機(jī)制,可以有效且精確地針對(duì)邊緣完成腹部多器官分割的任務(wù)。
為了驗(yàn)證該模型的性能,與分割標(biāo)簽通過三維重建進(jìn)行了比較。將預(yù)處理的切片結(jié)果還原為nii格式的三維數(shù)據(jù)集,比較三維下的分割結(jié)果。
圖3為模型在Synapse數(shù)據(jù)集不同層下的分割結(jié)果。第一列為Synapse數(shù)據(jù)集切片后的原始圖像,第二列為標(biāo)簽結(jié)果,第三列為本文提出的網(wǎng)絡(luò)模型的分割結(jié)果,第四列為本文進(jìn)行三維重建的分割結(jié)果。從圖3中可以看出,與第二列專家分割的結(jié)果進(jìn)行對(duì)比,第三列本文的實(shí)驗(yàn)分割結(jié)果在細(xì)節(jié)上分割精確,基本上與標(biāo)簽分割結(jié)果相似。并且將切片數(shù)據(jù)進(jìn)行三維重建之后,邊緣圓滑。
圖3 三維重建的Synapse分割結(jié)果
2.3.2 在ACDC數(shù)據(jù)集上的驗(yàn)證
為了驗(yàn)證模型的泛化性,將SwinEA在ACDC 數(shù)據(jù)集來進(jìn)行醫(yī)學(xué)圖像分割。表2中總結(jié)了不同算法下的分割實(shí)驗(yàn)數(shù)據(jù)結(jié)果,其中第2列為平均DSC值,后3列為右心室、心肌、左心室的具體DSC值。使用MRI模式的圖像數(shù)據(jù)作為輸入,SwinEA仍然能夠達(dá)到優(yōu)秀的性能,平均準(zhǔn)確率達(dá)到91.34%。這表明我們的方法具有良好的泛化能力和穩(wěn)健性。如圖4中,第一行為ACDC數(shù)據(jù)集切片后的原始圖像,第二行為分割標(biāo)簽,第三行為本文提出的網(wǎng)絡(luò)模型的分割結(jié)果,第四行為本文三維重建的分割結(jié)果。從圖4中可以看出,通過與第二行標(biāo)簽結(jié)果的結(jié)果進(jìn)行對(duì)比,第三行本文的實(shí)驗(yàn)分割結(jié)果在細(xì)節(jié)上分割精確,基本上與標(biāo)簽分割結(jié)果相似,并且將切片數(shù)據(jù)進(jìn)行三維重建之后,可以看到分割邊緣圓滑,分割區(qū)域清晰。
表2 不同算法在ACDC心臟分割數(shù)據(jù)集的評(píng)價(jià)指標(biāo)結(jié)果
圖4 不同算法在ACDC 數(shù)據(jù)集的分割結(jié)果
2.3.3 消融實(shí)驗(yàn)
為了探究不同因素對(duì)模型性能的影響,本文在Synapse數(shù)據(jù)集上進(jìn)行了消融實(shí)驗(yàn)研究。
(1)邊緣感知模塊的位置的影響
本文的EA模塊被添加到1/4、1/8和1/16分辨率刻度的位置在圖中分別代表邊緣感知模塊1、邊緣感知模塊2、邊緣感知模塊3,探討不同位置對(duì)分割效果的影響。從圖5中可以看出模型的分割效果,邊緣感知模塊1、邊緣感知模塊2代表的黑色白色部分大部分表現(xiàn)比邊緣感知模塊3的灰色更準(zhǔn)確。因此,為了使模型更加魯棒,本文將不使用邊緣感知模塊3的情況。
圖5 邊緣感知模塊的位置的影響
在實(shí)驗(yàn)過程中發(fā)現(xiàn)此模塊在不同位置上的表現(xiàn)不同,位于淺層時(shí)效果更好,可以更好提取細(xì)節(jié)空間信息。在添加塊數(shù)上,塊數(shù)更多會(huì)隨之帶來提取信息過于冗余,表現(xiàn)一般,平衡位置與添加塊數(shù)的平衡,在實(shí)驗(yàn)中驗(yàn)證添加兩塊EA模塊,且位于淺層位置可以達(dá)到一個(gè)更優(yōu)的結(jié)果。
(2)邊緣感知模塊的數(shù)量的影響
本文的邊緣感知模塊數(shù)量的討論結(jié)合位置限制,所以最多可以有3塊,最少可以為0塊。探討數(shù)量對(duì)分割性能的影響。從表3中可以數(shù)量看出模型的分割性能,3個(gè)邊緣感知模塊造成冗余,效果一般。兩塊效果普遍比一塊效果更優(yōu)。因此,為了使模型更加健壯,本文將采用兩塊邊緣感知模塊。
表3 邊緣感知模塊的數(shù)量的影響
本文通過消融實(shí)驗(yàn),得出關(guān)于邊緣感知模塊模塊數(shù)量影響的結(jié)論,添加3塊邊緣感知模塊模塊時(shí)造成提取信息過于冗余,造成效果不佳,而針對(duì)于分割邊緣低維度信息特征更加有用,所以兩塊時(shí),添加處于低維度特征的邊緣感知模塊1、邊緣感知模塊2模塊時(shí),效果更好,將高分辨率的低水平特征和低分辨率的高水平特征融合起來,幫助劃分出詳細(xì)清晰的分割邊界。
(3)損失函數(shù)的影響
本文采用的損失函數(shù)包括常用在分割中的Dice損失和交叉熵?fù)p失,兩項(xiàng)損失進(jìn)行消融實(shí)驗(yàn)后,實(shí)驗(yàn)結(jié)果見表4,結(jié)果驗(yàn)證兩項(xiàng)損失聯(lián)合訓(xùn)練的方法可以得到更好的分割效果。針對(duì)分割邊緣的清晰與精度,這兩種損失提供了損失的信息,更好解決了固定尺度卷積操作在目標(biāo)分割中的局限性。
表4 損失函數(shù)的影響
本文提出了一種針對(duì)分割邊緣的利用Swin Transformer融合邊緣感知的醫(yī)學(xué)圖像分割網(wǎng)絡(luò)(SwinEA),設(shè)計(jì)了一個(gè)基于上下文金字塔的邊緣感知(Edge-Aware)模塊,用于融合全局多尺度的上下文信息,主要針對(duì)邊緣和角落等局部特征。并且側(cè)重于解決U形網(wǎng)絡(luò)中全局多尺度上下文信息捕獲和整合的弱點(diǎn)。新型邊緣感知模塊以上下文金字塔為基礎(chǔ)的邊界感知模塊插入U(xiǎn)形框架,開發(fā)和融合豐富的全局多尺度上下文信息。將高分辨率的低水平特征和低水平的特征融合起來。高分辨率的低層次特征和低分辨率的高層次特征來幫助劃定詳細(xì)的分割邊緣。
本文對(duì)不同類型的醫(yī)學(xué)圖像分割任務(wù)進(jìn)行了綜合實(shí)驗(yàn),以驗(yàn)證所提出的SwinEA的有效性和泛化性,包括腹部多器官分割任務(wù)和ACDC數(shù)據(jù)集。本文方法可以通過進(jìn)一步的處理達(dá)到更好的性能,未來可以擴(kuò)展應(yīng)用到3D圖像以及其它醫(yī)學(xué)圖像分割任務(wù),這也是我們近期的工作。