基于改進U-net的醫學圖像分割模型

［摘 要］ U-net模型在醫學圖像分割中取得了很多的發展，但是其忽略了低層視覺特征和高層語義特征之間的差距，限制了網絡的特征提取能力。針對這個問題，首先設計了一個特征增強模塊，以更大的感受野來增強網絡的特征提取能力，并使用跳躍連接將網絡中低級特征與高級特征融合。同時引入三重注意力機制來融合空間維度和通道維度的特征，從而達到增強有用信息的目的。最后，設計了一種新的損失函數，引入二元交叉熵和Dice損失函數，來抑制數據集中類別不平衡的問題。該模型在BraTs數據集上進行了實驗，取得了0.8592的平均Dice。實驗結果表明，該模型性能優于其他常用的分割網絡。

［關鍵詞］ 醫學圖像分割； 特征增強； 注意力機制； 損失函數

［中圖分類號］ TP391 ［文獻標識碼］ A

近年來，隨著深度學習的快速發展，許多研究開始將深度學習應用到醫學圖像分割領域，并取得了巨大的成功[1]，展現了卷積神經網絡（CNN）在醫學圖像分割領域的優勢[2-3]。CNN將醫學圖像作為輸入，中間卷積層和池化層負責提取醫學圖像的特征，整個過程不需要人為干預，避免了傳統機器學習中繁雜的操作[4-5]。隨著研究的深入，有學者開發出了完全卷積神經網絡（FCN）[6]，它以驚人的分割精度迅速成為了最流行的分割方法之一。后來，Ronneberger[7]等提出了一種名為U-net的全卷積神經網絡用于醫學圖像分割并取得了巨大的成功。U-net由編碼器和解碼器構成，在網絡的內部還使用跳躍連接來融合不同層的特征，極大的提高了網絡的特征提取能力。然而，在U-net網絡及其各種改進的網絡中都是使用相同尺寸的卷積塊來提取特征，這在一定程度上限制了網絡的特征提取能力。為解決這個問題，Zhang等[8]在U-Net中引入InceptionResNet和密集連接，形成了一種新的語義分割網絡，在肺、血管和腦瘤醫學圖像數據集上取得了較好的分割性能。但該網絡同樣具有一定的局限性，它忽略了不同特征之間的語義差距，在編碼過程中容易丟失微小特征的信息。在此基礎上，Yu等[9]提出了一種深層聚合網絡（DLANet）來解決編碼過程中信息丟失的問題。該網絡以深層聚合的方式來提取特征。在他們的語義分割任務中，深層聚合的思想是將兩個模塊迭代使用，從而形成一個層次逐漸加深、分辨率逐漸提高的解碼器。U-net++[10]也使用了深度聚合思想，但是與DLANet不同的是，U-net++不僅采用迭代深度聚合策略，還增加了密集連接，從更深層來聚合語義特征，形成高度密集的特征聚合模型，達到減少信息丟失的目的。同時，Hu等[11]在神經網絡中引入注意力機制，通過注意力機制來放大有用信息抑制無用信息，從而提升網絡分割性能。

綜上所述，為了解決基于U-net的醫學圖像分割網絡中卷積核尺寸單一且容易忽略不同特征之間語義關系和丟失微小特征信息的問題，本文提出了一種帶三重注意力機制的特征增強分割模型（EFTA-Unet），該模型能夠充分利用醫學圖像中微小特征信息和捕獲不同特征之間的語義關系并且能夠克服醫學圖像數據集中類別不平衡的問題。

1 基于改進U-net的醫學圖像分割模型

本文提出了具有特征增強和注意力機制的醫學圖像分割模型（U-net with enhance feature and triplet attention module， EFTA-Unet），其框架如圖1所示。本文所提出的模型由編碼器和解碼器兩部分組成。在編碼器中通過級聯的卷積層以及池化層來提取醫學圖像的語義特征。解碼器則將這些特征恢復成像素級別的分割掩碼。此外，EFTA-Unet使用跳躍連接將編碼器中的特征與解碼器中的特征融合，減少編碼過程中信息丟失。

1.1 特征增強模塊

圖2是特征增強（enhance feature，EF）模塊的結構。當特征輸入時，會進行五個并聯操作，在前面四個并聯分支中，由兩個卷積層組成。第一層的卷積核尺寸為N×1，第二層的卷積核尺寸為1×N。本文使用這種級聯的非對稱卷積，相對于N×N的對稱卷積，能夠以較小的計算量捕獲同樣多的特征。第五個分支是一個跳躍連接。最后將這五個分支的輸出相加得到一個新的特征矩陣，新的特征矩陣經過卷積核大小為3×3和1×1的卷積層，最終得到整個分支的總輸出。EF模塊具有不同尺寸的感受野，使網絡具有更強的特征提取能力。

1.2 三重注意力機制（triplet attention，TA）

在通道注意力機制中，會對不同特征的權重進行均勻縮放。然而，這種方法需要通過全局平均池化將特征矩陣在空間上分解為單個像素來確定這些通道的權重。但是，由于它沒有考慮到通道和空間之間的相互關系，所以會導致大量的空間信息丟失。為了解決這個問題，三重注意力機制中提出跨維度交互的技術，通過建模空間維度和通道維度之間的相關性來更新權重，達到緩解空間信息的丟失的目的。在這里，三個分支分別計算（C，H）、（C，W）和（H，W）之間的相關性。圖3為三重注意力機制（triplet attention，TA）的結構，它具有三個并行分支。

如圖3所示，三重注意力機制的第一個分支計算通道維度C和空間維度W相關性，第二個分支計算通道C和高度H之間的相關性，最后一個分支計算高度H和寬度W之間的相關性，即空間注意力。三個分支輸出的特征矩陣尺寸相同。最后將三個分支的輸出相加取平均值作為整個模塊的輸出。在三層注意力機制中，Z-pool層的目的是減小特征維度數，Z-pool操作能夠在保留特征矩陣信息量的前提下降低特征矩陣的維度。Z-pool的數學表達式：

其中，T∈RC×H×W表示卷積層的輸出，C表示通道數，H表示張量的高度，W表示張量的寬度。MaxPool0dT表示在第0維執行最大池化操作，AvgPool0d（T）］表示在第0維執行平均池化操作。對于尺寸為（C×H×W）的張量，經過Z-pool操作后張量的尺寸變為（2×H×W）。

在第一個分支中，對于一個輸入特征T∈RC×H×W，進入第一個分支后，首先將特征矩陣沿H軸逆時針方向旋轉90°（[AKT^]*1）計算通道維度C與（H×W）之間的依賴關系，得到尺寸為（W×H×C）的特征矩陣。將特征矩陣經過Z-pool操作，得到尺寸為2×H×C的特征矩陣。再將這個特征矩陣經過一個標準化卷積層和一個歸一化操作，最終得到一個尺寸為（1×H×C）的特征矩陣。對這個輸出執行Sigmoid激活函數，獲得注意力權重。最后為了保證輸出矩陣與輸入矩陣尺寸一致，將得到的注意力權重矩陣沿著H軸順時針方向旋轉90°（[AKT^]1）。

第二個分支的作用是建模通道維度C和寬度維度W之間的相關性。首先將特征矩陣沿W軸逆時針旋轉90°（[AKT^]*2）得到尺寸為（H×C×W）的特征矩陣，再將這個特征矩陣經過Z-pool操作得到尺寸為（2×C×W）的特征矩陣并將這個特征矩陣經過一個標準卷積層和歸一化操作，得到形狀為（1×C×W）的注意力權重矩陣。最后，將得到的注意力權重經過Sigmoid操作和旋轉操作，最終得到第二個分支的注意力權重矩陣，并沿著W軸順時針方向旋轉90（2）。

如圖3所示，第三個分支與前兩個分支不同。這個分支不執行旋轉操作，對于輸入的特征矩陣直接執行Z-pool操作，這個分支最終輸出的特征矩陣形狀為（1×H×W）。最后，將三個分支的輸出融合成尺寸為（C×H×W）的輸出作為整個注意力模塊的輸出。對于輸入的特征矩陣T∈RC×H×W，當經過三重注意力機制后得到的注意力矩陣S的過程如下：

1.3 損失函數

圖像分割實際上是像素級別的分類任務，只需要將每個像素進行分類即可完成分割任務。對于腦膠質瘤圖像分割任務而言，由于整個圖像中除了腫瘤區域以外，其他的區域都可以算作是背景。所以二元交叉熵（binary cross entropy，BCE）可以作為損失函數的一部分，如

式中：yi表示數據中真實的標簽值，[AKy^]i表示分割模型的預測值，n為模型訓練的批次數。由于腦膠質瘤MRI中腫瘤區域只占整個圖像的2%左右，背景占比太大，存在像素類別不平衡的問題。如果只用二元交叉熵作為損失函數，則會導致網絡中在訓練過程中偏向學習背景特征而疏于學習腫瘤區域的特征，進而影響網絡的整體分割效果。而對于緩解類別不平衡問題最常用的損失函數為Dice，所以也可以將Dice作為本文的損失函數的一部分，其公式如：

式中：yi表示真實標簽值，[AKy^]i表示模型預測的值，ε是一個平滑算子，防止出現分子分母為0的情況。考慮到以上兩個方面，本文結合這兩個損失函數組成一個新的損失函數，公式如：

2 實驗設置與結果

2.1 數據集

為了驗證EFTA-Unet模型，本文在公共數據集BraTs上進行了多次實驗，并對實驗結果進行評估。使用BraTs2018的訓練集作為本文的訓練集，共285例樣本（HGG為210例，LGG為75例）。測試集則采用BraTs2019的訓練集在BraTs2018基礎上新增的50例樣本（HGG為49例，LGG為1例）。每一例樣本都包含四種模態，分別為Flair、T1、T1ce和T2。MRI的尺寸為244×244×155，前兩個數字尺寸表示圖像的寬度和高度，第三個數字表示切片的數量。本文將三維的MRI沿軸位切片，分解為155張圖片，然后丟棄黑色切片。需要注意的是，每位患者MRI切片中黑色圖像的數量不一致，所以每位患者腦MRI的有用切片數量也不一致。因此，對所有患者腦MRI處理后，用于實驗的訓練集圖像有18923張，測試集的圖像有3138張。用于實驗的圖像如圖4所示（綠色代表水腫，黃色代表增強腫瘤，紅色代表非增強腫瘤）。

腦膠質瘤的主要成分包括水腫（peritumoral edema，ED）、增強腫瘤（enhancing tumor，ET）、非增強腫瘤（non-enhancing tumor，NET）。本文的分割任務主要為整個腫瘤（WT，ED+ET+NET）， 增強腫瘤（ET），和核心腫瘤（TC，ET+NET）。三個分割任務中，核心腫瘤（TC）和增強腫瘤（ET）相對于整個腫瘤（WT）而言屬于小尺度腫瘤。

2.2 評價指標

本文使用Dice相關系數和豪斯多夫距離（Hausdorff95 distance）來衡量模型的性能。Dice相關系數計算公式如：

其中，TP為真陽性樣本數， FP為假陽性樣本數，FN為假陰性樣本數。

豪斯多夫距離（Hausdorff95 distance）測量的是真實區域與預測區域之間的距離，它對分割區域的邊界更為敏感，其計算公式如：

其中，sup表示最大值，inf表示最小值，T表示真實的腫瘤區域（Ground Truth），P表示預測的腫瘤區域， dt，p為距離函數，計算點t到p之間的距離。

2.3 實驗細節

所有實驗均在Intel（R） Core（TM） i910940X CPU @ 3.30GHz，64GB RAM，NVIDIA GeForce RTX 3090上進行。所有模型的優化器均為Adam，學習率為0.0001， 批量大小為8，本文的模型訓練了10000個epoch，損失函數中λ為0.5。

2.4 消融實驗

為驗證本文所提網絡（EFTA-Unet）的性能，進行了一系列的消融實驗。模型的基礎網絡為U-net，首先考慮只增加EF模塊，記為U-net+EF，表1為對比實驗結果。從表1中可以看出，引入EF模塊后平均Dice值提升了1.3%，性能更優。

在前面實驗的基礎上，繼續驗證TA模塊的有效性，名稱標記方式與前面相似，實驗結果如表2所示。從表2中可以看出，隨著模塊的增加，分割精度也在逐步提升，EFTA-Unet表現出了最佳的分割性能。EFTA-Unet的平均Dice值提升了1.7%，性能更優。

為了驗證本文損失函數的有效性，進行了一系列的對比試驗。首先本文單獨使用Dice和二元交叉熵作為損失函數，記為EFTA-Unet+Dice和EFTA-Unet+BCE，然后使用這兩個損失函數的線性組合作為對比試驗的損失函數，記為EFTA-Unet+Loss。實驗結果如表3所示，使用本文所提損失函數的模型性能明顯優于單獨使用兩個損失函數時的性能。

2.5 實驗結果

本文將EFTA-Unet與U-net[7]，DeepResUnet[14]， U-net++[10]， Dense_Unet[12]，" Unet3plus [13]在BraTs數據集上進行實驗，實驗結果如表4所示，圖5為分割結果可視圖。從表4中可以看出，EFTA-Unet模型在BraTs數據集上取得了0.8592的平均Dice和2.2498的平均HD，EFTA-Unet模型的分割性能明顯優于常用的模型。

從圖5中也可以看出，本文提出的模型在處理細節方面明顯優于其他網絡。整體來看，本文所提模型的分割性能最優。

3 總結

本文提出了一種名為EFTA-Unet的模型用于醫學圖像分割，該模型能夠學習醫學圖像中微小特征信息和建模不同特征間的相關性，并且能夠克服醫學圖像數據集中類別不平衡問題。本文提出的特征增強模塊（EF模塊）具有不同尺寸的卷積方式，以不同大小的感受野提取特征，在一定程度上防止微小特征信息丟失。其次，引入三重注意力機制，它通過通道注意和空間注意來捕捉跨維度的特征并建立相關性。此外，本文設計了一種新的損失函數用來解決醫學圖像中類別不平衡問題。最后，將本文提出的EFTA-Unet模型與幾種流行的醫學圖像分割模型在BraTs數據集上進行實驗，可以發現與其他流行的分割模型相比，EFTA-Unet模型在分割任務上表現最好。未來將進一步優化模型，希望在不影響分割性能的前提下減小模型的計算量。

［ 參 考 文 獻 ］

［1］ XIE X， NIU J， LIU X， et al. A survey on incorporating domain knowledge into deep learning for medical image analysis[J]. Medical Image Analysis， 2021， 69： 101985.

[2] WAN W， ZHONG Y， LI T， et al. Rethinking feature distribution for loss functions in image classification[C]∥Proceedings of the IEEE conference on computer vision and pattern recognition， 2018： 9117-9126.

[3] YANG H M， ZHANG X Y， YIN F， et al. Robust classification with convolutional prototype learning[C]∥Proceedings of the IEEE conference on computer vision and pattern recognition， 2018： 3474-3482.

[4] ZHENG Z， ZHONG Y， WANG J， et al. Foreground-aware relation network for geospatial object segmentation in high spatial resolution remote sensing imagery[C]∥Proceedings of the IEEE/CVF conference on computer vision and pattern recognition， 2020： 4096-4105.

[5] PATIL P W， BIRADAR K M， DUDHANE A， et al. An end-to-end edge aggregation network for moving object segmentation[C]∥proceedings of the IEEE/CVF conference on computer vision and pattern recognition， 2020： 8149-8158.

[6] LONG J， SHELHAMER E， DARRELL T. Fully convolutional networks for semantic segmentation[C]∥Proceedings of the IEEE conference on computer vision and pattern recognition， 2015： 3431-3440.

[7] RONNEBERGER O， FISCHER P， BROX T. U-net： Convolutional networks for biomedical image segmentation[C]∥International Conference on Medical image computing and computer-assisted intervention. Springer， Cham， 2015： 234-241.

[8] ZHANG Z， WU C， COLEMAN S， et al. DENSE-INception U-net for medical image segmentation[J]. Computer methods and programs in biomedicine， 2020， 192： 105395.

[9] YU F， WANG D， SHELHAMER E， et al. Deep layer aggregation[C]∥Proceedings of the IEEE conference on computer vision and pattern recognition， 2018： 2403-2412.

[10] XIAO X， LIAN S， LUO Z，et al. Weighted res-unet for high-quality retina vessel segmentation[C]. In： 2018 9th international conference on information technology in medicine and education （ITME）， IEEE， 2018： 327-331.

[11] HU J，SHEN L，ALBANIE S，et al．Squeeze-and-excitation networks[J]．IEEE Transactions on Pattern Analysis and Machine Intelligence，2020，43（08）：2011-2023.

[12] PAN Y， LIU J， XU T， et al. Hippocampal segmentation in brain mri images using machine learning methods： A survey[J]. Chinese Journal of Electronics， 2021， 30（05）： 793-814.

[13] HUANG H，LIN L，LONG R， et al. Unet 3+： A full-scale connected unet for medical image segmentation[C]. In： ICASSP 2020-2020 IEEE International Conference on Acoustics， Speech and Signal Processing （ICASSP），"" IEEE， 2020：1055-1059.

[14] ZHANG Z， LIU Q， WANG Y.Road extraction by deep residual u-net[J]. IEEE Geoscience and Remote Sensing Letters， 2018，15（05）： 749-753.

Medical Image Segmentation Model Based on Improved U-net

YANG Yongpu1，GAN Haitao1，XIA Wei2，YANG Zhi1，YE Zhiwei1

（1 School of Computer Science， Hubei Univ. of Tech.， Wuhan 430068， China；2 Department Wuhan Children’s Hospital （Wuhan Maternal and Child Healthcare Hospital），Tongji Medical College，Huazhong Univ. of Sci.amp; Tech.， Wuhan 430070， China）

Abstract： U-net model has made a lot of progress in medical image segmentation， but it ignores the gap between low-level visual features and high-level semantic features， which limits the feature extraction ability of the network. To solve this problem， a feature enhancement module is designed to enhance the feature extraction capability of the network with a larger receptive field， and a jump connection is used to fuse low-level features with high-level features in the network. At the same time， the triple attention mechanism is introduced to integrate the features of spatial dimension and channel dimension so as to enhance the useful information. Finally， a new loss function was designed， which introduced binary cross entropy and Dice loss function to suppress the category imbalance in the data set. The model was tested on the BraTs data set and the average Dice of 0.8592 was achieved. Experimental results show that the performance of this model is better than other common segmentation networks.

Keywords： medical image segmentation; feature enhancement; attention mechanism; loss function

［責任編校： 裴 琴］

［收稿日期］ 2022-09-29

［基金項目］ 湖北工業大學高層次人才基金（GCRC2020016）；湖北省自然科學基金（2021CFB282）；省部共建生物催化與酶工程國家重點實驗室開放基金項目（SKLBEE2021020 和 SKLBEE2020020）

［第一作者］ 楊永譜（1996-）， 男， 湖北黃岡人， 湖北工業大學碩士研究生， 研究方向為計算機視覺。

［通信作者］ 甘海濤（1985-）， 男， 江西豐城人， 湖北工業大學教授， 研究方向為機器學習、 計算機視覺和醫學影像處理。