改進U-Net的多級邊緣增強醫學圖像分割網絡

胡帥,李華玲,郝德琛

(中北大學軟件學院,山西 太原 030051)

0 引言

隨著現代醫學的不斷發展,各種醫療影像設備廣泛應用于臨床診療,計算機斷層掃描(CT)、X射線、超聲、正電子發射斷層掃描(PET)和核磁共振成像(MRI)等均為醫師診斷治療發揮著重要的作用。醫學圖像分割的主要目的是從各類醫學圖像中自動分割出病灶區域,輔助醫師更快更準確地作出診斷,極大地提高了醫師的診斷效率。因此,醫學圖像分割模型的分割效果顯得格外重要。與此同時,醫學圖像相對于自然圖像背景環境更復雜,具體表現在結構邊緣模糊、形狀多變、紋理結構細節豐富,使得醫學圖像分割任務更具挑戰性。隨著深度學習技術在醫學圖像分割中的廣泛應用,醫學圖像分割精度有了顯著提升。卷積神經網絡(CNN)的出現,尤其是全卷積神經網絡的出現,使深度學習模型具備了高速度、高精度和全自動的特點。RONNEBERGER等[1]提出用于醫學圖像分割的卷積網絡(U-Net)。鑒于U-Net的編解碼器結構在醫學圖像分割中的優異效果,多數性能優異的醫學圖像分割模型都是基于U-Net的主要結構進行改進,例如:用于醫學圖像分割的嵌套U-Net(UNet++)[2],注意力U-Net(Attention U-Net)[3],用于醫學圖像分割的全尺寸連接UNet(UNet 3+)[4],基于多層感知器的快速醫學圖像分割網絡(UNeXt)[5]等模型。

目前,醫學圖像分割模型仍存在以下問題:1)普通卷積運算的局部性導致多數模型未能有效地提取圖像的全局信息;2)對圖像邊緣信息的忽略導致病灶分割邊緣區域模糊。為解決以上問題,本文在保留U-Net的高效編解碼器結構的同時,在U形結構跳躍連接部分引入多級邊緣增強(MEE)模塊和細節特征關聯(DFA)模塊,提出改進U-Net的多級邊緣增強醫學圖像分割網絡(MDU-Net)模型,改善了醫學圖像分割時邊緣模糊的問題,同時提高了分割精度,在LiST2017肝臟腫瘤分割挑戰數據集[6]、胸部X射線掩膜和標簽數據集[7-8]2個數據集上進行實驗驗證。

1 相關工作

目前,主流的醫學影像分割模型主要分為CNN和Transformer 2種結構。CNN自提出以來已成為圖像處理的主流結構。2015年,SHELHAMER等[9]提出全卷積網絡(FCN),開啟了其在語義分割領域的應用,FCN的基本思想是使用卷積層替換全連接層,使用轉置卷積進行上采樣,分割精度優于傳統的分割方法。在FCN的基礎上,RONNEBERGER等[1]通過引入跳躍連接提出一種編解碼器結構的U-Net,編解碼器之間通過跳躍連接將相同層級的特征信息進行拼接,U-Net在各種類型的醫學圖像分割任務中均表現出了十分優異的分割性能,可以有效地融合低層和高層圖像特征,緩解語義間隙問題,在醫學圖像上顯示出很好的檢測能力。由于U-Net在醫學圖像分割領域的突出貢獻,隨即出現了多種改進U-Net的醫學圖像分割網絡。OKTAY等[3]提出Attention U-Net,能夠關注不同形狀和大小的目標結構,同時能夠抑制圖像中的不相關區域并突出有用的顯著特征。ALOM等[10]提出一種用于醫學圖像分割的遞歸殘差U-Net(R2U-Net),使用遞歸卷積代替傳統卷積,以保證更好的特征表示。林志潔等[11]使用內卷操作替換了原始的卷積操作提升了U-Net對局部特征的學習能力。劉文等[12]對U-Net架構中的卷積模塊進行改進,同時引入了批量歸一化層提出Concat-UNet,增強了網絡的特征提取能力。LIN等[13]提出一種改進FCN的網絡架構(RefineU-Net),通過結合淺層特征中的全局上下文信息和深層特征中的語義信息,將其進行全局細化,使用生成的判別性注意力特征與解碼器特征進行連接,便于局部細化操作。HUANG等[14]提出一種密集卷積網絡(DenseNet),緩解了梯度消失問題,加強了特征傳播,提高了特征的重復利用。受此思想的啟發,AZAD等[15]使用雙向卷積長短期記憶(BConvLSTM)網絡并以非線性的方式將從相應層次的編碼器與解碼器中提取的特征圖結合起來,采用密集連接的卷積操作實現特征重用,提出一種用于醫學圖像分割的雙向網絡(ConvLSTM U-Net),優化跳躍連接,同時結合BConvLSTM塊,以捕獲更有鑒別性的信息。ZHOU等[2]提出UNet++,以縮小編解碼器之間的特征圖語義差距。HUANG等[4]提出UNet 3+,將全尺寸跳躍連接與不同尺度的特征信息圖的高級語義相結合,并通過深度監督從全尺寸聚合特征圖中學習分層表示。HE等[16]提出殘差結構,增加了從輸入到輸出的路徑,有效緩解了梯度消失等問題,提高了網絡的收斂速度和運行效率。為了解決卷積塊堆疊造成的有效感受野減小的問題,XU等[17]通過在普通卷積塊之前嵌入由深度可分離卷積組成并具有殘差結構的卷積模塊,提出一種用于醫學圖像分割的更深入緊湊的分散注意力U-Net(DCSAU-Net)。

在醫學圖像分割領域,自動化的分割模型不僅對分割的準確率有較高的要求,而且對于邊緣區域的分割也要求嚴格。對于醫學圖像分割邊緣區域模糊的問題,學者們進行了大量研究。LIU等[18]結合邊緣特征提取模塊提出一種用于醫學圖像分割的多層邊緣注意力網絡(MEA-Net),通過從不同的編碼器階段采用卷積操作獲取的邊緣特征序列融合到最后的分割掩碼中,提高邊緣區域的分割效果。ZENG等[19]提出一種用于醫學圖像分割的密集邊緣注意力UNet(DEA-UNet),該模型從U形網絡最底部上采樣層開始使用反向注意力模塊逐級提取特征,并結合邊緣引導模塊更大限度地獲取邊緣特征。ZHANG等[20]在編碼器部分加入了邊緣注意力引導模塊來獲取邊緣信息表示,最后將邊緣特征與多尺度編碼器的輸出特征信息相結合生成最后的分割掩碼。HAO等[21]提出一種基于圖的邊緣注意力門醫學圖像分割模型(EAGC-UNet++),將圖論的思想引入醫學圖像分割任務中,使用殘差圖卷積塊以及傳統卷積塊共同作為UNet++網絡的編碼器部分,在模型中加入邊緣注意力門來獲取邊緣信息。

在上文提到的各種改進U-Net模型中,對于醫學圖像的邊緣信息處理還有待改進,醫學圖像的細節特征信息利用還不夠完善。為了解決上述問題,本文提出一種MDU-Net,由擴張門控注意力(DGA)、多級邊緣增強、細節特征關聯3個新設計的模塊組成。

2 MDU-Net模型

2.1 整體架構

本文提出的MDU-Net模型的整體架構如圖1所示。以編解碼器結構的U-Net模型作為主體結構,并嵌入MEE模塊和DFA模塊。多級邊緣增強模塊的設計是為了充分捕獲低級特征圖中的邊緣特征信息,融合不同層次的圖像邊界特征,為最終的分割效果提供邊界補充信息。細節特征關聯模塊接收雙層特征圖信息,逐步引導多尺度上下文特征信息的融合,高效利用圖像的深層次細節特征。在主體結構中加入各個模塊以學習圖像的邊界以及細節特征,能獲取更準確的分割結果。

圖1 MDU-Net模型整體架構Fig.1 Overall architecture of MDU-Net model

2.2 多級邊緣增強模塊

在編碼器結構中,低層特征具有豐富的空間細節特征,高層特征具有豐富的語義信息。孫軍梅等[22]僅在U-Net基準網絡的解碼器的最后3層進行了邊緣特征的提取,忽略了編碼器低級特征層中豐富的邊緣特征信息,會導致大量邊緣信息缺失。不同層次特征信息的融合不僅實現了多層次特征間的優勢互補,而且有效地提高了網絡的遠程上下文學習。在U-Net網絡中,在編碼器結構中的下采樣操作會造成邊緣特征信息的嚴重缺失。為了補充缺失的邊緣信息和高效利用低級特征中豐富的邊緣信息,提出MEE模塊,該模塊對低級特征的邊緣信息進行深入提取并補充至解碼器特征信息中。如圖2所示,該模塊將編碼器結構中低級特征塊E1輸入擴張門控注意力模塊獲得高級邊緣信息塊E′1,將特征塊E2先輸入改進的注意力融合(AF)模塊在獲得特征信息后進行一次上采樣操作得到E′2,再與高級邊緣信息塊E′1進行元素相乘得到邊緣增強信息Ed。

圖2 多級邊緣增強模塊結構Fig.2 Structure of multistage edge-enhanced module

上述過程如式(1)～式(3)所示:

E′1=DGA(E1)

(1)

E′2=UpSampling(AF(E2))

(2)

Ed=E′1?E′2

(3)

其中:DGA為輸入擴張門控注意力模塊;UpSampling表示上采樣操作;?表示基于元素的乘法。

在改進的并行注意力模塊中,將通道注意力塊和空間注意力塊進行并聯操作,通過參數α和β控制注意力塊的加和比例,參數α、β的總和為1。經過多次實驗驗證,設置α、β同時為0.5可獲得最佳分割效果。

2.3 擴張門控注意力模塊

醫學圖像分割屬于密集預測任務,同時獲取全局和局部信息是提高分割性能的關鍵。全局信息有助于模型了解目標區域的整體結構及其與背景的關系,從而更準確地定位目標區域。李翠云等[23]將邊緣增強的注意力模塊嵌入解碼器結構,未注意到低級特征中包含的邊緣信息。本文注意到低級特征中的豐富邊緣信息,同時采用不同擴張率的深度可分離卷積獲取不同的特征信息。首先,通過擴張率為1和3的卷積獲得局部特征信息,通過擴張率為5和7的卷積獲得全局特征信息。然后,將取得的局部特征信息和全局特征信息進行逐像素相加,并加入殘差信息。這樣能有效建立圖像的上下文信息,充分利用不同維度的特征信息對全局信息進行補充。

上述過程可由式(4)～式(7)表示:

Edi=DSCd=i(E)

(4)

Eattg=Sigmoid(Ed7)⊙Ed5

(5)

Eattl=Sigmoid(Ed3)⊙Ed1

(6)

Edga=E+Eattl+Eattg

(7)

其中:E表示輸入特征;DSCd=i(i=1,3,5,7)表示擴張率為i的深度可分離卷積操作;⊙表示基于元素的乘法;Eattg和Eattl為計算出的全局特征注意力圖和局部特征注意力圖。

擴張門控注意力模塊結構如圖3所示。

圖3 擴張門控注意力模塊結構Fig.3 Structure of dilated gated attention module

2.4 細節特征關聯模塊

醫學圖像分割目標的尺度有很大變化,同時一些分割目標結構邊界模糊,形狀復雜,紋理異質,非常容易產生錯誤的分割結果。多尺度上下文信息可以提高語義分割任務的性能,由于目標的尺度是多樣的,因此有效利用多尺度的上下文信息是實現目標精確定位的關鍵。受此啟發,在細節特征關聯模塊中,嵌入多尺度特征提取(MFE)模塊來獲取特征圖的多尺度特征信息,并聯平均池化和最大池化操作,將池化結果相加后進行Sigmoid操作生成池化特征注意力圖Epoolatt,如圖4所示。

圖4 細節特征關聯模塊Fig.4 Detail feature association module

WOO等[24]認為:通道注意力和空間注意力在抑制不相關信息方面發揮著重要作用。王曉援等[25]在高級特征信息塊中提取了多尺度特征信息,但未考慮模型對高級特征信息的注意程度,會提取過多冗余的信息。本文在特征注意力融合模塊中并聯通道注意力模塊以及空間注意力模塊最終獲取注意力融合特征圖,以此提高該模塊對相關信息的注意程度,最后將池化特征注意力圖與MFE模塊得到的多尺度特征信息進行逐像素相乘,之后與注意力融合特征圖進行逐像素相加獲得該模塊的結果。

上述過程可由式(8)～式(14)表示:

Epoolatt=

Sigmoid(AvgPool(Ei)+MaxPool(Ei))

(8)

M1=Conv1(Ei+1)

(9)

M3=Conv1(Conv3(Ei+1))

(10)

M5=Conv1(Conv5(Ei+1))

(11)

M7=Conv1(Conv7(Ei+1))

(12)

Esc=Concat(M1,M3,M5,M7)

(13)

Eout=Epoolatt⊙Esc+AF(Ei+1)

(14)

其中:Ei表示輸入特征;AvgPool和MaxPool分別表示平均池化操作和最大池化操作;Concat表示特征圖通道連接操作;⊙表示基于元素的乘法;AF表示注意力融合操作;Eout表示該模塊最后的輸出結果。

3 實驗與結果分析

3.1 數據集與數據預處理

使用2個數據集來驗證MDU-Net模型的通用性和有效性。由于獲取圖像的方式大相徑庭,因此在使用同一模型進行訓練時需要對原有的圖像進行預處理,將圖像的長和寬處理為相同尺寸。在使用數據時,將圖像統一設置為256×256×3,即圖像長度和寬度分別為256像素,通道數為3。

3.1.1 LiTS2017肝臟腫瘤分割挑戰數據集

LiTS2017肝臟腫瘤分割挑戰數據集包含131組由世界各地醫院提供的CT圖像數據,選擇其中前30組CT圖像數據進行實驗。在實驗中,將數據集按80∶20劃分訓練集和測試集,其中,訓練集包含1 540張圖像,測試集包含385張圖像。

3.1.2 胸部X射線掩膜和標簽數據集

胸部X射線掩膜和標簽數據集包含138張X射線圖片,其中,80張X射線圖片正常,58張X射線圖片存在結核病異常。這個數據集有丟失標簽的問題,經過分析和處理后,將數據集按80∶20劃分訓練集和測試集,其中,訓練集包含396張圖像,測試集包含99張圖像。

3.2 實驗處理

3.2.1 實驗環境

模型使用Python 3.8和PyTorch實現。模型訓練和測試是在1臺服務器上進行的,使用NVIDIA RTX 3060Ti顯卡和16 GB內存。將數據集中的所有圖像及其對應的標簽大小分別調整為256×256×3。使用的優化器是RMSProp,學習率為0.000 1。用于訓練的批量大小為2,輪次數設置為50。

3.2.2 評價指標

為了評估MDU-Net模型的有效性,采用Dice相似系數、靈敏度、準確率、特異性、查準率以及Jaccard相似系數作為評價指標,計算公式分別如式(15)～式(20)所示:

(15)

(16)

(17)

(18)

(19)

(20)

其中:NTP、NTN、NFP、NFN分別表示真陽性、真陰性、假陽性、假陰性樣本數量。

上述評價指標的數值越大,代表模型分割效果越好。如果這些指標的數值接近1,則意味著輸出掩碼與標簽真值掩碼重疊。

3.3 實驗結果

3.3.1 模型比較

1)LiTS2017肝臟腫瘤分割挑戰數據集。分別使用U-Net[1]、UNet++[2]、DCSAU-Net[17]、利用深度殘差U-Net的道路提取(ResUNet)模型[26]、用于二維醫學圖像分割的上下文編碼網絡(CE-Net)模型[27]、用于圖像分割的深度卷積編解碼器體系結構(SegNet)[28]、用于醫學圖像分割的Transformers強編碼器(TransUNet)模型[29]、用于醫學圖像分割的Transformer編碼器和CNN解碼器的整合(LeViT-UNet)模型[30]等在LiTS2017數據集上進行實驗,實驗結果的數據展示如表1所示,其中最優指標值用加粗字體標示,下同,對應的分割結果如圖5所示。由表1可以看出,MDU-Net實現了最好的效果,在準確率、靈敏度、Jaccard相似系數、Dice相似系數等指標上均名列首位,尤其在靈敏度、Dice相似系數指標上有明顯的提升,分別達到了96.24%、0.926 5,在準確率、Jaccard相似系數指標上均有提升。由圖5可以看出,MDU-Net對目標區域的分割效果在完整度以及邊緣分割細節上都明顯優于其他模型,再次證明了MDU-Net相較于其他模型更加高效。

表1 LiTS2017肝臟腫瘤分割挑戰數據集上的實驗數據Table 1 Experimental data on the LiTS2017 liver tumor segmentation challenge dataset

圖5 LiTS2017肝臟腫瘤分割挑戰數據集上的分割結果比較Fig.5 Comparison of segmentation results on the LiTS2017 liver tumor segmentation challenge dataset

2)胸部X射線掩膜和標簽數據集。使用不同模型在胸部X射線掩膜和標簽數據集上進行實驗,最終的實驗數據和實驗分割效果分別在表2和圖6中進行展示。由表2可以看到,MDU-Net在準確率、特異性、查準率、Jaccard相似系數、Dice相似系數評價指標上都有提升,分別達到了97.96%、98.53%、95.46%、0.920 4、0.957 8。由圖6可以看出,MDU-Net在分割目標邊緣區域具有優秀的分割效果,主要原因為MEE模塊提供了更多的邊緣信息,使得MDU-Net在邊緣處的效果顯著優于其他模型。

表2 胸部X射線掩膜和標簽數據集上的實驗數據Table 2 Experimental data on chest X-ray mask and label dataset

圖6 胸部X射線掩模和標簽數據集上的分割結果比較Fig.6 Comparison of segmentation results on the chest X-ray mask and label dataset

此外,還對MDU-Net的參數量、計算量指標與不同模型進行比較分析,結果如表3所示。MDU-Net為了提取醫學圖像中更為復雜豐富的特征信息,提出的MEE、DGA和DFA模塊結構相對復雜,導致模型的參數量和計算量相對較高,但這是一個值得權衡的問題。在醫學圖像分割領域,分割精度對于診斷和治療意義重大,因此提高分割精度是至關重要的目標。結合圖5和圖6的分割結果可以看出,增加的參數對于模型的分割效果起到了關鍵作用。

表3 模型參數量以及計算量分析比較Table 3 Analysis and comparison of the model parameter and calculation quantity

3.3.2 消融實驗

為了進一步驗證所提模塊的有效性,在實驗條件相同的情況下,以U-Net為基準網絡,逐步添加多級邊緣增強模塊、擴張門控注意力模塊和細節特征關聯模塊進行消融研究,并在肝臟腫瘤分割挑戰數據集和胸部X射線掩模和標簽數據集上分別進行實驗,實驗數據如表4和表5所示,其中“√”表示使用該模型或模塊。

表4 LiTS2017肝臟腫瘤分割挑戰數據集上的消融實驗分割結果Table 4 Segmentation results of ablation experiments on the LiTS2017 liver tumor segmentation challenge dataset

分析兩組消融數據可知,相比于基準網絡U-Net,在其添加DFA模塊后的準確率、靈敏度、Jaccard相似系數、Dice相似系數指標上均有所提升。DFA、DGA和MEE 3種模塊的組合在整體上的分割效果最好,基于3種模塊的網絡模型在胸部X射線掩模和標簽數據集的準確率、特異性、查準率、Jaccard相似系數、Dice相似系數上表現均為最佳。

4 結束語

本文注重高效保留醫學圖像的邊緣信息以及細節特征的提取與融合,提出改進U-Net的MEE醫學圖像分割網絡模型。在具有豐富邊緣特征信息的低級特征塊中,利用MEE模塊提取邊緣特征,嵌入DGA模塊用于抑制非重要特征的獲取。在此基礎上,在模型跳躍連接上嵌入DFA模塊用于補充圖像的全局和局部特征信息。實驗結果表明,該模型在邊緣區域分割以及整體分割效果上表現優異,優于目前先進的醫學圖像分割模型。下一步將繼續優化所提模型,降低模型參數量,使其可應用于實際醫學圖像分割任務。