作者簡介:普鐘(1997-),女(彝族),云南楚雄人,碩士研究生,主要研究方向為深度學習、醫學圖像處理;張俊華(1976-),女(通信作者),云南昆明人,博導,博士,主要研究方向為醫學影像處理和分析、模式識別等(jhzhang@ynu.edu.cn);黃昆(1998-),男,江西吉安人,碩士研究生,主要研究方向為醫學圖像處理、計算機視覺;周奇浩(1996-),男,云南曲靖人,碩士研究生,主要研究方向為醫學圖像處理、計算機視覺.
摘 要:針對脊椎CT、MR圖像分割模型分割性能不高的問題,基于U型網絡提出了脊椎分割網絡MAU-Net。首先引入坐標注意力模塊,使網絡準確捕獲到空間位置信息,并嵌入到通道注意力中;然后提出基于Transformer的雙支路通道交叉融合模塊代替跳躍連接,進行多尺度特征融合;最后提出特征融合注意力模塊,更好地融合Transformer與卷積解碼器的語義差異。在脊柱側凸CT數據集上,Dice達到0.929 6,IoU達到0.859 7。在公開MR數據集SpineSagT2Wdataset3上,與FCN相比,Dice提高14.46%。實驗結果表明,MAU-Net能夠有效減少椎骨誤分割區域。
關鍵詞:脊椎圖像分割;U型網絡;坐標注意力;雙支路通道Transformer;Transformer-Convolution融合注意力
中圖分類號:TP391.4;TP183 文獻標志碼:A
文章編號:1001-3695(2023)04-048-1256-07
doi:10.19734/j.issn.1001-3695.2022.07.0409
Abstract:Considering the existing spinal computer tomography(CT)and magnetic resonance(MR)image segmentation models have limitations in segmentation performance,this paper proposed a spinal segmentation method MAU-Net based on U-shaped network.Firstly,this paper introduced coordinate attention module into the encoder of U-shaped network,which enabled the network accurately capture the spatial position information and embedded it into the channel attention.Secondly,this paper proposed dual-branch channel cross fusion module based on Transformer,it could replace the skip connection for multi-scale feature fusion.Finally,this paper proposed a feature fusion attention module to better fuse the semantic differences between Transformer and Convolution network.On scoliosis CT dataset,Dice reached 0.929 6,IoU reached 0.859 7.On the public MR dataset SpineSagT2Wdataset3,compared with FCN,Dice improved by 14.46%.Experimental results show that this method can effectively reduce the 1 segmentation area of vertebrae.
Key words:spine image segmentation;U-shaped network;coordinate attention;dual-branch channel Transformer;Transformer-Convolution fusion attention
0 引言
脊椎圖像分割是椎骨定量分析的關鍵技術,是脊椎圖像配準[1]、三維重建[2]等后續研究的基礎和前提,直接影響模型的精準程度和重建效果[3]。
由于脊椎解剖結構復雜,周圍被肋骨、肌肉叢、韌帶、血管和神經等組織包繞[4],且椎骨不同軸位的影像形狀大小各異,以及病理變化、圖像偽影造成椎骨的外觀發生改變,脊椎圖像容易存在邊緣不清晰、對比度不明顯等問題,導致椎骨分割結果不準確。國內外研究者針對深度學習的圖像分割方法展開研究,目前存在兩類較為熱門的分割方法。
基于卷積神經網絡(convolutional neural network,CNN)搭建的U型網絡分割方法發展較為成熟。Ronneberger等人[5]提出了U-Net。Zhou等人[6]提出了UNet++,在U-Net之上,嵌套密集跳躍連接,搭建一個不同深度的U-Net集合,融合了編碼器與解碼器之間特征映射的語義差距。Oktay等人[7]提出了Attention U-Net(AttUNet),使用注意力門控模型,來抑制模型學習與任務無關的特征,加強與任務相關特征的學習。Jha等人[8]設計了Double U-Net,該網絡有兩個編碼器和兩個解碼器,同時在網絡兩個分支中都加入了空洞空間金字塔池化模塊[9]和壓縮激勵(squeeze-and-excitation,SE)[10]模塊,進一步提升U型網絡在不同分割任務中的性能。陳志等人[11]通過混合門控軸向自注意力和卷積模塊,搭建了雙編碼器和解碼器結構的多通道分割網絡,提升網絡特征提取能力。高紅霞等人[12]借鑒UNet3+[13],提出了單編碼—雙解碼結構的DCAW-Net,融合密度連接與自適應加權損失,設計了新的損失約束,并結合了VAE模塊,更好地提取分割目標的邊界和輪廓。上述基于CNN的網絡雖取得一定效果,但感受野有限,不利于捕獲全局信息。
近年基于Transformer架構搭建的U型網絡分割方法開始興起。Zhang等人[14]提出TransFuse,構建了BiFusion融合模塊,有效融合并行的CNN支路和Transformer支路。Cao等人[15]提出了Swin-UNet,基于Swin Transformer[16]模塊,搭建了類似U-Net的純Transformer對稱網絡,對局部和全局語義特征進行學習。Gao等人[17]提出了UTNet,采用了卷積與高效多頭自注意力模塊的混合層結構設計,結合相對位置編碼,降低模型復雜度。Hatamizadeh等人[18]提出了UNETR,采用收縮—膨脹模式,由純Transformer構成編碼器,提取編碼器不同層的序列,分別通過不同分辨率的跳躍連接與解碼器合并,捕獲全局多尺度特征。李擎等人[19]提出了雙分支U型網絡UConvTrans,將CNN分支與Transformer分支的輸出拼接,實現全局和局部特征交互融合。上述卷積和Transformer相結合的方法,雖然能取得一定效果,但忽略了空間位置信息與通道信息的融合,邊緣輪廓的分割效果有待提升。
為解決上述問題,本文提出融合多注意力的脊椎圖像分割網絡MAU-Net(multi-attention U-Net)。本文的主要工作和貢獻如下:a)在編碼部分引入了坐標注意力(coordinate attention,CA)[20]模塊,利用位置信息對通道信息進行編碼,更好地分割椎骨邊緣輪廓和微小細節;b)設計了雙支路通道Transformer(Dual-branch Channel Transformer,DCT)模塊代替跳躍連接,增強圖像的全局特征提取能力;c)提出了特征融合注意力模塊T-CFA(Transformer-Convolution fusion attention),引導Transfor-mer操作的特征層和卷積操作的特征層更好地融合。
1 本文方法
1.1 MAU-Net結構
針對脊椎CT和MR圖像的數據特點,本文設計了MAU-Net,以U-Net作為主體結構,主要包括編碼路徑、坐標注意力模塊、DCT模塊、T-CFA模塊、解碼路徑。網絡結構如圖1所示,圖中D為編碼階段下采樣前的層級,T為經過DCT模塊生成的特征層,U為解碼階段卷積上采樣后的層級。
在編碼器中,分為4個層級,每個層級包括2個卷積層和1次2×2的最大池化下采樣。解碼器也分為4個階段,每個階段包括2個卷積層和1次2×2的上采樣操作。首先,將編碼器第1階段的第2個3×3卷積替換為坐標注意力模塊。大小為224×224×3的圖像經過一次卷積操作來提取初始特征,經過坐標注意力模塊,以捕獲椎骨圖像的空間長程依賴和空間位置信息,更好地學習底層的細微特征。然后,將跳躍連接部分替換為DCT模塊,沿著通道軸做Transformer操作,使用面向通道的特征融合注意力模塊T-CFA對捕獲到的信息進行通道交叉融合,以消除通道Transformer和解碼器卷積操作之間的語義差異。最后,將T-CFA模塊校準后的特征層,與解碼卷積得到的特征層拼接后進行上采樣操作。CNN根據像素點之間的運算對圖像進行局部感知,具有平移不變性。而Transformer結構的多頭自注意力機制可以構建像素間長距離的依賴關系,具有強大的全局信息提取能力。脊柱胸椎和肋骨相連,神經、血管、肌肉叢密集分布,各類組織邊界之間比較模糊,有效地處理局部和全局信息能夠對分割精度提升帶來很大幫助。利用Transformer操作和卷積操作的互補性,融合全局和局部信息,使網絡更好地提取椎骨邊緣特征,提高模型的分割能力。
DCT模塊共有兩個輸入,分別對輸入的兩個特征層進行通道的Transformer操作。首先給每個特征層切分patch,切分時要保證編碼器不同層級的patch在通道維度操作時能夠映射到特征相同的區域。然后送入多頭通道交叉注意力模塊,進行多頭注意力操作。最后進入殘差結構的多層感知器中,DCT模塊包含多組Transformer注意力操作。
相比于ViT網絡中Transformer的單輸入單輸出結構實現同一個特征層特征提取,本文所提出的DCT模塊結構設計了雙輸入雙輸出結構,能引導輸入的兩個特征層從通道角度進行多尺度交叉融合,捕獲到更多全局依賴關系。
1.3.3 多頭通道注意力模塊改進
為了使DCT模塊能從通道維度同時對兩個不同特征層進行Transformer操作,受多頭注意力結構啟發,本文設計了多頭通道注意力機制(multi-head cross attention,MCA)。
對比原始注意力機制在特征層被劃分為若干個patch后直接對patch進行注意力操作,本文設計的DCT模塊是從通道軸進行注意力操作,將4個特征層分為兩組,每組根據通道數分別劃分patch后送入DCT模塊中,跨層的兩個特征層沿著通道軸進行Transformer操作,并且使用了實例歸一化。圖7展示了原始自注意力機制和本文的DCT模塊注意力的區別。
2 實驗與分析
2.1 實驗數據集
本文在兩個脊椎病理數據集上評估網絡,包括脊柱側凸CT數據集和MR數據集SpineSagT2Wdataset3。圖像受到與椎骨相鄰的非椎骨組織(肌肉組織、肋骨等)影響,部分椎骨的CT圖像的邊緣輪廓不清晰。因為脊椎發生病變導致椎骨受到擠壓,部分椎骨發生形變,相鄰椎骨的圖像重疊,以上問題都可能影響脊椎圖像的分割效果。
數據集1 實驗使用的脊柱側凸CT數據集如圖10所示。數據由昆華醫院提供,共包含2 385張已消除病例信息的脊柱側凸CT軸位面圖,包括人體12個胸椎和5個腰椎在內的17個椎骨圖像,標簽由4名經驗豐富的專業醫師共同標注。實驗中按照8:1:1的比例將訓練集、驗證集和測試集分別隨機劃分為1 908張、238張和239張,圖像大小均為512×512。
數據集2 實驗使用的脊椎MR數據集為競賽數據集SpineSagT2Wdataset3,數據均為矢狀位T2加權的MR數據,由同一臺磁場強度為3.0 T的MR設備采集,包括人體部分胸椎、5個腰椎以及骶骨尾骨圖像,共215位患者(腰椎間盤突出、腰椎間胖退行性疾病等)的脊椎MR圖像,包含195套帶標簽數據和20套無標簽數據。實驗使用帶標簽的195套MR圖像,共2 460張,將圖像隨機劃分為2 360張訓練數據和100張測試數據,圖像大小不一,訓練階段統一調整為224×224。
2.2 實驗環境與參數設置
實驗使用Windows 10系統和PyCharm編輯器,在型號為NVIDIA GeForce RTX 2080Ti的GPU計算平臺上進行,CPU型號為AMD Ryzen7 3800X 8-Core processor,內存為8 GB×2,所有程序由PyTorch1.8.0框架以及Python 3.8接口實現(本文代碼可在https://github.com/puzhong-hug/MAUNet中獲取)。
本文在實驗訓練階段,將輸入的圖像大小統一調整為224×224,并且通過像素歸一化至[0,1]。在網絡的訓練過程中使用動量為0.9的Adam優化器,學習率的初始值設置為10-3,權重衰減系數為0.000 5。在每個epoch前,對輸入圖像做隨機翻轉處理以減輕網絡過擬合。本文網絡batchsize設置為8。訓練過程中使用早停法,patience設置為50,當Dice系數連續50個epoch內不再提升,則自動停止迭代過程。實驗通過驗證集選取最佳權重模型。
在實現過程中,經DCT模塊,編碼器和解碼器的參數變化如表1所示。DCT1的PatchSize設置為(32,8),DCT2的PatchSize為(16,4)。每個DCT模塊Transformer層數設置為4。MCA頭數都設置為4,dropout率為0.1,維度擴展比為4。Embedding操作的dropout率設置為0.1。DCT模塊沿著通道軸操作,使得梯度平滑的傳播,更好地發揮Transformer長程依賴建模優勢。因為Transformer操作本身并不改變特征層的通道數,所以編碼端的特征層在通過DCT模塊以后通道數不變。
觀察五組實驗可知,與原U型網絡相比,加入本文提出的各個模塊后,評價指標均得到不同程度提升。其中,U型網絡中同時使用坐標注意力、DCT模塊和T-CFA模塊時分割效果最好,相比U-Net,Dice分別提升2.27%、2.82%,IoU分別提升2.28%、3.98%,分割性能優于其他方案。結果表明,加入坐標注意力模塊后,網絡結合了脊椎特征自身重要性和特征位置信息有助于提高分割精度。Transformer提取的全局信息特征補充到解碼端比直接的跳躍連接能夠有效提升分割精度。加入T-CFA模塊,進一步提升DCT模塊與解碼端CNN的融合效果,使得分割結果進一步提升。
2.4.4 脊椎CT圖像分割
為了驗證本文網絡的性能和分割效果,在脊椎側凸CT數據集上,將本文網絡MAU-Net與FCN[26]、U-Net[5]、UNet++[6]、DeeplabV3+[27]、AttUNet[7]和Swin-UNet[15]的分割結果進行對比,Dice分別提升了4.34%,2.27%,1.39%,3.69%,0.35%,1.48%,IoU分別提升了6.05%,2.28%,1.35%,4.56%,0.6%,1.82%,各方法得到的分割性能指標如表5所示。實驗結果表明本文方法能夠有效提升脊椎CT圖像分割的精度。
本文提出的網絡與多個網絡的分割結果進行對比,脊椎側凸CT數據集的分割視覺效果如圖11所示,每列分別顯示原圖、原圖對應的真實標簽以及對比網絡的分割圖片。
方框內為分割結果較差的區域,可以看出,UNet++分割結果存在椎骨邊緣的誤分割。AttUNet較UNet++分割缺失區域少。Swin-UNet分割結果存在部分椎骨邊緣不清晰、椎骨分割錯誤信息。本文網絡得到的分割圖像椎骨邊緣輪廓清晰,誤分割區域更小,能夠有效地消除肋骨以及其他組織造成的影響,在保留椎骨細微形狀信息的同時實現更精細的分割。本文方法與對比模型的分割視覺效果圖相比更加接近真實標簽,可以產生更好的分割結果,能夠有效解決椎骨病變導致的椎骨形變、邊緣模糊等問題。
2.4.5 脊椎MR圖像分割
為了進一步驗證本文網絡性能,本文選取了脊椎MR分割數據集SpineSagT2Wdataset3進行實驗。實驗分割結果如表6所示,其中FCN、U-Net、UNet++、DeeplabV3+使用了于文濤等人[28]的實驗數據。
在SpineSagT2Wdataset3數據集上,本文方法較FCN、U-Net、UNet++、DeeplabV3+、文獻[28]、AttUNet和Swin-UNet,Dice分別提升了14.46%,8.7%,7.72%,5.6%,3%,4.56%,1.92%,IoU分別提升14.06%,8.48%,5.88%,3.48%,2.28%,2.7%,1.76%。與表中其他模型對比,脊椎MR數據集的分割結果均得到提升。本文網絡在脊椎MR圖像分割任務中具有較好的泛化性能。
本文網絡與多個網絡在脊椎MR圖像上的分割視覺效果分別如圖12所示。
可以觀察到,UNet++中幾幅圖存在椎骨的缺失。AttUNet中雖存在分割信息的缺失,但是缺失區域比UNet++少,分割效果更好。Swin-UNet有部分椎骨粘連,但分割錯誤信息較少。本文方法幾乎不存在椎骨分割缺失問題,誤分割區域更少,邊緣輪廓更加清晰,分割后圖片內部連貫性更好。
2.5 討論
綜合客觀評價指標和分割效果圖可知,本文網絡在Dice系數、IoU評價指標上獲得了最優的表現,且在視覺效果上優于文中對比方法。本文方法利用輕量級坐標注意力模塊聯系圖像上下文特征信息,獲取位置信息和通道信息,在通道維度進行注意力融合,以擴大感受野,專注于目標分割區域,從而有效地解決了脊椎圖像椎骨形變重疊、解剖結構復雜造成的椎骨分割缺失等影響分割效果的問題。
3 結束語
本文提出了一種融合多注意力機制的網絡模型,該網絡引入坐標注意力模塊和DCT模塊,捕獲空間維度相關性和通道維度依賴性,使網絡準確地利用重要通道信息。同時多尺度特征信息融合使網絡在不同的通道維度自適應地選擇感受野,將層級間捕獲到的通道信息進行交叉融合,解決椎骨邊緣模糊造成的分割不準確等問題。在脊椎側凸數據集、SpineSagT2Wdataset3數據集上,本文方法的Dice分別達到了0.929 6、0.891 9,IoU分別達到了0.859 7、0.807 9,分割效果優于對比方法。未來研究中,可以通過設計更輕量、視野更大的Transformer模塊,尋求Transformer與卷積操作之間更高效的結合方式。
參考文獻:
[1]Valsasina P,Horsfield M A,Meani A,et al.Improved assessment of longitudinal spinal cord atrophy in multiple sclerosis using a registration-based approach:relevance for clinical studies[J].Journal of Magnetic Resonance Imaging,2022,55(5):1559-1568.
[2]Liu Zheng,Su Zhihai,Wang Min,et al.Computerized characterization of spinal structures on MRI and clinical significance of 3D reconstruction of lumbosacral intervertebral foramen[J].Pain Physician,2022,25(1):E27.
[3]Reyneke C J F,Lyuthi M,Burdin V,et al.Review of 2-D/3-D reconstruction using statistical shape and intensity models and x-ray image synthesis:toward a unified framework[J].IEEE Reviews in Biomedical Engineering,2018,12:269-286.
[4]Telera S,Raus L,Pipola V,et al.Surgical and radiological anatomy of the thoracolumbar spine for vertebroplasty and kyphoplasty[M]//Vertebral Body Augmentation,Vertebroplasty and Kyphoplasty in Spine Surgery.Berlin:Springer,2021:19-37.
[5]Ronneberger O,Fischer P,Brox T.U-Net:convolutional networks for biomedical image segmentation[C]//Proc of International Conference on Medical Image Computing and Computer-Assisted Intervention.Berlin:Springer,2015:234-241.
[6]Zhou Zongwei,Rahman S M M,Tajbakhsh N,et al.UNet+:a nested U-Net architecture for medical image segmentation[M]//Deep Lear-ning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support.Berlin:Springer,2018:3-11.
[7]Oktay O,Schlemper J,Folgoc L L,et al.Attention U-Net:learning where to look for the pancreas[EB/OL].(2018-05-20).https://arxiv.org/abs/1804.03999.
[8]Jha D,Riegler M A,Johansen D,et al.DoubleU-Net:a deep convolutional neural network for medical image segmentation[C]//Proc of the 33rd IEEE International Symposium on Computer-Based Medical Systems.Piscataway,NJ:IEEE Press,2020:558-564.
[9]Chen L C,Papandreou G,Schroff F,et al.Rethinking atrous convolution for semantic image segmentation[EB/OL].(2017-12-05).https://arxiv.org/abs/1706.05587.
[10]Hu Jie,Shen Li,Sun Gang.Squeeze-and-excitation networks[C]//Proc of IEEE Conference on Computer Vision and Pattern Recognition.Piscataway,NJ:IEEE Press,2018:7132-7141.
[11]陳志,李歆,林麗燕,等.引入門控軸向自注意力的多通道病理圖像分割[J/OL].計算機應用.(2022-09-02)[2022-09-18].http://kns.cnki.net/kcms/detail/51.1307.tp.20220902.0906.004.html.(Chen Zhi,Li Xin,Lin Liyan,et al.Multi-channel pathological image segmentation with gated axial self-attention[J/OL].Journal of Computer Applications.(2022-09-02)[2022-09-18].http://kns.cnki.net/kcms/detail/51.1307.tp.20220902.0906.004.html.)
[12]高紅霞,郜偉.融合密集連接與自適應加權損失的血管壁圖像分割[J].計算機應用研究,2022,39(6):1905-1910.(Gao Hong-xia,Gao Wei.Vessel wall image segmentation based on dense connection and adaptive weighted loss[J].Application Research of Computers,2022,39(6):1905-1910.)
[13]Huang Huimin,Lin Lanfen,Tong Ruofeng,et al.UNet 3+:a full-scale connected U-Net for medical image segmentation[C]//Proc of IEEE International Conference on Acoustics,Speech and Signal Processing.Piscataway,NJ:IEEE Press,2020:1055-1059.
[14]Zhang Yundong,Liu Huiye,Hu Qiang.TransFuse:fusing transformers and CNNs for medical image segmentation[C]//Proc of International Conference on Medical Image Computing and Computer-Assisted Intervention.Berlin:Springer,2021:14-24.
[15]Cao Hu,Wang Yueyue,Chen J,et al.Swin-UNet:UNet-like pure transformer for medical image segmentation[EB/OL].(2021-05-12).https://arxiv.org/abs/2105.05537.
[16]Liu Ze,Lin Yutong,Cao Yue,et al.Swin Transformer:hierarchical vision transformer using shifted windows[C]//Proc of the IEEE/CVF International Conference on Computer Vision.Piscataway,NJ:IEEE Press,2021:10012-10022.
[17]Gao Yunhe,Zhou Mu,Metaxas D N.UTNet:a hybrid Transformer architecture for medical image segmentation[C]//Proc of International Conference on Medical Image Computing and Computer-Assisted Intervention.Berlin:Springer,2021:61-71.
[18]Hatamizadeh A,Tang Yucheng,Nath V,et al.UNETR:transformers for 3D medical image segmentation[C]//Proc of IEEE/CVF Winter Conference on Applications of Computer Vision.Piscataway,NJ:IEEE Press,2022:574-584.
[19]李擎,皇甫玉彬,李江昀,等.UConvTrans:全局和局部信息交互的雙分支心臟圖像分割[J/OL].上海交通大學學報.(2022-09-16)[2022-09-21].http://doi.org/10.16183/j.cnki.jsjtu.2022.088.(Li Qing,Huangfu Yubin,Li Jiangyun,et al.UConvTrans:a dual-flow cardiac image segmentation network by global and local information integration[J/OL].Journal of Shanghai Jiaotong University.(2022-09-16)[2022-09-21].http://doi.org/10.16183/j.cnki.jsjtu.2022.088.)
[20]Hou Qibin,Zhou Daquan,Feng Jiashi.Coordinate attention for efficient mobile network design[C]//Proc of IEEE/CVF Conference on Computer Vision and Pattern Recognition.2021:13713-13722.
[21]Woo S,Park J,Lee J Y,et al.CBAM:convolutional block attention module[C]//Proc of European Conference on Computer Vision.Berlin:Springer,2018:3-19.
[22]Vaswani A,Shazeer N,Parmar N,et al.Attention is all you need[EB/OL].(2017-06-12).https://arxiv.org/abs/1706.03762.
[23]Dosovitskiy A,Beyer L,Kolesnikov A,et al.An image is worth 16×16 words:transformers for image recognition at scale[EB/OL].(2021-06-03).https://arxiv.org/abs/2010.11929.
[24]Chen Jieneng,Lu Yongyi,Yu Qihang,et al.TransUNet:transformers make strong encoders for medical image segmentation[EB/OL].(2021-02-08).https://arxiv.org/abs/2102.04306.
[25]Valanarasu J M J,Oza P,Hacihaliloglu I,et al.Medical transformer:gated axial-attention for medical image segmentation[C]//Proc of International Conference on Medical Image Computing and Computer-Assisted Intervention.Berlin:Springer,2021:36-46.
[26]Long J,Shelhamer E,Darrell T.Fully convolutional networks for semantic segmentation[C]//Proc of IEEE Conference on Computer Vision and Pattern Recognition.Piscataway,NJ:IEEE Press,2015:3431-3440.
[27]Azad R,Asadi-Aghbolaghi M,Fathy M,et al.Attention DeeplabV3+:multi-level context attention mechanism for skin lesion segmentation[C]//Proc of European Conference on Computer Vision.Berlin:Springer,2020:251-266.
[28]于文濤,張俊華,梅建華,等.脊柱MR圖像自動分割方法的研究[J].計算機工程與應用,2022,58(22):203-209.(Yu Wentao,Zhang Junhua,Mei Jianhua,et al.Research on automatic segmentation method of spinal MR images[J].Computer Engineering and Applications,2022,58(22):203-209.)