SwinEA：融合邊緣感知的醫學圖像分割網絡

葉晉豫，李 嬌，鄧紅霞，張瑞欣，李海芳

(太原理工大學 信息與計算機學院，山西 晉中 030600)

0 引 言

關于醫學圖像分割[1，2]的研究方法層出不窮，最先進的醫學圖像分割框架大多基于U-Net[3]或其變體[4]。遵循這一技術路線有許多算法，如U-Net的各種變體[5-8]，通過跳躍連接保存了編碼器層的單粒度信息，但忽略了豐富的多尺度空間信息，使得大量邊緣信息丟失，這極大地影響了其在分割任務中的性能。

最近，受Transformer在自然語言處理方面取得巨大成功的啟發[9]，其也應用于圖像處理領域[10，11]。利用Transformer的關注全局的關鍵優勢結合CNN的局部特性來構建一個更優的分割網絡。目前，初步研究已經嘗試將Transformer在醫學圖像分割領域應用[12-14]。一個常見的策略是將視覺Transformer直接應用于醫學圖像。ViT[15](vision Transformer)是第一個純粹基于Transformer的架構。Touvron等[16]描述了幾種培訓策略，使ViT在ImageNet數據集上訓練良好。Transformer也被應用于醫學圖像分割領域[17-22]，利用了Transformer和CNN的特性來提高模型的分割能力。例如，Chen等[2]和Hata等[1]將二維或三維解碼器添加到類似ViT的編碼器中。Hu等[11]使用類似Swin Transfomer[23]的編碼器-解碼器結構用于分割。現在，CNN和Transformer架構都展示了各自的獨特優勢，結合兩種架構的優勢進行綜合分析是非常有意義的。

所以，本文提出了一種混合模型SwinEA(Swin Transformer fusion edge-aware medical image segmentation network)，即上下文金字塔結構結合Swin Transformer的網絡，基于Swin Transformer的邊緣感知上下文金字塔融合多尺度特征的醫學圖像分割網絡。該網絡利用上下文金字塔結構融合捕捉豐富的多尺度上下文信息，計算相鄰像素間的局部相關性的卷積，用于提取邊緣和角落等局部特征；并且結合設計用于序列到序列預測的Swin Transformer，利用其具有的先天性全局自注意力機制的替代體系結構，用來解決CNN的歸納偏置問題。

1 本文方法

本文提出的SwinEA的總體架構如圖1(a)所示。SwinEA由編碼器、解碼器和邊緣感知模塊組成。其中編碼器、解碼器基本單元是Swin Transformer塊。

圖1 SwinEA模型框架

首先，編碼器要將輸入的醫學圖像分割為4×4的非重疊補丁塊，再經過一個線性嵌入層(linear embedding)將特征維度投影到任意維度(表示為C)，轉換后的補丁塊通過4層Swin Transformer塊和補丁合并層以生成分層不同尺度的特征。具體來說，補丁合并層(patch merging)負責下采樣和增加維度，Swin Transformer塊負責特征的學習。受U-Net[3]的啟發，設計了一個基于Swin Transformer的對稱解碼器，該解碼器由Swin Transformer塊和補丁擴展層(patch expanding)構成，與補丁合并層相反，補丁擴展層專門設計用于執行上采樣，補丁擴展層將相鄰維度的特征圖重塑為一個2倍分辨率的上采樣特征圖，再使用最后的補丁擴展層進行4倍上采樣，恢復圖像的分辨率并將特征映射到輸入的分辨率大小(W×H)，然后再將這些上采樣的特征輸出像素級分割預測。

在編碼器解碼器結構的中間，本文不以跳躍連接進行連接，而是引入了一種上下文金字塔模塊進行連接，將提取的上下文特征融合多尺度的特征通過邊緣感知(edge-aware module)模塊，以彌補損失多尺度空間信息，更好提取邊緣和角落等局部特征，通過重構跳躍連接為解碼器提供不同級別的全局上下文信息。

1.1 Swin Transformer模塊

Swin Transformer塊與傳統的多頭自注意(MSA)模塊不同，它是基于滑動窗口構造的，由兩個連續的Transformer組成，結構如圖1(b)每一個Swin Transformer模塊由LN(LayerNorm)層、多頭自我注意模塊、殘差連接和具有GELU非線性的2層MLP(多層感知器)。兩個連續的Transformer塊分別采用基于窗口的多頭自注意(W-MSA)模塊和基于移動窗口的多頭自注意(SW-MSA)模塊。基于這樣的窗口劃分機制，Swin Transformer塊可以表示為

(1)

(2)

(3)

(4)

self-Attention的計算方式如下

(5)

1.2 編碼器

補丁合并層：輸入的補丁分為4部分，并通過補丁合并層連接在一起。通過這樣的處理，特征分辨率將被2倍下采樣。由于連接操作導致特征維度增加了4倍，因此在連接的特征上應用線性層將特征維度統一為原始維度的2倍。補丁合并層無需卷積或插值操作即可實現下采樣和特征維數的增加。

1.3 解碼器

對應于編碼器構建一個對稱的解碼器，它也是基于Swin Transformer塊構建的。因此，對應于編碼器中使用的補丁合并層，我們在解碼器中構建了一個補丁擴展層，用于提取深度特征，進行上采樣。補丁擴展層將相鄰維度的特征圖重塑為更高分辨率的特征圖(2倍上采樣)相應地將特征維度減少到原始維度的一半。補丁擴展層將相鄰維度的特征地圖重新塑造成一個2倍上采樣分辨率的特征圖。最后，利用網絡結構最后一個補丁擴展層進行4倍上采樣，恢復圖像的分辨率到輸入分辨率，然后經過線性投影層，將這些上采樣特征輸出像素級的分割預測。

1.4 邊緣感知模塊

為了解決固定尺度卷積操作在分割中的局限性，本文設計了邊緣感知模塊，通過不同尺度的空間信息提高了網絡結構的魯棒性。此外，U形網絡中通常使用的跳躍連接，在接收字段不匹配會引入不相關的混亂和有語義差距。在本文中，提出全局上下文金字塔結構的邊緣感知(EA)模塊來解決這些問題，如圖1(c)所示。在EA模塊，會對跳躍連接進行重構。在解碼過程中，淺層特征展示詳細的邊界信息，同時也會帶來一些背景噪音。因此，邊緣感知模塊用于提取邊緣特征，并進一步引導解碼器，同時抑制淺層噪聲并細化物體的輪廓。

在邊緣感知模塊中，利用不同大小的卷積核運算在多個分支中提取不同尺度的特征，通過融合不同粒度的空間信息，使網絡能夠學習更多的上下文信息。基于U形結構，在編碼器和解碼器之間設計了多個邊緣感知(EA)模塊，旨在通過重構跳轉連接為解碼器提供不同級別的全局上下文多尺度信息。眾所周知，高層特征表達更多的語義信息，而淺層特征攜帶更多細節信息。因此，淺層深度主干的特征用于產生邊緣特征，所以提出的邊緣感知模塊可以盡可能多提取淺層特征用于產生邊緣特征。在EA模塊中，將這一階段的特征圖與所有更高一級階段的特征圖相結合，重構跳過連接。

如圖1(c)所示，以第二階段的邊緣感知模塊為例，將各階段的不同尺度的特征映射到與第二階段相同的通道空間中的3×3卷積。接下來，生成的特征映射第三階段特征地圖和第四階段特征地圖上采樣到與第二階段特征地圖相同的大小并連接。然后為了從中提取全局上下文信息不同層次的特征圖，利用3種可分離卷積[24](Dsconv@1，Dsconv@2，Dsconv@4)，它們具有不同的膨脹率(1，2和4)進行平行使用，并且在這里采用可分離卷積可以減少模型參數。值得注意的是平行路徑和擴張的數量速率隨融合階段的數目而變化。最后，通過卷積得到最終的特征地圖。

以上，可以總結出各個階段的EA模塊(為了簡化公式，忽略正則卷積)

(6)

其中，Gk為插入第k階段，Fk為編碼器中第k階段的特征圖，?2i-k為上采樣操作速率為2i-k，c表示串接操作和Dsconv@2i-k為可分離的擴張卷積膨脹率為2i-k。

通過引入多個邊緣感知模塊在編碼器和解碼器之間，全局語義高層次的空間信息流可以逐層引導到不同階段。

1.5 損失函數

醫學圖像分割的一個主要挑戰是分類分布不平衡。為了進一步優化我們的模型，也為了更好地評估標簽輪廓和預測的輪廓之間的相似性，我們將組合一個新的損失函數，由交叉熵損失和Dice Loss組成。

交叉熵損失函數lce計算如下

(7)

其中，σθ(x) 是深度神經網絡的輸出，x是網絡的輸入，而y代表標簽。θ表示神經網絡的可訓練參數。

Dice Loss函數lDice計算如下

(8)

其中，|X∩Y| 表示X和Y集合的交集，|X| 表示元素個數，|Y| 表示分割任務中的標簽及預測結果。

我們采用由Dice損失lDice和交叉熵損失函數lce計成的聯合損失來執行所有分割任務。

公式如下

ltotal=λlDice+(1-λ)lce

(9)

其中，λ是Dice損失和交叉熵之間的權衡損失。為了公平比較，我們實驗中的每個單獨的任務都使用相同的損失函數來執行。

2 實驗和結果分析

2.1 實驗數據集

Synapse多器官分割數據集：該數據集包括3779張軸向腹部臨床CT圖像。訓練集包括18個樣本，測試集包括12個樣本。其中每個CT圖像樣本由85～198個512×512像素的切片組成，體素空間分辨率為 ([0.54～0.54]×[0.98～0.98]×[2.5～5.0]) mm3。 本文對8個腹部臟器(主動脈、膽囊、左腎、右腎、肝、胰、脾、胃)進行實驗驗證得到其平均DSC和平均Hausdorff距離(HD)。

ACDC(心臟自動診斷挑戰)：ACDC是一個公共心臟MRI數據集，包括100例樣本。一系列短軸切片，覆蓋心臟從左心室基部到心尖，厚度為5～8 mm。短軸平面內空間分辨率從0.83到1.75 mm2/pixel，對應的標簽包括左心室(LV)、右心室(RV)和心肌(MYO)。數據集分為70個訓練樣本(1930個軸向切片)、10個驗證樣本和20個測試樣本。

2.2 實驗設置

2.2.1 環境

SwinEA網絡是基于Python 3.6和Pytorch 1.7.0實現的。對于所有的訓練案例，都使用數據增強，如翻轉和旋轉，以增加數據的多樣性。輸入的圖像大小設置為224×224，patch大小設置為4。本文在內存為32 GB的Nvidia V100 GPU、Linux操作系統上訓練本文的模型。

在訓練期間，批量大小為24，使用的是SGD優化器，動量為0.9，權值衰減為1e-4，對本文的反向傳播模型進行優化。實驗設置300個epoch。

2.2.2 評估指標

本文使用兩種類型的指標來評估所有模型。

評價預測與分割標簽重疊程度的Dice相似性系數(DSC)

(10)

其中，P為預測的分割圖，G為標簽。

HD(Hausdorff distance)，測量兩個分割貼圖之間最大的對稱距離

(11)

式中：d(·) 為歐幾里得距離，sup和inf分別為上極值和下極值。P為預測的分割圖，G為標簽。本文采用95%HD，以消除一個非常小的子集的異常值的影響。

2.3 實驗結果

2.3.1 在Synapse多器官CT數據集上的驗證

在Synapse多器官CT數據集上，本文提出的SwinEA與之前較新的方法的比較見表1，其中2-3列分別為平均DSC和HD評價指標結果，后8列為8個器官的具體DSC指標結果。實驗結果表明，本文的基于Swin Transformer融合多尺度上下文金字塔的方法分割精度達到了79.65%(DSC↑)和20.97%(HD↓)。與最近的TransUnet、Swin-Unet、MTUnet[25]方法相比，本文的算法在DSC評價指標和HD評價指標上有著提升，這表明本文的方法可以實現更好的分割邊緣預測。

表1 不同算法在Synapse多器官分割數據集的評價指標對比結果

不同方法對Synapse多器官CT數據集的分割結果如圖2所示。從圖中白色框標出的地方可以看出，本文方法在一些分割邊緣上比其它一些先進方法均有著分割邊緣清晰，區域明顯的效果。因此，在本文工作中，驗證了基于Swin Transformer的上下文金字塔融合多尺度特征結合的醫學圖像分割網絡通過邊緣融合模塊捕捉了豐富的多尺度上下文信息，計算相鄰像素間的局部相關性的卷積，在提取邊緣和角落等局部特征方面表現得很好，也獲得更好的分割結果。

圖2 不同算法在Synapse數據集上的分割對比結果

表1對不同算法的實驗結果進行了定量評估。U-Net為原始深度卷積網絡的方法；Transunet為生成模型的編碼結構使用Transformer，解碼器中使用卷積結構；SwinUnet為編解碼器結構均使用純的Transformer方法；SwinEA為本文提出的方法。從表中可以看出，通過使用上下文金字塔結構與Swin Transformer方法共同結合，DSC提升到79.65%，HD提升到20.97%。數據方面也驗證，在Swin Transformer網絡中引入上下文金字塔機制，可以有效且精確地針對邊緣完成腹部多器官分割的任務。

為了驗證該模型的性能，與分割標簽通過三維重建進行了比較。將預處理的切片結果還原為nii格式的三維數據集，比較三維下的分割結果。

圖3為模型在Synapse數據集不同層下的分割結果。第一列為Synapse數據集切片后的原始圖像，第二列為標簽結果，第三列為本文提出的網絡模型的分割結果，第四列為本文進行三維重建的分割結果。從圖3中可以看出，與第二列專家分割的結果進行對比，第三列本文的實驗分割結果在細節上分割精確，基本上與標簽分割結果相似。并且將切片數據進行三維重建之后，邊緣圓滑。

圖3 三維重建的Synapse分割結果

2.3.2 在ACDC數據集上的驗證

為了驗證模型的泛化性，將SwinEA在ACDC 數據集來進行醫學圖像分割。表2中總結了不同算法下的分割實驗數據結果，其中第2列為平均DSC值，后3列為右心室、心肌、左心室的具體DSC值。使用MRI模式的圖像數據作為輸入，SwinEA仍然能夠達到優秀的性能，平均準確率達到91.34%。這表明我們的方法具有良好的泛化能力和穩健性。如圖4中，第一行為ACDC數據集切片后的原始圖像，第二行為分割標簽，第三行為本文提出的網絡模型的分割結果，第四行為本文三維重建的分割結果。從圖4中可以看出，通過與第二行標簽結果的結果進行對比，第三行本文的實驗分割結果在細節上分割精確，基本上與標簽分割結果相似，并且將切片數據進行三維重建之后，可以看到分割邊緣圓滑，分割區域清晰。

表2 不同算法在ACDC心臟分割數據集的評價指標結果

圖4 不同算法在ACDC 數據集的分割結果

2.3.3 消融實驗

為了探究不同因素對模型性能的影響，本文在Synapse數據集上進行了消融實驗研究。

(1)邊緣感知模塊的位置的影響

本文的EA模塊被添加到1/4、1/8和1/16分辨率刻度的位置在圖中分別代表邊緣感知模塊1、邊緣感知模塊2、邊緣感知模塊3，探討不同位置對分割效果的影響。從圖5中可以看出模型的分割效果，邊緣感知模塊1、邊緣感知模塊2代表的黑色白色部分大部分表現比邊緣感知模塊3的灰色更準確。因此，為了使模型更加魯棒，本文將不使用邊緣感知模塊3的情況。

圖5 邊緣感知模塊的位置的影響

在實驗過程中發現此模塊在不同位置上的表現不同，位于淺層時效果更好，可以更好提取細節空間信息。在添加塊數上，塊數更多會隨之帶來提取信息過于冗余，表現一般，平衡位置與添加塊數的平衡，在實驗中驗證添加兩塊EA模塊，且位于淺層位置可以達到一個更優的結果。

(2)邊緣感知模塊的數量的影響

本文的邊緣感知模塊數量的討論結合位置限制，所以最多可以有3塊，最少可以為0塊。探討數量對分割性能的影響。從表3中可以數量看出模型的分割性能，3個邊緣感知模塊造成冗余，效果一般。兩塊效果普遍比一塊效果更優。因此，為了使模型更加健壯，本文將采用兩塊邊緣感知模塊。

表3 邊緣感知模塊的數量的影響

本文通過消融實驗，得出關于邊緣感知模塊模塊數量影響的結論，添加3塊邊緣感知模塊模塊時造成提取信息過于冗余，造成效果不佳，而針對于分割邊緣低維度信息特征更加有用，所以兩塊時，添加處于低維度特征的邊緣感知模塊1、邊緣感知模塊2模塊時，效果更好，將高分辨率的低水平特征和低分辨率的高水平特征融合起來，幫助劃分出詳細清晰的分割邊界。

(3)損失函數的影響

本文采用的損失函數包括常用在分割中的Dice損失和交叉熵損失，兩項損失進行消融實驗后，實驗結果見表4，結果驗證兩項損失聯合訓練的方法可以得到更好的分割效果。針對分割邊緣的清晰與精度，這兩種損失提供了損失的信息，更好解決了固定尺度卷積操作在目標分割中的局限性。

表4 損失函數的影響

3 結束語

本文提出了一種針對分割邊緣的利用Swin Transformer融合邊緣感知的醫學圖像分割網絡(SwinEA)，設計了一個基于上下文金字塔的邊緣感知(Edge-Aware)模塊，用于融合全局多尺度的上下文信息，主要針對邊緣和角落等局部特征。并且側重于解決U形網絡中全局多尺度上下文信息捕獲和整合的弱點。新型邊緣感知模塊以上下文金字塔為基礎的邊界感知模塊插入U形框架，開發和融合豐富的全局多尺度上下文信息。將高分辨率的低水平特征和低水平的特征融合起來。高分辨率的低層次特征和低分辨率的高層次特征來幫助劃定詳細的分割邊緣。

本文對不同類型的醫學圖像分割任務進行了綜合實驗，以驗證所提出的SwinEA的有效性和泛化性，包括腹部多器官分割任務和ACDC數據集。本文方法可以通過進一步的處理達到更好的性能，未來可以擴展應用到3D圖像以及其它醫學圖像分割任務，這也是我們近期的工作。