基于Unet的多注意力腦腫瘤圖像分割算法

吳 量，付殿臣，程 超

(長春工業(yè)大學(xué) 計算機(jī)科學(xué)與工程學(xué)院，吉林 長春 130012)

0 引 言

醫(yī)學(xué)圖像分割是醫(yī)學(xué)圖像分析領(lǐng)域的一個重要內(nèi)容,通常所說的分割指的是語義分割[1],即通過對圖像的每一個像素標(biāo)記上語義信息,來實現(xiàn)像素級別的預(yù)測,屬于計算機(jī)視覺的范疇。與自然圖像分割相比,醫(yī)學(xué)圖像分割任務(wù)由于邊界模糊、數(shù)據(jù)集小、語義簡單等特點,實現(xiàn)的難度會更大。為了提高醫(yī)學(xué)圖像分割的效率和準(zhǔn)確性,人們開展了越來越多的嘗試,其中一個很重要的方向就是腦腫瘤MRI圖像的分割[2]。近年來,腦腫瘤在社會上已經(jīng)成為一種常見的疾病,對人們的生命造成了巨大的威脅。腦腫瘤個體間形狀、紋理差異較大,很難提升分割的精度。目前,腦腫瘤醫(yī)學(xué)診斷主要依賴于醫(yī)生個人知識的積累和實戰(zhàn)經(jīng)驗的判斷,這種依賴存在較大的潛在問題,比如耗時長,或在高強(qiáng)度的重復(fù)工作條件下醫(yī)生過度疲勞而由此引發(fā)的漏診和誤診。因此,構(gòu)建一種精確、快速的腦腫瘤分割方法對腦瘤的診斷和早期治療具有很重要的意義。

1 相關(guān)工作

目前腦腫瘤分割算法有很多,主要分為傳統(tǒng)圖像分割方法和基于深度學(xué)習(xí)的分割方法。

1.1 傳統(tǒng)圖像分割方法

在傳統(tǒng)的圖像分割中,單獨使用一種分割方法得到的分割精度較低,研究者經(jīng)常將傳統(tǒng)方法與其他的方法結(jié)合使用以提高分割的精度。

閾值法通過設(shè)定不同的閾值將分割目標(biāo)與背景的像素區(qū)分,達(dá)到分割的目的,特點是計算速度快,但是適用性有限。賴小波等人結(jié)合傳統(tǒng)機(jī)器學(xué)習(xí)方法隨機(jī)森林和閾值法提出了一種多模態(tài)MRI圖像和多特征融合的腦腫瘤分割方法[3],首先進(jìn)行圖像校準(zhǔn)和偏置場校正,融合提取的多模態(tài)特征構(gòu)建RF模型進(jìn)行粗分割,最后根據(jù)先驗知識設(shè)定閾值進(jìn)行分割?；趨^(qū)域的方法優(yōu)點是簡單，能夠正確分割具有相似屬性的區(qū)域,缺點是容易造成過度分割。候冬奧等人結(jié)合區(qū)域生長法提出了一種全自動的腦腫瘤分割方法[4],該方法利用人腦結(jié)構(gòu)對稱信息,通過區(qū)域生長法實現(xiàn)粗略分割,然后使用改進(jìn)的測地線活動輪廓(geodesic active contours，GAC)模型對粗略分割區(qū)域進(jìn)行精細(xì)分割。童云飛等人采用快速模糊C均值聚類和自動閾值法對腫瘤區(qū)域進(jìn)行粗略輪廓提取[5],而后使用混合水平集算法來提取病灶的邊界。張騰達(dá)等人針對傳統(tǒng)水平集(Level Set)方法對腫瘤MRI圖像進(jìn)行分割時弱邊緣分割缺失的現(xiàn)象[6],采用模糊聚類算法與Level Set結(jié)合的方式增強(qiáng)了傳統(tǒng)方法對腦腫瘤分割的魯棒性和準(zhǔn)確性。

1.2 基于深度學(xué)習(xí)的圖像分割方法

近年來,深度學(xué)習(xí)技術(shù)快速發(fā)展,人們結(jié)合深度學(xué)習(xí),提出了很多優(yōu)秀的腦腫瘤MRI圖像分割方法。研究者首先應(yīng)用卷積神經(jīng)網(wǎng)絡(luò)(CNN)對語義分割進(jìn)行了探索,并將CNN應(yīng)用于腦腫瘤MRI圖像分割[7],得到了良好的分割結(jié)果。但是,由于CNN本身的輸出全連接層結(jié)構(gòu),模型在分割圖像時耗費巨大的內(nèi)存,計算效率很低。另外,CNN只可以對輸入圖像進(jìn)行像素的簡單分類,且全連接層限制了輸入圖像必須固定尺寸大小,通常模型的預(yù)測圖像較輸入圖像的分辨率會小一些,這大大影響了分割的效率和精度,因此得到的預(yù)測分割結(jié)果往往不盡如人意。

為解決這一問題,2015年Long等人[8]提出了一種專門適用于圖像分割任務(wù)的全卷積神經(jīng)網(wǎng)絡(luò)(FCN)模型。FCN模型不僅避免了CNN全連接層在計算時重復(fù)冗余卷積造成的存儲開銷過大的問題,也解決了圖像尺度的可變換問題,使得網(wǎng)絡(luò)可以接受任意大小的圖像,模型的輸出圖像可以恢復(fù)到與原圖大小的分辨率,實現(xiàn)了圖像在真正意義上的語義分割。然而,醫(yī)學(xué)圖像數(shù)據(jù)對細(xì)節(jié)較為敏感,傳統(tǒng)FCN的分割精度是遠(yuǎn)遠(yuǎn)不夠的。邢波濤等人基于FCN模型[9],將多模態(tài)MRI圖像進(jìn)行預(yù)處理和特征融合,通過結(jié)合條件隨機(jī)場(CRF)來細(xì)化FCN的粗病灶分割結(jié)果。針對醫(yī)學(xué)圖像的分割任務(wù),Ronneberger等人[10]基于FCN提出更加適用于醫(yī)學(xué)圖像分割的Unet模型。與FCN模型相比,Unet模型結(jié)構(gòu)對稱,其特有的Skip Connection結(jié)構(gòu)充分結(jié)合底層空間特征和高層語義特征,對圖像的細(xì)節(jié)特征提取更加敏感,針對醫(yī)學(xué)圖像的小目標(biāo)分割更加精確。Dong等人[11]將Unet模型應(yīng)用于腦膠質(zhì)瘤的分割,在BraTS 2015[12]腦腫瘤MRI數(shù)據(jù)集上取得了優(yōu)良的分割效果。

因Unet在醫(yī)學(xué)圖像分割任務(wù)中表現(xiàn)出的優(yōu)良性能,近幾年來提出的醫(yī)學(xué)圖像分割模型[13-17]大多基于U型對稱結(jié)構(gòu)和應(yīng)用Unet的特征融合方式,研究者分別在模型深度和特征融合方式等方面對Unet進(jìn)行了改進(jìn)。Ibtehaz等人[14]通過結(jié)合深度殘差網(wǎng)絡(luò)(ResNet)和Unet模型,加深了模型的深度，并將算法應(yīng)用于醫(yī)學(xué)圖像分割,在5個醫(yī)學(xué)數(shù)據(jù)集上較Unet提升了分割的精度。最近,Cao等人[15]充分利用密集連接網(wǎng)絡(luò)[16](DenseNet)特征融合特性,將密集連接(Dense Connection)與Unet模型結(jié)合應(yīng)用電磁圖像的分割任務(wù),并通過實驗驗證了DenseUnet分割性能的提升。醫(yī)學(xué)圖像多為3D數(shù)據(jù),2D卷積不能充分利用3D數(shù)據(jù)的空間結(jié)構(gòu)信息,而3D卷積具有較高的計算成本和GPU內(nèi)存消耗。為了解決這些問題,Li等人[17]提出了一種新的混合密集連接的Unet,并取得了優(yōu)良的分割指標(biāo)。針對Unet不同大小的數(shù)據(jù)集自適應(yīng)層數(shù)問題,Zhou等人[13]提出了Unet++模型,模型可根據(jù)不同的數(shù)據(jù)集自適應(yīng)調(diào)整網(wǎng)絡(luò)的提取層數(shù)和減少訓(xùn)練參數(shù)量。隨著深度學(xué)習(xí)理論的完善,人們對不同領(lǐng)域的任務(wù)也進(jìn)行了交叉研究。有研究者嘗試把與人類視覺機(jī)制相符合的注意力機(jī)制應(yīng)用于圖像處理領(lǐng)域。Li等人[18]對注意力門輸出層的通道結(jié)合方式進(jìn)行了改進(jìn),并用DenseNet結(jié)合基礎(chǔ)的Unet模型。該文基于Unet重新設(shè)計網(wǎng)絡(luò)提取模型并結(jié)合改進(jìn)后的多注意力機(jī)制,提出了一種基于Unet的多注意力腦腫瘤MRI圖像自動分割算法。

該文的主要創(chuàng)新點總結(jié)如下:

(1)提出了一種快速、準(zhǔn)確、端到端的腦腫瘤MRI圖像自動分割方法。

(2)在主干特征提取網(wǎng)絡(luò)部分,重新設(shè)計Unet網(wǎng)絡(luò)的結(jié)構(gòu),模型與ResBlock結(jié)合并且增加并行Dalited Convolution特征提取模塊,提升了網(wǎng)絡(luò)特征提取的性能。

(3)針對小型病灶分割精度低的問題,引入空間注意力和層注意力多注意力機(jī)制,使得特征提取網(wǎng)絡(luò)更加注意某些特征層和空間區(qū)域的關(guān)鍵特征,抑制了冗余區(qū)域的無效特征。

(4)在BraTS數(shù)據(jù)集上對比了其他分割網(wǎng)絡(luò),驗證了該文的模型,得到了較好的分割效果。

2 算法設(shè)計

文中算法對四種模態(tài)圖像數(shù)據(jù)進(jìn)行并行處理,最后得到多類腫瘤的分割預(yù)測結(jié)果,整體的算法流程結(jié)構(gòu)示意圖如圖1所示。

圖1 算法流程示意圖

2.1 主干特征提取網(wǎng)絡(luò)的設(shè)計

模型基于經(jīng)典的U型網(wǎng)絡(luò),去掉了Unet的池化層并整體采用ResBlock[19]對網(wǎng)絡(luò)進(jìn)行加深處理,使得網(wǎng)絡(luò)能夠充分融合淺層和深層語義信息。為保證引入ResBlock后不增加大量的參數(shù),模型去除了兩層網(wǎng)絡(luò),將各層的通道數(shù)調(diào)整為原來的1/2。另外,模型增加了一條并行的特征提取模塊,提取的特征與最后一層下采樣層進(jìn)行特征疊加?？紤]到大幅度的下采樣會導(dǎo)致有效特征的丟失,在特征融合前,并行模塊采用了膨脹系數(shù)為2的Dalited Convolution[20]對第一層提取到的特征進(jìn)行了二次特征提取,通過增大特征提取卷積的感受野,使得特征下采樣過程更加平滑。并行的特征提取網(wǎng)絡(luò)結(jié)構(gòu)如圖2所示。

圖2 網(wǎng)絡(luò)模型結(jié)構(gòu)

2.2 多注意力機(jī)制Masc模塊

該文提出的Masc注意力機(jī)制模塊,延續(xù)了傳統(tǒng)注意力機(jī)制對輸入圖像進(jìn)行通道級別的注意力特征圖像級聯(lián)空間注意力模塊的設(shè)計思想。在此基礎(chǔ)上對模型的輸入特征單獨應(yīng)用空間注意力模塊進(jìn)行空間注意特征提取,對兩次注意力機(jī)制獲取的空間注意力特征做了特征融合,使模型更加注重提取圖像的二維空間信息。

注意力模塊的結(jié)構(gòu)如圖3所示,具體的算法流程接下來會詳細(xì)描述。假設(shè)Masc模塊的輸入是任意一個中間特征F∈RH*W*C,F首先同時輸入Masc模塊的通道注意力模塊和空間注意力模塊,分別得到一維通道注意圖Fc∈R1*1*C和二維空間注意圖Fs2∈RH*W*1,在得到Fc之后,Fc會輸入到級聯(lián)的空間注意力模塊得到Fs1∈RH*W*1,最后采取特征相疊加的特征融合方式合并Fs1和Fs2得到最終的注意圖Fsc∈RH*W*1?？傮w的注意圖提取過程可以公式化如下：

圖3 Masc注意力機(jī)制模塊

Fc=Mc(F∈RH*W*C)

(1)

Fs1=Ms1(F∈RH*W*C)

(2)

Fs2=Ms2(Fsc∈R1*1*C)

(3)

Fsc=Fs1+Fs2

(4)

其中,M*(*)表示特征圖進(jìn)入相應(yīng)模塊進(jìn)行提取注意特征與原圖進(jìn)行矩陣元素相乘運算。Fsc∈RH*W*1是最終注意力模塊的輸出結(jié)果。

2.2.1 通道注意模塊

(5)

其中，σ1表示softmax激活函數(shù),σ0表示relu激活函數(shù),w*,b*分別表示卷積層的權(quán)重和偏置,? 表示元素乘積運算。下面公式參數(shù)同式(5)。

2.2.2 空間注意模塊

(6)

GlobalMaxPool(Fi)))?Fi=

(7)

(8)

該文的Masc模塊分別對輸入的圖像進(jìn)行通道注意力和雙空間注意力的提取,更加注重圖像2維空間的信息。

3 實 驗

3.1 數(shù)據(jù)集介紹及預(yù)處理

數(shù)據(jù)集采用的是公開數(shù)據(jù)集BraTS 2019腦膠質(zhì)瘤MRI數(shù)據(jù),其中包含259例高級別膠質(zhì)瘤(HGG)和76例低級別膠質(zhì)瘤(LGG)兩類患者數(shù)據(jù)。HGG患者包括水腫(edema),壞死(necrosis),增生(enhancing tumor)腫瘤,非增生(non-enhancing tumor)腫瘤部分,而LGG患者增生(enhancing tumor)可能存在也可能不存在。BraTS 2019數(shù)據(jù)集[12]包括4類標(biāo)簽,如圖4(e)所示,分別是正常組織(標(biāo)簽0)、壞死和非增生腫瘤(標(biāo)簽1)、水腫(標(biāo)簽2)、增生腫瘤(標(biāo)簽4)。分割任務(wù)包括三個：(1)完整腫瘤區(qū)域(Complete tumor)，包括所有三個腫瘤內(nèi)類,標(biāo)簽是1、2和4。(2)核心腫瘤區(qū)域(Core tumor)，不包括“edema”區(qū)域,標(biāo)簽是1和4。(3)增強(qiáng)腫瘤區(qū)域(Enhancing tumor),只包括標(biāo)簽4。實驗將BraTS 2019的訓(xùn)練數(shù)據(jù)分成訓(xùn)練集、驗證集和測試集,比例分別是0.6、0.2和0.2。原始數(shù)據(jù)集是由T1、T2、T1C和FLAIR四個序列構(gòu)成的3D數(shù)據(jù),如圖4(a)～圖4(d)所示。每例數(shù)據(jù)由155張厚度為1 mm×1 mm×1 mm,大小為240×240的切片組成。在數(shù)據(jù)預(yù)處理方面,該文取3D數(shù)據(jù)的2D切片,只對數(shù)據(jù)做了主要部分的裁剪和空白切片的去除,具體的做法是將(155,240,240)大小的切片處理成(146,192,152)大小,除此之外未做任何的數(shù)據(jù)增強(qiáng)。另外,該文沒有對每種病灶對應(yīng)相關(guān)模態(tài)單獨處理,四種模態(tài)都是單通道圖片數(shù)據(jù),直接將數(shù)據(jù)合并成四通道數(shù)據(jù)送入網(wǎng)絡(luò),這種數(shù)據(jù)處理方式能有效地結(jié)合每種模態(tài)的特性,更利于模型的訓(xùn)練。

圖4 BraTS腦腫瘤數(shù)據(jù)

3.2 訓(xùn)練環(huán)境及模型參數(shù)設(shè)置

該文的實驗環(huán)境是Nvidia Tesla V100 GPU,所用框架為基于Tensorflow的Keras框架。模型參數(shù)設(shè)置：網(wǎng)絡(luò)優(yōu)化器采用的是Adam優(yōu)化器,初始學(xué)習(xí)率LR設(shè)置為1e-3,Batch_Size設(shè)置為4,Epoch設(shè)置為200。根據(jù)不同的實驗?zāi)Ｐ?參數(shù)略有調(diào)整。

該文的損失函數(shù)(Loss Function)采用Sudre[21]提到的Weighted Cross-Entropy (WCE)和Generalized Dice Loss (GDL)相結(jié)合的方式。

(9)

(10)

3.3 評價指標(biāo)

在MRI腦腫瘤分割任務(wù)中,最重要的評價指標(biāo)是Dice相似系數(shù)。Dice相關(guān)系數(shù)有效地衡量了算法分割結(jié)果與真實標(biāo)注標(biāo)簽的重疊程度,具體計算如公式(11)：

(11)

其中，TP、FP、FN分別為真陽性、假陽性、假陰性。Dice相關(guān)系數(shù)越大表示腦腫瘤的分割結(jié)果越好。

3.4 實驗結(jié)果

如表1所示，文中模型MascParallelUnet在參數(shù)量并未增加的條件下，分割性能較原始Unet 模型分割Complete Dice提升了1.56%,Core Dice提升了4.06%,Enhancing Dice提升了9.58%。

表1 對比實驗Dice指標(biāo)對比

分割效果如圖5所示，可以看到文中模型得到的分割結(jié)果對于腦腫瘤病灶的細(xì)節(jié)處理得更好。與現(xiàn)有的模型對比,文中模型比DenseUnet模型的分割精度略低,但是DenseUnet的參數(shù)量較該文網(wǎng)絡(luò)參數(shù)量多了2倍左右。

圖5 對比實驗分割效果圖

該文分別測試了結(jié)合Masc模塊和不結(jié)合Masc模塊對模型分割精度的影響,并與其他注意力模塊進(jìn)行了對比。如表2所示,可以發(fā)現(xiàn)Masc注意力結(jié)構(gòu)加入Unet,Complete Dice沒明顯提升,而Core Dice和Enhancing Dice有顯著提升,分割效果如圖6所示。該文在對比實驗中加入CBAMUnet和AGUnet,實驗結(jié)果表明兩種注意力模塊的Complete Dice和Core Dice都比傳統(tǒng)Unet略低,在Enhancing Dice上有所提高。該文提出的新的注意力模塊結(jié)構(gòu)較Unet的Core Dice提升了2.03%,在Enhancing Dice上提升了8.73%,可以發(fā)現(xiàn)該文提出的Masc模塊對小目標(biāo)病灶的分割更加精確,證明了該文改進(jìn)的注意力模塊的有效性。另外,不加入Masc模塊的并行特征提取網(wǎng)絡(luò)ParallelUnet較Unet模型在Core Dice和Enhancing Dice上也有顯著的提升,說明該文提出的并行特征提取網(wǎng)絡(luò)對總體腦腫瘤分割精度的提升起到了主導(dǎo)作用。

表2 注意力模塊Dice指標(biāo)對比

圖6 注意力模塊性能分割效果圖

4 結(jié)束語

為了提升腦腫瘤多類型病灶分割的分割精度,該文提出了一種基于Unet網(wǎng)絡(luò)的快速、準(zhǔn)確的腦腫瘤自動分割模型。首先,重新設(shè)計了模型編碼網(wǎng)絡(luò)結(jié)構(gòu),引入并行Dalited Convolution提取模塊,融合淺層空間特征與高層語義特征,使特征提取更加充分。其次,模型結(jié)合多注意力機(jī)制并對空間注意力模塊進(jìn)行了特征融合加強(qiáng),使模型更加注意圖像的二維空間信息。最后,對比其他基于Unet改進(jìn)的方法,該文模型的分割精度得到了一定程度的提升。此外,模型是基于2D數(shù)據(jù)集進(jìn)行分割的,并沒有充分利用數(shù)據(jù)集3D的空間信息。因此,下一步研究重點是把2D模型拓展到3D模型。