基于多尺度高階注意力機制的視網膜血管分割①

姜璐璐,李思聰,曹加旺,孫司琦,馮 瑞,,鄒海東,4

1(復旦大學 工程與應用技術研究院,上海 200433)

2(上海交通大學附屬第一人民醫院,上海 200080)

3(復旦大學 計算機科學技術學院,上海 200433)

4(蘇州市產業技術研究院,蘇州 215011)

1 引言

血管系統是視網膜的基本結構,其形態學和拓撲結構的變化可以用來識別和分類系統性代謝和血液疾病的嚴重程度,例如糖尿病和高血壓[1].糖尿病性視網膜病變 (DR)是糖尿病的一種常見并發癥,是由視網膜微血管滲漏和阻塞導致的一系列眼底病變.DR 可引起新血管的生長,是否有異常新生血管也是判斷增殖性DR 與非增殖性DR 的標準[2].高血壓視網膜病變 (HR)是另一種常見的由高血壓引起的視網膜疾病[3].在高血壓患者中,可以觀察到血管彎曲度增加或血管狹窄[4].通過視網膜血管獲得的血管形狀和分叉的信息,可以增強對DR 或者HR 的監測.因此,分割視網膜血管對于一些嚴重疾病的早期診斷與治療具有重要意義.

現有的眼底視網膜成像技術有以下幾類: 彩色眼底照相 (FP)技術、眼底熒光素血管造影 (FFA)、光學相干斷層掃描 (OCT)以及眼底相干光層析血管成像(OCTA).彩色眼底照相是最常用的視網膜成像技術,其優點是獲取方式簡單、圖像易于觀察.

傳統的無監督方法一般包括: 濾波匹配法、區域生長、血管跟蹤、閾值分割和圖像形態學處理等.這些傳統的無監督方法不需要人工標注,但這些方法依賴于手工提取特征進行血管表示與分割.此外,此類算法存在分割精度不夠、泛化性較差等局限性.

與傳統的無監督方法相比,深度卷積神經網絡方法具有更強大的特征表征和學習能力,在醫學圖像分割任務中取得了最高水平[5].自2015年引入U-Net[6]以來,它已成為醫學影像分割中最具影響力的深度學習框架[7-10].其整體網絡采用編碼器-解碼器的結構,通過“跳躍連接”將不同分辨率的特征圖進行通道融合產生較好的分割效果.盡管U-Net 具有良好的表示能力,但它依賴于多級級聯卷積神經網絡.這種方法在重復提取低層特征時會導致計算資源的過度和冗余使用[11].

注意力機制被提出用于解決以上問題,其模仿了人類視覺所特有的大腦信號處理機制,令網絡從大量信息中重點關注對任務結果更重要的區域,而抑制其他不重要的部分[12].在視網膜血管分割任務中,背景像素占比較大,而血管像素的占比小,因此可以采用注意力機制關注血管區域.卷積神經網絡可以利用不同類型的注意力機制以關注重要的區域或者特征通道[13-18].例如,空間注意力機制[11,18]利用特征的空間關系生成空間注意力圖從而使網絡關注具有豐富信息的區域,通道注意力機制[13]通過顯式建模通道間的依賴關系來提高模型的性能.空間注意力和通道注意力的融合[15]也已成功地應用于醫學分割領域.

然而,這些常用的方法是一階注意力機制,難以提取圖像中一些更為抽象的高階語義信息且不能充分利用到全圖像的信息,導致在處理形狀和結構復雜的目標時發生退化[19].尤其在視網膜血管分割任務中,由于血管形態結構多變,以上方法仍欠缺對復雜和高階特征信息的捕獲能力.

本文提出了一種基于多尺度高階注意力機制的視網膜圖像分割方法(MHA-Net),可以明顯提高視網膜血管的分割精度.該方法采用改進的U-Net 結構,并引入多尺度高階注意力模塊,對編碼器提取到的深層特征進一步處理,聚焦于圖像的高階語義信息,從而改進模型處理醫學圖像分割時尺度不變的缺陷.經過在DRIVE[20]數據集上的實驗證明,該方法有效地提高了分割的精度,同時對細小血管的分割也更為精細.

2 相關工作

2.1 空洞卷積在圖像分割中的應用

空洞卷積(dilated convolution)[21,22]通過在卷積核相鄰兩個元素之間插入零值,在不增加參數量和計算成本的同時擴大了感受野.受空洞空間金字塔池化(ASPP)[23]在語義圖像分割中的應用啟發,空洞卷積在醫學圖像分割中同樣得到了廣泛的應用[17,24].但是,基于空洞卷積的分割方法都存在一個共同問題,稀疏采樣會造成詳細信息的丟失,從而導致像素級分類不準確.D-LinkNet[25]利用“短路連接(shortcut)”結合了文獻[21]的級聯模型與文獻[1]的并行模型.

之前的研究主要集中在通過增加在不同尺度特征圖上的感受野,從而直接提高分割網絡的性能.我們的工作與上述方法不同,我們利用空洞卷積對不同尺度的特征圖進行采樣,并通過聚合這些多尺度的特征圖產生高階注意力圖,從而進一步使網絡聚焦于更加抽象和全面的語義信息.

2.2 高階注意力

注意力機制的思想核心是通過計算權重矩陣而使網絡有選擇地關注具有重要信息的部分[12].Okty 等人[11]提出了用于醫療影像分割的注意力門控(attention gate,AG)模型,該模型可以自動學習區分目標的外形和尺寸,在小目標分割任務中效果尤其顯著.不同于在跳躍連接中添加注意力門控(AG)的方法,SA-UNet[14]引入了一個空間注意力模塊,通過在空間維度計算注意力權重矩陣并與輸入的特征圖相乘,實現自適應地細化特征.該方法是注意力模塊在U 形分割網絡降采樣后的深層特征圖上的一種應用.Chen 等[19]首先提出了高階注意力模型,并將其應用于行人重識別建模.該模型利用注意機制中形成的復雜高階統計量,捕捉行人之間的細微差異,從而產生區別性的關注建議.Ding 等[26]利用圖的傳遞閉包進一步優化高階注意力模塊,在此基礎上提出具有自適應感受野和動態權重的high-order attention (HA)模塊.HA 模塊通過圖的傳遞閉包構建注意力圖,從而捕獲高階的上下文相關信息.

之前的一些工作(如文獻[13])通過在U 型網絡的底部引入注意力機制來進一步挖掘深層次的特征.然而,這些網絡更多地關注了局部信息,而忽略深層特征中的全局信息.這導致盡管在提取深層特征時添加了幾種不同類型的注意力模塊,也不能有效地提高醫學圖像分割任務的性能.相反,模型的性能甚至會略有下降.

本文的工作是在上述注意力機制[14,19,26]上的改進.在U 形網絡的多個降采樣塊之后所得的深層特征的噪聲相對較小,因此注意力模塊需要盡可能地挖掘深層特征中的全局信息.另一方面,與淺層特征相比,在深層特征中引入噪聲會對整個模型造成更大的損害.因此,本文設計了多尺度高階注意力(MHA)模塊,其在不引入噪聲的前提下引導網絡提取深層特征中的更為全局的信息,有效提高了視網膜血管中分割性能.

3 基于多尺度高階注意力機制的分割網絡

3.1 網絡框架

圖1 給出了基于多尺度高階注意力機制的視網膜圖像分割方法(MHA-Net)的網絡架構,其遵循了編碼器-解碼器的U 型結構.編碼器包含若干個下采樣塊和MHA 模塊,其中每個下采樣塊由1 個3×3 的卷積層、1 個批處理規范化層和一個ReLU 激活函數層組成,3 個下采樣塊連接在一起后緊跟一個2×2 的最大池化操作.在下采樣完成之后,將提取到的圖像深層次特征輸入到MHA 模塊進行細化,MHA 模塊的位置放置于網絡底部,即U 型收縮路徑和擴張路徑之間.在此處加入attention 模塊的原因是在靠前位置采集到的為低層次結構信息,包含有許多噪聲.此外,加權的shortcut 被引入以保留原本的上下文信息.最后,經過融合得到的特征圖通過編碼器產生最終的分割結果.解碼器部分使用反卷積[27]進行上采樣操作.

圖1 MHA-Net 架構圖

3.2 多尺度高階注意力模塊

本文提出的多尺度高階注意力模塊如圖2 所示.在編碼器的底部,原始的特征圖Xin∈RH×W×C通過并行的共享權重的空洞卷積(膨脹率r分別為1,2,4,8),產生新的多尺度特征圖分為為Xr(r=1,2,4,8),通過1×1 卷積得到的特征圖為X*.將這些多尺度特征圖使用式(1)計算得到融合的多尺度注意力矩陣:

圖2 多尺度高階注意力(MHA)模塊

其中,1/C是用來控制數值爆炸的縮放因子.之后,利用圖的傳遞閉包計算了多尺度高階注意力矩陣A,m∈{1,2,···,n}.具體計算的細節將在第3.3 節討論.最后,將特征圖X*與歸一化的高階注意力矩陣相乘得到細化的特征圖Xm,如式(2):

Γθ代表1×1 卷積.在多尺度高階注意力模塊之后,將細化后的特征圖Xm乘上自適應因子 α以抵消縮放因子1/C的偏移影響,如式(3):

深層特征圖在通過多尺度高階注意力模塊之后,提取了更加高階抽象的語義特征,也更具有區分力,從而更聚焦于血管的分割.之后,再通過解碼器模塊,逐漸從低分辨率重構至高分辨率.

3.3 高階注意力傳播機制

根據文獻[26],最初的多尺度注意力融合矩陣可以看做圖的鄰接矩陣,圖中的邊表示連接的兩個節點屬于同一類.如圖3 所示,給定注意力圖,通過閾值化刪去低置信度的邊后形成下采樣的圖如式(4):

圖3 三階高階注意力傳播原理圖: 以黃色點為中心點通過圖的傳遞閉包進行傳播

其中,δ代表閾值,設置為0.5.如圖4 所示,根據圖的傳遞閉包,可以通過鄰接矩陣自乘m-1 次得到:

其中,m表示鄰接矩陣冪次的整數,代表注意力傳播的階數.因此,不同層次的注意力信息通過解耦成不同的注意圖并得到高度相關的鄰居.生成的高階注意圖用于聚合多層次的上下文信息.

4 實驗驗證與結果分析

4.1 數據集與實驗參數

本文使用的數據集是DRIVE (digital retinal image for vessel extraction)[20].該數據集包含40 張圖像像素尺寸為584×565 的彩色眼底圖像,其中訓練集與測試集各20 張.為擴充數據,避免訓練樣本過少可能造成的過擬合問題,我們對訓練樣本隨機采樣256×256 的patch.此外,使用隨機翻轉、隨機旋轉、彈性形變等方法進行數據增強.本文使用PyTorch 框架[28],批量設置為 16,采用Adam 算法[29]優化模型,學習率設置為0.000 1.動量和權重衰減因子分別設置為0.9 和0.999.

4.2 評價指標

為了對實驗結果進行客觀的定量分析,選取以下指標進行計算: Dice 系數(DSC)、準確率(ACC)、敏感度(SE)、特異性(SP)和ROC 曲線下面積AUC.AUC 的范圍在0-1 之間,AUC 越逼近 1,其模型預測能力越高.評價指標的計算方式如下:

其中,X代表金標準,Y代表預測結果.真陽性TP為正確分類的血管像素個數,真陰性TN正確分類的背景點像素個數,假陽性FP為背景像素誤分成血管像素的個數,假陰性FN為血管像素誤分成背景像素的個數.

4.3 實驗結果分析

本文算法性能在DRIVE 數據集上評估,圖4 展示了部分分割結果.圖4(a)為原始圖像,圖4(b)為金標準圖像,圖4(c)為本文算法的分割結果,從結果可以看出,本文算法整體分割效果良好,平滑度也優于金標準.同時,本文算法細節上表現優秀,保持了微血管的連通性,說明本文中采取的注意力機制能夠關注到重要的血管區域.

圖4 DRIVE 數據集分割結果

為了驗證本文所提出的模型性能的優越性,表1將本文算法與近兩年最先進的血管分割算法的各項指標進行對比,其中加粗字體部分為每項最優指標.

表1 DRIVE 數據集上不同算法分割性能比較

結果表明,本文提出的多尺度高階注意力方法MHA-Net 取得了優異的表現,其Dice 系數、靈敏度和AUC 分別達到了0.826 6、0.831 2 和0.988 3,在所有方法中表現最優.本文算法在保證高準確率的同時,有著良好的敏感度,這意味著分割結果盡可能地保留血管信息,分割得到的血管連續完整.綜上所述,本文算法整體性能優于現有算法.

為了證明提出的多尺度高階注意力(MHA)模塊的有效性,在DRIVE 數據集上還進行了消融實驗.表2展示了U-Net、U-Net+MHA、Backbone、Backbone+HA 以及MHA-Net 的分割性能.其中U-Net+MHA 表示在U-Net 基礎上引入MHA 模塊的網絡,Backbone表示在本文使用的骨干網絡,Backbone+HA 表示與本文相同的骨干網絡上引入原始的高階注意力HA 模塊,MHA-Net 為本文算法,相當于Backbone+MHA,即在本文使用的骨干網絡上引入多尺度高階注意力MHA 模塊.

表2 DRIVE 數據集上的消融實驗

結果表明: (1)U-Net+MHA 比U-Net 有更好的性能,準確率提高0.07%,敏感度提高1.56%,AUC 提高0.20%,這證明了本文提出的多尺度高階注意(MHA)模塊的有效性.(2)MHA-Net 在準確率、靈敏度和AUC 指標上都優于Backbone+HA,這表明多尺度高階注意力模塊對多尺度上下文特征信息捕捉能力更強,對復雜結構的血管圖像有更強的特征提取能力.(3)本文提出的MHA-Net 在大多數指標上都表現最好,在視網膜血管分割領域全面優于U-Net,說明該網絡模型的合理性和優越性.

此外,對實驗結果進行了可視化分析,如圖5 所示,從左至右依次是原始圖、金標準、U-Net 分割結果、Backbone+HA 分割結果以及MHA-Net (本文)分割結果.我們放大了微血管,本文提出的MHA-Net 分割結果更加精細,在血管末也未出現粘連或者斷裂的情況.

圖5 DRIVE 數據集分割結果對比

5 結論與展望

本文針對視網膜血管分割任務中血管粗細不勻、形狀多變、微小血管易斷裂等問題,本文提出多尺度高階注意力(MHA)機制以自適應地挖掘深層次特征.MHA-Net 以端到端方式進行視網膜血管分割訓練,并通過MHA 模塊學習到具有鑒別性的特征.在DRIVE上的實驗表明,本文提出的算法取得了優越的分割性能.同時,MHA 模塊可以即插即用,在各種醫學影像分割任務中適用.后續的工作將嘗試把多尺度高階注意力機制運用到三維的影像分割中.