孫軍梅 葛青青 李秀梅 趙寶奇
(杭州師范大學信息科學與技術學院 杭州 311121)
近年來,深度學習技術已廣泛應用于醫學圖像領域[1–3],其中如何自動識別和分割醫學圖像中的病灶是最受關注的問題之一。由于人體器官多樣、病灶形狀復雜、圖像噪聲干擾等諸多原因,器官病灶等一些待分割物體容易出現邊緣分割不清晰、缺失值大等狀況。
目前已有較多學者針對醫學圖像分割方法展開相關研究,其中UNet[4]是最典型也是應用最廣泛的方法,它利用收縮路徑獲取特征信息,利用擴張路徑實現精確定位,在各類數據集上均有較好的表現。其后出現了許多在UNet上進行改進的網絡模型。Shankaranarayana等人[5]將殘差連接思想和UNet結合提出殘差UNet(Res-UNet)。Oktay等人[6]提出注意力UNet(Attention U-Net),通過集成注意力門捕捉顯著特征。Zhou等人[7]提出UNet++,整合不同層次的特征,使用靈活的網絡結構配合深監督,使得參數量巨大的深度網絡能夠在可接受的精度范圍內大幅度縮減參數量。Jafari等人[8]提出密集殘差UNet(Dense Residual U-net, DRU-Net),在殘差網絡(ResNet)[9]與密集卷積網絡(DenseNet)[10]的基礎上額外增加了跳躍連接,以更少的參數實現了更高的精度。此外,針對目前方法不能很好地分割腫瘤和器官的重疊部分這一問題,Li等人[11]提出了一種基于位置導向的可變形UNet,利用可變形卷積的空間變形能力處理器官和腫瘤的幾何變換,并引入了一個新的池化模塊,來保留傳統最大池化操作中丟失的位置信息。為了自動完善分割網絡的分割效果,Kitrungrotsakul等人[12]提出了一種用于醫學圖像分割的交互式深度優化網絡RefineNet,由分割主干UNet和優化網絡兩部分組成,結合了多尺度初始細分的功能和用戶的種子點。網絡可以在訓練時自行生成種子點,但在測試時需要用戶提供種子點,較為不便,且隨機誤差較大。
醫學圖像通常對比度較低,組織特征具有可變性,不同組織之間或者組織和病灶之間邊緣較模糊,且微細結構(如血管、神經等)分布復雜。圖像邊緣是醫學圖像最基本的特征之一,往往攜帶著一幅圖像的重要信息。因此,邊緣增強處理是醫學圖像處理的關鍵技術之一,對于進行高層次的特征描述、識別和理解等有著重大的影響。針對醫學圖像邊緣增強問題,Zhang等人[13]提出了邊緣注意網絡(Edge aTtention Network, ET-Net),嵌入邊緣注意表示來指導分割網絡。具體地,邊緣引導模塊用于學習早期編碼層中的邊緣注意表示,然后將其傳輸到多尺度解碼層,并使用加權聚合模塊進行融合,對邊緣信息做到了很好的獲取。Valanarasu等人[14]提出了一個完整的級聯網絡KiU-Net進行腦解剖分割,可以同時利用Ki-Net的底層精細邊緣特征圖和U-Net的高級形狀特征圖,不僅提高了分割精度,而且對于小的解剖標志和模糊的噪聲邊緣也實現了快速收斂。針對分割區域及邊緣的不確定性與模糊性,Lee等人[15]提出新穎的邊緣關鍵點選擇算法和結構邊緣保持分割框架,使得模型能夠自動提取和學習結構關鍵點信息,但在網絡中增加了較多參數,不易于模型部署。針對區域邊緣和區域內的不連續性,Chu等人[16]提出一種利用簡單邊緣檢測器定位所有不連續點,并對這些區域進行額外監控的方法,有效提高了分割精度。
上述工作都為后續醫學圖像分割技術的研究提供借鑒意義,但仍有許多問題亟待解決。首先,分割精度較低。細胞和器官等形狀復雜多樣,且可能存在團簇重疊狀態,致使錯誤分割的概率很大,影響醫學判斷。其次,模型結果的提升伴隨著測試時參數的增加。越來越多的網絡模型基于UNet做出改進[17–20],或采用雙流卷積神經網絡結構增強邊緣信息[21–23],雖然在一定程度上可有效提升分割結果,但存在測試時參數冗余現象。
為了解決上述問題,本文基于UNet提出一種具有邊緣增強特點的醫學圖像分割網絡(Add-and-Subtract UNet, AS-UNet),Add是指訓練時在UNet的基礎之上增加邊緣注意模塊(Boundary Attention Block, BAB), Subtract是指測試時舍棄BAB。AS-UNet通過掩膜邊緣提取算法得到掩膜邊緣圖,在UNet擴張路徑的最后3層引入結合多尺度特征圖的BAB,強化邊緣,減少缺失值;并提出使用組合損失函數,結合基于區域的Dice損失(Dice Loss)[24]和基于邊緣的Boundary損失 (Boundary Loss)[25],在保證區域缺失值小的同時補充邊緣信息,提高分割精度;此外,由于組合損失函數包含了經過BAB前后的兩種輸出,在訓練時經過前后向反饋不斷更新AS-UNet中的網絡參數,使得訓練好的模型在測試時可以舍棄添加的BAB部分,保證分割精度的同時減少測試時參數。
本文所提具有邊緣增強特點的醫學圖像分割網絡AS-UNet結構如圖1所示。采用UNet網絡的編碼—解碼結構,輸入通過UNet得到的直接輸出稱為輸出1(output1)。將UNet擴張路徑中L7, L8, L9得到的特征圖R7,R8,R9分別作為輸入層送入BAB。R7經過BAB后采用2× 2的上采樣(Up-conv 2× 2)達到與R8相同的尺寸,作為補充層傳入L8的BAB中。同理,R8經過BAB后采用2× 2的上采樣達到與R9相同的尺寸,并作為補充層傳入L9的BAB中。而最后一層得到的特征圖R9經過BAB后則直接進行1× 1的卷積得到輸出2(output2)。此過程如式(1)所示

圖1 AS-UNet網絡結構

訓練時輸入圖像通過UNet生成輸出1,同時將R7, R8, R9通過BAB還原圖像邊緣細節信息,生成輸出2。將輸出1和輸出2分別與掩膜(Mask,即標簽(Ground Truth, GT))進行損失計算,在前后向傳播中更新參數,得到最終模型。
在UNet中,邊緣分割模糊以及部分區域缺失仍是最主要的不足。因此,本文所提AS-UNet網絡通過在UNet網絡中引入邊緣注意模塊BAB以強化邊緣信息,補充缺失區域。BAB結構如圖2所示。

圖2 BAB結構

在BAB中,改變掩膜邊緣圖大小并按通道連接至各特征圖之后,有效地提供了精確的圖形邊緣,增強圖像對于目標區域的邊緣信息獲取,使得后續網絡更重視這一部分的信息,從而彌補UNet對邊緣信息的忽視。
2.2.1 掩膜邊緣提取算法
為了引入像素級別的精確邊緣信息,本文通過掩膜邊緣提取算法將訓練圖像的掩膜導出為掩膜邊緣圖,以此作為邊緣信息的重要補充。掩膜邊緣提取算法過程可表述為:遍歷掩膜的每一個像素點(i,j),當遍歷的像素點值為0,且以該像素點為中心的九宮格內其余像素點不都為0時,將該像素點記作0,直到掩膜的所有像素點都遍歷結束后,生成掩膜邊緣圖。圖3為掩膜與經過掩膜邊緣提取算法導出的掩膜邊緣圖。

圖3 圖像掩模與對應邊緣圖
2.2.2 注意力模塊
受空間通道壓縮與激勵(spatial-channel Sequeeze & Excitation, scSE)[26]注意力模塊的啟發,本文提出一種新的注意力模塊,具體結構如圖4所示。對于輸入特征圖U ∈?w×h×c,分別在通道和空間上進行壓縮得到特征圖U?sCE∈?w×h×1和向量

圖4 注意力模塊

cSE和sCE順著空間和通道維度分別進行特征壓縮,將得到的特征權重向量分別加權到輸入特征圖上,完成不同維度上對原始特征的重標定。而scSE的工作僅僅是將cSE和sCE的輸出結果逐像素相加,同時完成空間以及通道上的壓縮與激勵。而本文所提出的注意力模塊首先將空間和通道上壓縮得到的特征圖進行相乘,得到與輸入相同大小的權重W,如此便可以為輸入特征圖的每個像素都提供各自的權重,且相乘能夠讓重要的位置更加突出,而價值較小的位置則可以被抑制。隨后再將這一權重與輸入逐像素相乘得到最終輸出。此方法相較于scSE同樣簡單且不增加任何參數,但能夠更多地關注重要位置的信息,獲取感興趣區域及邊緣特征信息。本文將其應用于邊緣注意模塊3×3的卷積層之后,幫助模型取得更好的分割效果。
在圖像分割領域,主要采用的損失函數包括基于分布的交叉熵損失函數(Cross-Entropy Loss)以及基于區域的Dice損失函數。交叉熵損失函數單獨評估每個像素矢量的類預測,然后對所有像素求平均值,使得圖像中的像素能夠被平等學習。但是,醫學圖像中常出現類別不均衡的問題,由此導致訓練會被像素較多的類主導,對于像素較少的物體很難學習到其特征,從而降低網絡的有效性。而Dice損失函數,本質上是衡量兩個樣本的重疊部分,雖然解決了類別不均衡的情況,但是未關注到圖像的邊緣信息。由于醫學圖像的邊緣信息尤為重要,因此本文提出結合基于區域的Dice 損失和基于邊緣的Boundary 損失的組合損失函數L,在兩個不同的側重維度上進行監督。Dice損失和Boundary損失的定義如式(4)、式(5)所示

損失函數L結合了基于區域的Dice損失和基于邊緣的Boundary損失,使得網絡同時關注區域與邊緣信息,在保證區域缺失值小的同時補充邊緣信息,從而提高分割精度。隨著神經網絡的不斷迭代,平衡系數α和β自學習調整更新,促使在前半部分的UNet網絡中,Dice損失占據比重較大,相對來說更加關注區域的信息,在后半部分的邊緣注意模塊中,Boundary損失占據比重較大,相對來說更加關注邊緣信息。本文使用組合損失函數發揮邊緣注意模塊的作用,達到關注區域信息的同時不損失邊緣信息,解決目前醫學圖像分割中存在的分割缺失值大、邊緣不清晰的問題。
BAB的引入提高了模型的分割精度,但也不可避免地增加了網絡參數,本文方法AS-UNet在測試時利用神經網絡結構的冗余性進行模型精簡,只需要選擇模型的輸出1,達到不增加網絡參數的目的。
這種方法可執行的依據是使用新的組合損失函數L,L中的Dice損失和Boundary損失都結合了y1和y2兩種輸出,因此,在將L總體最小化時必然會將針對y1部分的Loss最小化,如式(9)、式(10)所示。訓練時經過前后向反饋不斷更新AS-UNet中的網絡參數,修復未留意的邊緣信息。因此可實現在最終測試時即使舍棄邊緣注意網絡,僅采用y1作為測試結果,也不影響分割效果,在提高分割精度的同時極大地減少了模型參數,為實際應用提供了更大的可能性。

結合兩種輸出的損失函數L可以兼顧輸出1與輸出2,使其分別達到最好的效果。訓練時輸入需要經過基礎網絡+BAB,不斷更新優化網絡內部參數,而測試時可以舍棄BAB,在訓練好的模型基礎上選擇輸出1即可得到最終結果,在提高分割精度的同時可以減少參數量。
數據集1:Glas。Glas是MICCAI2015腺體分割挑戰賽的公開數據集,由165幅圖像組成,這些圖像來自16個蘇木精和伊紅(Hematoxylin-Eosin,H&E)染色載玻片的大腸癌組織切片。原始圖像大小不一,大多數為775×522。將數據集分為訓練集和測試集,100幅圖像屬于訓練集,65幅圖像屬于測試集。
數據集2:DRIVE。DRIVE發布于2003年,是一個用于血管分割的數字視網膜圖像數據集,由40幅圖像組成,其中7張顯示出輕度早期糖尿病視網膜病變跡象。將數據集分為訓練集和測試集,30幅圖像屬于訓練集,10幅圖像屬于測試集。
數據集3:ISIC2018。ISIC2018[27,28]是用于黑色素瘤檢測的皮膚病變分析數據集,其中用于醫學圖像分割的部分包括2594幅RGB皮膚病變圖像,采用二進制標注。將數據集分為訓練集和測試集,2074幅圖像屬于訓練集,520幅圖像屬于測試集。
本文使用平均Dice(Mean Dice)[29]、F1值、Hausdorff距離[30]作為評價指標,衡量分割結果的優劣,更為公正地對分割結果做出評價。其中Mean Dice更多關注區域信息,用來衡量兩組樣本之間的相似性;F1值是精確率(Precision, P)和召回率(Recall, R)的加權調和平均,用來衡量單個樣本目標的檢測精度;Hausdorff距離更多關注邊緣信息,用來衡量基于邊緣的分割精度。3類評價指標計算方式如式(12)—式(14)所示。

實驗使用的深度學習框架是Tensorflow 1.14,計算機操作系統為64位Ubuntu 20.04.1,同時使用GPU加速網絡模型的訓練和測試。采用Adam優化器,因為 Adam 算法可以在訓練時自適應地調節學習率,且有更快的收斂速度,初始學習率設置為0.001。考慮到計算資源的局限性,訓練階段批大小(batch size)設置為2,測試階段批大小設置為1,輸入圖像統一轉換為512×512,并進行增強對比度處理,灰度化彩色圖像后將灰度像素轉化為0~1,減小輸入特征的尺度,再送入網絡模型進行訓練和測試。
(1)不同模型對比實驗。為驗證本文方法的有效性,在3個不同類型的數據集上,分別和UNet及其他3個也在UNet基礎上做改進的較新穎的圖像分割方法UNet++, DRU-Net以及KiU-Net進行對比實驗。其中UNet是醫學圖像分割的基礎網絡;UNet++在不同層之間重新設計跳躍連接將編碼器和解碼器相連,并且通過減輕未知的網絡深度和設計一個修剪方案來提高模型性能;DRU-Net提出了一個簡單但高效的網絡模塊,以更少的參數實現更高的精度;KiU-Net則基于UNet進行模塊結構的改進,通過改變感受野捕捉精細邊緣特征圖和高級形狀特征圖,更好地實現邊緣細節捕獲與分割。
各模型的實驗結果如表1所示。相比于其他4種網絡,本文方法在不同數據集下3個評價指標均有較好的提升。與UNet相比,本文方法幾乎不增加測試時參數量,在Glas數據集上Mean Dice和F1值分別提高了2.19%和2.21%,Hausdorff距離減小了31.8;在DRIVE數據集上Mean Dice和F1值分別提高了2.16%和2.64%,Hausdorff距離減小了10.89;在ISIC2018數據集上Mean Dice和F1值分別提高了1.53%和2.91%,Hausdorff距離減小了7.53。由此可見本文方法的有效性和較好的泛化能力。

表1 不同模型在不同數據集上分割結果對比

圖5 遷移模型
圖6—圖8給出了5種網絡模型在3個數據集上的分割效果。從中可以看出本文所提AS-UNet在分割結果中更加貼近原圖的GT,尤其是在邊緣的處理上更加細致。在Glas數據集中,細胞較大,細胞之間的粘合問題較少,本文方法處理得到的分割圖存在的缺失值較少,且邊緣較平整;在DRIVE數據集中,人眼血管較細,分布復雜,分割難度較大,本文方法處理得到的分割圖相較于其他模型在細血管的連續性問題上處理得更好,其他模型分割得到的細血管很多都變成了點狀而非連續,本文模型借助邊緣注意模塊在一定程度上仍能保持血管的基本連續;在ISIC2018數據集中,病變區域形狀各異,各樣本間差異較大,其他模型均存在明顯的缺失或多余塊,而本文方法較為精細地分割邊界,且區域完整。

圖6 不同模型在Glas數據集上的分割結果對比

圖7 不同模型在DRIVE數據集上的分割結果對比

圖8 不同模型在ISIC2018數據集上的分割結果對比
(2)消融實驗。采用消融實驗分別驗證本文所提邊緣注意模塊BAB以及損失函數的有效性、可移植性和參數情況。所有實驗訓練配置均相同。參照圖5,基礎網絡可用UNet和FCN等網絡代替,因此分別以UNet和FCN為基礎網絡設計實驗。其中,測試時,輸出1可作為舍棄BAB后的直接輸出,也可用于生成掩膜邊緣圖,繼而送入BAB生成另一個輸出2。將訓練時添加BAB且測試時采用輸出2的方法記作+BAB;將訓練時添加BAB且測試時舍棄BAB采用輸出1的方法記作+Sub。
如表2所示,各種對比方法在兩種基礎網絡之上的效果相似。添加BAB后,通過引入掩膜邊緣圖解決分割邊緣不清晰問題,并對網絡關鍵特征的提取有一定的引導意義。相較于單純的UNet和FCN,在各數據集上Mean Dice提升1%~4%,F1值提升1%~3%,Hausdorff距離也基本上有所減小,可見本文設計的BAB具有可移植性與非依賴性;對比+BAB與+Sub,兩者在訓練時都融合了BAB,但在測試時,前者使用輸出2,后者使用輸出1,它們的Mean Dice, F1值和Hausdorff距離并無較大改變,基本保持在同一水平線,但Sub的參數量減少了1 M左右,在保證分割精度的同時減少測試時參數。

表2 消融實驗
(3)不同注意力模塊對比實驗。為驗證本文提出注意力模塊的有效性,分別將BAB中的注意力模塊舍棄或替換為scSE進行對比實驗。輸出采用輸出1,實驗結果如表3所示。對比本文方法與不添加注意力機制、添加scSE的分割結果,本文方法的提升效果最為明顯。相較于scSE,在各數據集上Mean Dice提升0.3%左右,F1值提升1%左右,Hausdorff距離也都有所減小。可見本文設計注意力模塊的有效性,可強化學習重點特征,彌補缺失值,融入在BAB中,進一步提高網絡分割精度。

表3 不同注意力模塊在不同數據集上的分割結果對比
(4)不同損失函數對比實驗。為了驗證本文提出組合損失函數的有效性,在AS-UNet網絡中分別對基于區域的Dice損失、基于邊緣的Boundary損失和本文所提Dice + Boundary 損失這3者進行對比實驗。實驗采用Mean Dice, F1值和Hausdorff距離作為評價指標。3個數據集上的實驗結果如表4所示。從表4中結果可見,組合損失函數在各數據集上都能達到更好的分割效果,相較于單一的Dice損失和Boundary損失,Dice + Boundary損失的Mean Dice和F1值均提高了1%左右,Hausdorff距離均減小2左右。這是因為本文方法引入掩膜邊緣圖,增強邊緣信息,結合區域以及邊緣的損失函數能夠更真實地反映預測圖與GT之間的差異,從而進行參數校準與更新。

表4 不同損失函數在不同數據集上的分割結果對比
本文針對醫學圖像分割缺失值大、邊緣不清晰、模型參數冗余等問題提出AS-UNet。訓練時引入邊緣注意模塊,將掩膜邊緣圖連接到UNet擴張路徑的最后3層上以補充邊緣信息,解決邊緣分割不清晰的問題;同時在邊緣注意模塊中引入新的注意力模塊,并結合多尺度特征圖以強化重要特征學習,彌補部分缺失值;提出了新的組合損失函數優化網絡參數,并實現測試時只需要選擇訓練好模型的輸出1即可得到最終結果,在提高分割精度的同時減少測試時參數。實驗結果表明,本文模型在Glas, DRIVE, ISIC2018 3個不同類型的分割數據集上都有較好表現,驗證了本文方法的有效性和泛化能力。論文相關代碼地址:https://github.com/gqq1210/AS-UNet。