基于RT-Unet的醫(yī)學(xué)圖像分割算法

趙佳美，吳迪康，王志芳

(黑龍江大學(xué) 電子工程學(xué)院，黑龍江 哈爾濱 150080)

0 引言

隨著計算機技術(shù)和人工智能的快速發(fā)展，深度學(xué)習(xí)在自然語言處理、機器視覺任務(wù)和多模態(tài)分析等各個方面都取得了突破性的進展[1-3]。深度學(xué)習(xí)能夠自動且準確地提取特征并具有強大的建模能力，相較于傳統(tǒng)算法在醫(yī)學(xué)圖像分割中得到了更廣泛的應(yīng)用。但由于不同患者間器官的結(jié)構(gòu)、紋理和大小等方面存在差異，醫(yī)學(xué)圖像分割精度高度依賴于成像設(shè)備與醫(yī)生的專業(yè)知識，快速且準確地完成分割任務(wù)仍然具有挑戰(zhàn)性[4]。

目前，U-Net[5]是基于深度學(xué)習(xí)的醫(yī)學(xué)圖像分割算法的主流方法，由編碼器與解碼器構(gòu)成的U型對稱結(jié)構(gòu)，通過跳過連接使用簡單的拼接操作完成多尺度特征的融合，達到減少空間信息丟失的目的。這種U型結(jié)構(gòu)使得U-Net在各種醫(yī)學(xué)影像處理方面都取得了巨大的成功，許多算法都是在此基礎(chǔ)上進行的改進，諸如U-Net++[6]，3D U-Net[7]和Res-UNet[8]等。但是由于卷積運算固有的局限性，圖像上的對角信息很難直接進行卷積運算，無法有效地提取全局信息，使U-Net在遠程關(guān)系建模方面存在一定限制。為了解決這個問題，此前方法大多是引入圖像金字塔[9]、深度卷積層[10]或注意力機制[11]等，但會造成算法計算量的直線上升。

近兩年，許多研究者將自然語言處理領(lǐng)域的Transformer引入到機器視覺領(lǐng)域[12]。2020年Dosovitskiy等[13]首次提出基于Transformer的ViT算法完成圖像識別任務(wù)，取得了與卷積神經(jīng)網(wǎng)絡(luò)(Convolutional Neural Network，CNN)相當?shù)男阅埽菂s需要較大的數(shù)據(jù)集來完成模型的訓(xùn)練。隨著研究的深入，CNN與Transformer二者結(jié)合在圖像處理領(lǐng)域上取得了較大的突破。2021年Chen等[14]開創(chuàng)性地提出了TransUNet，利用CNN提取細節(jié)信息，將其送入到Transformer中提取全局信息，通過跳過連接融合不同分辨率的多尺度特征。該算法在腹部多器官分割、心臟分割數(shù)據(jù)集中取得了優(yōu)異的成績，充分證明了CNN與Transformer組合的有效性，但卻忽略了跨尺度特征之間的依賴性。2021年Cao等[15]提出了基于Transformer的Swin-Unet算法應(yīng)用于醫(yī)學(xué)圖像分割任務(wù)，極大地減少了算法的計算量并取得了較高的分割準確率。但Swin-Unet在訓(xùn)練過程中只能學(xué)習(xí)單尺度的上下文特征，缺乏處理局部信息與多尺度信息的內(nèi)在歸納偏置[16]。針對上述問題，本文在Swin-Unet的基礎(chǔ)上提出了RT-Unet算法，可以充分地學(xué)習(xí)多尺度特征信息，提高算法的訓(xùn)練精度與分割準確率。

1 研究方法

1.1 整體結(jié)構(gòu)

RT-Unet是一種由編碼器、解碼器和跳過連接組成的U型對稱拓撲結(jié)構(gòu)，在殘差連接的過程中加入局部歸納偏置模塊構(gòu)建RESwin Transformer模塊，協(xié)助完成特征信息的學(xué)習(xí)，通過疊加聚合不同感受野的卷積層來提取各個層間或?qū)觾?nèi)的多尺度特征。然后，在位置嵌入與編碼模塊使用4層連續(xù)卷積層，并選擇GELU作為激活函數(shù)。RT-Unet不但可以有效地學(xué)習(xí)局部與遠程信息，而且可以獲得更豐富的多尺度特征，增強網(wǎng)絡(luò)的泛化能力，RT-Unet的總體結(jié)構(gòu)如圖1所示。

圖1 RT-Unet的結(jié)構(gòu)Fig.1 Structure of RT-Unet

首先，編碼部分將醫(yī)學(xué)影像數(shù)據(jù)輸入圖像分塊處理層(Patch Partition)中，將其切分成大小相等的非重疊圖像塊，并對圖像塊進行位置嵌入與編碼(Linear Embedding)，完成轉(zhuǎn)換特征維度等操作；然后，將轉(zhuǎn)換后的圖像塊輸入RESwin Transformer模塊進行特征學(xué)習(xí)，通過圖像塊合并層(Patch Merging)實現(xiàn)下采樣功能。與之對稱的解碼部分，采用圖像塊擴展層(Patch Expanding)完成上采樣操作，逐步恢復(fù)圖像的大小。同時，為了減少空間信息的丟失，編-解碼部分利用跳過連接完成多尺度上下文的特征融合；最后，將得到的特征信息應(yīng)用線性投影層完成最終的分割預(yù)測。

1.2 RESwin Transformer 模塊

圖2 RESwin Transformer模塊的結(jié)構(gòu)Fig.2 Structure of RESwin Transformer block

RESwin Transformer模塊加入LIB模塊后，一方面可以在尺度不變的情況下，通過疊加聚合具有不同感受野的卷積層獲取豐富的多尺度上下文特征。另一方面，RESwin Transformer模塊自身的注意力機制在遠程關(guān)系建模上具有一定優(yōu)勢，又通過模塊中的卷積計算相鄰像素的相關(guān)性獲得邊緣與角落的局部細節(jié)信息。因此，RESwin Transformer模塊可以更有效地建模局部和遠程依賴關(guān)系，較好地提升網(wǎng)絡(luò)的魯棒性與分割準確率。RESwin Transformer模塊的計算過程總結(jié)如下：

(1)

(2)

(3)

(4)

LIB模塊嵌入到RESwin Transformer模塊的殘差連接過程中,具體結(jié)構(gòu)如圖3所示。首先，Seq2Img層完成簡單的重塑作用，將輸入的一維序列轉(zhuǎn)換為二維特征圖像塊。然后，依次將圖像塊輸入感受野為1×1,3×3,1×1組成的3個卷積層(Conv)中提取局部多尺度特征。隨后，每層再進行批量歸一化(BN)和ReLU激活函數(shù)處理。最后，由Img2Seq層將具有多尺度特征的二維圖像塊轉(zhuǎn)換回一維序列并入到MLP中，協(xié)助RESwin Transformer模塊學(xué)習(xí)局部特征與遠程依賴關(guān)系。

圖3 LIB模塊的結(jié)構(gòu)Fig.3 Structure of the LIB block

1.3 位置嵌入與編碼模塊

位置嵌入與編碼模塊的主要作用是將圖片轉(zhuǎn)換為RESwin Transformer模塊可以處理的一維序列，并對序列進行相對位置編碼。與自然語言處理領(lǐng)域中的詞匯不同，圖像與圖像之間是沒有位置順序的。因此，RT-Unet首先將一張醫(yī)學(xué)影像劃分為若干個不重疊的圖像塊，每個圖像塊充當一個詞匯。然后，通過位置嵌入與編碼模塊將得到的圖像塊展平，再利用線性映射將二維圖像轉(zhuǎn)換為一維向量并對其逐一進行編碼，使圖像塊可以嵌入位置信息。

位置嵌入與編碼模塊由4個連續(xù)的卷積層完成像素級空間信息的編碼，每個卷積層后跟隨一個GELU激活函數(shù)和歸一化層，結(jié)構(gòu)如圖4所示。

圖4 位置嵌入與編碼模塊的結(jié)構(gòu)Fig.4 Structure of position embedding and encoding block

激活函數(shù)對網(wǎng)絡(luò)上一層輸出的神經(jīng)元進行處理，將結(jié)果傳遞到下一神經(jīng)元中，給神經(jīng)元引入非線性因素，使網(wǎng)絡(luò)可以應(yīng)用到眾多的非線性模型中。激活函數(shù)ReLU輸入值為正時，網(wǎng)絡(luò)可以正常收斂；但為負時，ReLU的輸出值將始終為0，導(dǎo)致網(wǎng)絡(luò)參數(shù)的更新出現(xiàn)停滯，進而出現(xiàn)梯度消失問題。而GELU激活函數(shù)在ReLU的基礎(chǔ)上加入統(tǒng)計特性完成對輸入神經(jīng)元的概率描述。它是一種符合預(yù)期的隨機正則變換方式，能夠有效地避免梯度消失問題，增加算法的非線性。GELU激活函數(shù)的實驗效果要優(yōu)于ReLU，尤其是在Transformer模型中表現(xiàn)最好[17]。本文采用GELU激活函數(shù)，計算公式如下：

GELU(x)=xP(X≤x)=xφ(x)，

(5)

式中，φ(x)為正態(tài)分布的概率函數(shù)，當為標準正態(tài)分布時：

(6)

2 實驗結(jié)果分析

2.1 評價指標與Synapse數(shù)據(jù)集

本文采用2個評價指標：平均Dice系數(shù)(DSC)和Hausdorff距離(HD)[18]。二者都是描述2個樣本集合間的相似性度量，但DSC更加側(cè)重于內(nèi)部填充部分的分割準確性，而HD則對分割的邊界更為敏感。它們的定義如下：

(7)

HD(A,B)=max(h(A,B),h(B,A))，

(8)

(9)

(10)

式中，A代表醫(yī)學(xué)影像的標簽圖；B代表算法分割的預(yù)測圖。式(10)為雙向HD，式(11)和式(12)分別為從A到B和從B到A的單向HD。

本文在Synapse腹部多器官分割數(shù)據(jù)集上驗證RT-Unet的分割效果。Synapse數(shù)據(jù)集共包含30張腹部掃描CT，每張CT里含有8個器官：主動脈、膽囊、脾臟、左腎、右腎、肝臟、胰腺和胃。為了避免由于數(shù)據(jù)集較小而導(dǎo)致過擬合現(xiàn)象的發(fā)生，將3D醫(yī)學(xué)影像轉(zhuǎn)換成大小為224 pixel×224 pixel的2D切片，共獲得3 779張2D切片，其中2 212張用于算法的訓(xùn)練，剩余的完成算法的測試。

2.2 實驗結(jié)果分析

RT-Unet在Python3.6，Pytorch1.5.0和NVIDIA 2080Ti GPU的基礎(chǔ)上完成。損失函數(shù)為交叉熵損失函數(shù)，優(yōu)化器選擇隨機梯度下降(SGD)。在相同的實驗條件下，將RT-Unet在Synapse數(shù)據(jù)集上分割8個腹部器官得到的DSC，HD參數(shù)與經(jīng)典的分割網(wǎng)絡(luò)ViT，V-Net，U-Net，Swin-Unet，TU-Net進行對比分析，實驗結(jié)果如表1所示。

表1 不同方法在Synapse數(shù)據(jù)集上的分割結(jié)果Tab.1 Segmentation results of different algorithms on Synapse dataset

由表1可以看出，與基礎(chǔ)的ViT，V-Net相比，RT-Unet的DSC與HD參數(shù)都得到了較大的提升，證明CNN與Transformer組合的有效性。其次，與經(jīng)典的醫(yī)學(xué)圖像分割網(wǎng)絡(luò)U-Net相比，RT-Unet 的DSC，HD參數(shù)分別提升了2.23%與16.27 mm，說明RT-Unet對于較小器官與分割邊界更為敏感。由于實驗設(shè)備與參數(shù)設(shè)定等諸多因素，Swin-Unet與TU-Net的DSC，HD參數(shù)分別為：77.28%，77.95%和26.93，23.43 mm。RT-Unet較二者分別提升了1.8%，1.13%和3.5，1.27 mm，取得了DSC為79.08%，HD為23.43 mm較為準確的分割結(jié)果，表明RT-Unet同時注重了局部細節(jié)信息與遠程依賴關(guān)系，又促進了多尺度上下文信息的提取，有效地提升了算法的分割精度。

本文主要對Swin-Unet，TU-Net，RT-Unet算法的分割結(jié)果進行可視化對比分析，結(jié)果如圖5所示。由圖5的第1行可以看出，RT-Unet分割的邊界相較于其他算法更為清晰光滑；從第2行可以看出，Swin-Unet和TU-Net對于肝臟都存在漏分割的現(xiàn)象，并且TU-Net也并未分割出胰腺，相比之下RT-Unet更加注重細節(jié)信息的提取，對器官的分割也更為精準；從第3行可以看出Swin-Unet和TU-Net對于胃部都出現(xiàn)了過分割的現(xiàn)象，RT-Unet綜合豐富的多尺度上下文特征，此類現(xiàn)象出現(xiàn)的概率較小。

圖5 不同算法的分割結(jié)果Fig.5 Segmentation results of different algorithms

2.3 消融實驗

通過分別移除算法中的局部歸納偏置模塊或替換為原有的位置嵌入與編碼模塊來驗證它們的有效性，并在Synapse數(shù)據(jù)集上完成了消融實驗，實驗結(jié)果如表2所示。從表2可以看出，當算法中僅有位置嵌入與編碼模塊作用時，取得了DSC為77.70%的效果，在不影響整個算法復(fù)雜度的情況下，較小地提升了分割的準確性。當算法中僅有局部歸納偏置模塊作用時，取得了DSC為78.44%，提升了近1.2%，有效地促進了不同分辨率信息的融合。由此可證，二者對于算法性能的提升是不可缺少的，刪除其中任何一個都可能導(dǎo)致性能的損失。

表2 基于位置嵌入與編碼模塊和LIB模塊的消融實驗Tab.2 Ablation studies on linear embedding block and LIB block

3 結(jié)束語

本文提出了一種基于RESwin Transformer的U型對稱結(jié)構(gòu)算法——RT-Unet。該算法在RESwin Transformer模塊中引入了內(nèi)在的LIB，由具有不同感受野的卷積層獲取豐富的多尺度信息。卷積與Transformer二者組合，既可以關(guān)注局部細節(jié)信息，又可以有效地建模遠程依賴關(guān)系。相較于其他算法，RT-Unet更加注重器官邊界與小器官分割的準確性，較少會出現(xiàn)過分割與欠分割現(xiàn)象。而且，本文在位置嵌入與編碼模塊中，選擇GELU激活函數(shù)增強算法的泛化能力，極大地避免了梯度消失問題的出現(xiàn)。在Synapse腹部多器官分割實驗中，RT-Unet取得了DSC為79.08%，HD為23.43 mm的分割結(jié)果，證明了其有效性。