基于注意力機制的U-Net腦脊液細胞分割

代 茵, 劉維賓, 董昕陽, 宋雨朦

(1.東北大學(xué) 醫(yī)學(xué)與生物信息工程學(xué)院, 遼寧 沈陽 110169; 2.東北大學(xué) 教育部醫(yī)學(xué)影像與智能分析工程中心, 遼寧 沈陽 110169; 3.約克大學(xué) 計算機學(xué)院, 英國 約克郡 YO10 5DD)

目前國內(nèi)外針對自身免疫性腦炎的診斷研究是利用磁共振成像(magnetic resonance imaging, MRI)檢測以及結(jié)合其他影像學(xué)診斷方法，如正電子發(fā)射計算機斷層顯像(positron emission computed tomography, PET)、超聲波檢查，以及腦脊液(cerebrospinal fluid, CSF)檢測，而CSF檢測是最為標準的，通常作為醫(yī)生判斷疾病的“金標準”[1].免疫機制在精神疾病病理學(xué)中的作用越來越得到人們的關(guān)注，多項研究表明了自身免疫性疾病與精神分裂癥之間存在著必然的聯(lián)系.CSF分析對識別中樞神經(jīng)系統(tǒng)是否存在炎癥起著核心的作用[2].CSF檢測結(jié)果不僅可以作為判斷腦炎與精神分裂癥的依據(jù)，還可以作為臨床孤立綜合征[3]、額顳癡呆、小兒鞍上生殖細胞瘤、多發(fā)腦梗死性癡呆、神經(jīng)萊姆病、遺傳性多發(fā)腦梗死性癡呆、單發(fā)腦梗死性癡呆等疾病的重要依據(jù).

CSF細胞圖像有其獨有的特點，其存在細胞邊界模糊與背景不易區(qū)分、數(shù)據(jù)圖像存在一定的噪聲、相同類型的細胞通常在圖像中成簇存在[4]、樣本數(shù)據(jù)過少等問題.目前醫(yī)生大多還是通過手動分割CSF細胞的方式并通過經(jīng)驗來診斷相關(guān)疾病，工作強度大，效率低，而基于人工智能的傳統(tǒng)分割算法的魯棒性差，所以針對CSF細胞檢測的自動分割研究顯得尤為重要.近幾年深度學(xué)習(xí)發(fā)展迅猛，有研究人員在圖像分割與深度學(xué)習(xí)相結(jié)合的領(lǐng)域作了許多研究與貢獻.在CSF圖像的研究中，Xu等[4]提出一種層次非最大抑制(hierarchy-non maximum suppression,Hierarchy-NMS)算法，與非極大抑制算法和軟化非極大抑制算法進行了比較，證明了層次非最大抑制算法在細胞識別和計數(shù)方面均優(yōu)于非極大抑制的一些變體.但該研究更多的關(guān)注點是識別與記數(shù)，而沒有考慮CSF細胞邊界與背景精確分割的問題.

在醫(yī)學(xué)圖像領(lǐng)域內(nèi)，Jiang等[5]設(shè)計了一種基于雙通道雙向LSTM(dual channel bidirectional LSTM, DC-BLSTM)的神經(jīng)網(wǎng)絡(luò)，將3D圖像分割任務(wù)轉(zhuǎn)化為多個1D分割任務(wù)，比現(xiàn)有的基于卷積神經(jīng)網(wǎng)絡(luò)的3D神經(jīng)元圖像分割方法更容易標記訓(xùn)練樣本.另外極具代表性的是以U-Net模型為基礎(chǔ)框架，通過改造編碼器或解碼器實現(xiàn)其性能的提升.研究人員通過這些變體對許多細胞圖像進行了分割研究，例如對腺體細胞進行分割[6]，肌肉干細胞的分割[7]，肺、細胞輪廓、胰腺的分割研究[8]，多器官分割和心臟分割[9]，以及對ISBI細胞追蹤挑戰(zhàn)賽2019數(shù)據(jù)集的分割測試[10]等.注意力機制的優(yōu)秀性能與良好表現(xiàn)引起了研究人員的注意.Vaswani等[11]提出了Transformer網(wǎng)絡(luò)架構(gòu)，此架構(gòu)是完全基于注意力機制，并且在WMT 2014 English-German數(shù)據(jù)集上驗證了優(yōu)越性能.Zhang等[12]提出一種帶有注意力模塊的O形網(wǎng)絡(luò)結(jié)構(gòu)，可以在不分割的情況下檢測生物醫(yī)學(xué)圖像中的節(jié)點，使用該方法在兩個視網(wǎng)膜數(shù)據(jù)集和一個神經(jīng)元數(shù)據(jù)集中取得了良好的檢測效果.遷移學(xué)習(xí)可以提升模型泛化能力，使模型的初始性能得到提高.劉珍梅[13]將遷移學(xué)習(xí)算法與宮頸細胞核分割算法相結(jié)合，減少分割標注數(shù)據(jù)集制作的壓力，得到的模型收斂性更好.

本文將注意力機制與深度學(xué)習(xí)相結(jié)合，并引入遷移學(xué)習(xí)，對CSF細胞進行精確分割，取得了理想的結(jié)果.

1 方法描述

本文采用了基于注意力機制的U-Net卷積神經(jīng)網(wǎng)絡(luò)對CSF細胞進行精確分割，并引入預(yù)訓(xùn)練模型進行遷移學(xué)習(xí)[14]，在一定程度上提高了模型訓(xùn)練速度與性能.概括來說實驗流程共分為4個部分：

1) 數(shù)據(jù)集擴充：將從已有的CSF圖像中篩查出29例可靠圖像，然后再采用對其旋轉(zhuǎn)、裁剪、鏡像的方法進行擴充；

2) 數(shù)據(jù)預(yù)處理：將擴充后的數(shù)據(jù)集進行尺寸大小統(tǒng)一預(yù)處理，使其適配batch的讀取方式，網(wǎng)絡(luò)的感受野也就限定了統(tǒng)一的尺度范圍，方便統(tǒng)一做訓(xùn)練，并且訓(xùn)練集通過拉普拉斯算子進行了濾波處理；

3) 網(wǎng)絡(luò)訓(xùn)練：將處理好的數(shù)據(jù)集傳入帶有注意力機制的U-Net卷積神經(jīng)網(wǎng)絡(luò)模型中，并引入遷移學(xué)習(xí)進行訓(xùn)練；

4) 分割指標評估：訓(xùn)練完成后，使用訓(xùn)練模型對數(shù)據(jù)進行預(yù)測，使其預(yù)測圖與Ground Truth進行指標評估，評估指標包括相似系數(shù)值、平均交并比和類別平均像素準確率.

實驗流程示意圖如圖1所示.

圖1 實驗流程圖

2 基于注意力機制的U-Net網(wǎng)絡(luò) 結(jié)構(gòu)模型

2.1 總體設(shè)計

本文所使用的基于注意力機制的U-Net卷積神經(jīng)網(wǎng)絡(luò)總體分為編碼器和解碼器，在解碼器中加入注意力模塊，在對編碼器每個分辨率上的特征與解碼器中對應(yīng)特征進行拼接之前，使用Attention Gate重新調(diào)整了編碼器的輸出特征，該模塊生成一個門控信號，用來控制不同空間位置處特征的重要性[15].網(wǎng)絡(luò)的總體結(jié)構(gòu)如圖2所示.

圖2 基于注意力機制的U-Net

2.2 U-Net網(wǎng)絡(luò)結(jié)構(gòu)

2.2.1 編碼器結(jié)構(gòu)

作為主干特征提取網(wǎng)絡(luò)，部分是由卷積層和最大池化層組合而成的模塊，整體結(jié)構(gòu)與VGG[16]極為相似，所以本文采用VGG16的主干特征提取網(wǎng)絡(luò)，如此也方便使用預(yù)訓(xùn)練權(quán)重進行遷移學(xué)習(xí).

在編碼器中，共進行了5次卷積模塊的操作，輸入圖像大小為512×512×3，前2個卷積模塊操作分別為2次卷積核為3×3的64通道和128通道的卷積操作，獲得的初步有效特征層再分別進行一次2×2的最大池化，獲得相應(yīng)的特征層.后3個卷積模板操作分別為3次卷積核為3×3的256通道、512通道和512通道的卷積操作用來獲得初步有效特征層.第3個卷積模塊與第4個卷積模塊再分別進行一次2×2的最大池化，第5個卷積模板不再做最大池化操作.編碼器的結(jié)構(gòu)圖如圖3所示.

圖3 編碼器的結(jié)構(gòu)圖

2.2.2 解碼器結(jié)構(gòu)

經(jīng)過編碼器的主干特征提取獲得5個有效特征層，對這5個特征層進行特征融合，所采用的融合方式是對特征層進行上采樣并堆疊.為了方便搭建網(wǎng)絡(luò)以及考慮到通用性，本文設(shè)計的U-Net與Ronneberger等所提出的U-Net[17]略有區(qū)別.本文采用上采樣時先進行2倍的上采樣，然后進行特征融合，這樣做的優(yōu)勢是最終得到的特征層與輸入圖片大小相同.解碼器的結(jié)構(gòu)圖如圖4所示.

圖4 解碼器的結(jié)構(gòu)圖

2.3 注意力機制

在傳統(tǒng)的U-Net中，為了避免在解碼過程中丟失重要細節(jié)信息，使用了跳躍連接的方式，將編碼過程中提取得到的映射直接鏈接到解碼器相對應(yīng)的層.但是這樣操作有個缺點，即提取的低級特征會存在很多冗余信息.利用注意力機制[18]能夠抑制模型學(xué)習(xí)與任務(wù)不相關(guān)的部分，強化學(xué)習(xí)與任務(wù)相關(guān)的特征，從原理上分為空間注意力與通道注意力.

Jaderberg等[19]提出的STN網(wǎng)絡(luò)可以根據(jù)任務(wù)的不同，自適應(yīng)地將數(shù)據(jù)進行空間變換.此網(wǎng)絡(luò)是基于空間的注意力思想.SENet[20]在ILSVRC 2017分類中，贏得了第一名，此網(wǎng)絡(luò)本質(zhì)上是基于通道的注意力模型，它只使用全局平均池化來壓縮通道特征.通過學(xué)習(xí)的方式自動獲取每個特征通道的重要性，然后通過獲取到的通道重要程度來提升有用信息和抑制無用信息.

本文采用的是通道注意力模塊與空間注意力模塊相結(jié)合的方式，將模塊添加到網(wǎng)絡(luò)的解碼器中，通道的每一個高級特征都可以看作是一種特定的響應(yīng)．通過挖掘通道映射之間的相互依賴關(guān)系，可以強調(diào)相互依賴特征映射，改善特定語義的特征表示．通道注意力模塊與空間注意力模塊相結(jié)合的模塊示意圖如圖5所示[15].

圖5 注意力模塊

通道注意力專注于有意義的輸入圖像，為了有效地計算通道注意力，需要壓縮輸入特征圖的空間維度.對于空間信息的聚合，常使用的方法是平均池化，而最大池化可以收集到難以區(qū)別物體之間的重要特征線索，來獲取更加詳細的通道注意力，所以通道注意力模塊同時使用平均池化和最大池化.然后將它們送入一個權(quán)重共享的多層感知機中，最后再將這些輸出的特征按對應(yīng)位置相加.通道注意力模塊如圖6所示[15].

圖6 通道注意力模塊

空間注意力專注于輸入圖像的有效信息的位置，可以很好地彌補通道注意力的不足.為計算空間注意力，沿著通道軸使用平均池化和最大池化，將它們相連，生成一個有效特征描述符.最后通過卷積層生成空間注意力.空間注意力模塊如圖7所示[15].

圖7 空間注意力模塊

2.4 損失函數(shù)

本文采用交叉熵與Dice損失相結(jié)合的形式作為語義分割的損失指標.交叉熵損失函數(shù)常用于分類問題，本文是對圖像細胞做分割，本質(zhì)也是對細胞與背景像素點做二分類，在二分類情況下，模型最后需要預(yù)測的結(jié)果有兩種情況，對于每個類別預(yù)測出的概率為pi和1-pi，交叉熵的表達式為

.

(1)

其中：yi代表樣本i的標簽，正類為1，負類為0；pi代表樣本i預(yù)測為正的概率.Dice系數(shù)CDice是一種對比圖像相似度的度量函數(shù)，通常情況下用于計算兩個圖像數(shù)據(jù)的相似程度，它的取值范圍為[0,1]，計算公式為

(2)

其中：N為圖像數(shù)目；qi為模型預(yù)測中的某一像素；gi為相同位置上金標準中的某一像素.CDice越大表示其預(yù)測結(jié)果與真實結(jié)果的重合度越大，預(yù)測效果就越好.作為損失函數(shù)，其值越小，代表模型訓(xùn)練效果越好.Dice損失的表達式為

(3)

其中，F(xiàn)Dice_Loss為Dice損失.本文將交叉熵與Dice損失的值相加后作為損失函數(shù)來評估模型訓(xùn)練的情況，其權(quán)重系數(shù)為1[21].

3 實驗與分析

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

本文采用的CSF圖像信息數(shù)據(jù)集來自北京協(xié)和醫(yī)院，采用HE染色技術(shù)獲取，經(jīng)過篩查得到29例可用數(shù)據(jù)作訓(xùn)練.這些訓(xùn)練數(shù)據(jù)集的標簽是由專業(yè)人員手動標注而成，具有專業(yè)性.其標簽的位深度為24 b，每個通道具有相同的像素數(shù)，且細胞像素設(shè)置為255，背景像素設(shè)置為0，所以標簽的視覺效果是一個二值圖像.

圖像經(jīng)過旋轉(zhuǎn)、鏡像、裁剪等方法將數(shù)據(jù)擴充至232例訓(xùn)練圖像.并且將數(shù)據(jù)大小統(tǒng)一為512×512，以便訓(xùn)練.將擴充后的數(shù)據(jù)集經(jīng)過一次拉普拉斯高通濾波處理，突出細胞邊緣與背景的灰度差值，增強細節(jié)，其拉普拉斯核取為

3.2 環(huán)境配置

本文所使用的設(shè)備為i7-8700處理器，16 GB內(nèi)存，64位Windows 10操作系統(tǒng)，NVIDIA GeForce GTX 1080 GPU加速處理圖像.并且使用了TensorFlow-gpu 1.13.1與keras 2.1.5的開源深度學(xué)習(xí)框架，Python編程語言，Pycharm集成開發(fā)環(huán)境，還使用了Numpy計算庫以及PIL中的圖像處理方法等.

3.3 實驗設(shè)置

將擴充后的數(shù)據(jù)經(jīng)過預(yù)處理后傳入網(wǎng)絡(luò)中作訓(xùn)練.網(wǎng)絡(luò)參數(shù)的優(yōu)化器使用Adam optimizer[22]，此優(yōu)化方法對高維度的數(shù)據(jù)進行快速擬合，利用梯度的一階矩陣估計和二階矩陣估計動態(tài)調(diào)整每個參數(shù)的學(xué)習(xí)率，經(jīng)過偏置校正后，每次迭代學(xué)習(xí)率都有個確定范圍，使得參數(shù)比較平穩(wěn).

為了均衡訓(xùn)練速度與訓(xùn)練效果，每次將2張圖像傳入網(wǎng)絡(luò)中進行訓(xùn)練.加快網(wǎng)絡(luò)訓(xùn)練速度，改善資源利用率，先將網(wǎng)絡(luò)進行凍結(jié)，在凍結(jié)網(wǎng)絡(luò)的情況下訓(xùn)練迭代50次，將更多的資源放在訓(xùn)練后面部分的網(wǎng)絡(luò)參數(shù).經(jīng)過反復(fù)測試驗證，凍結(jié)時的初始學(xué)習(xí)率設(shè)置為1×10-4效果最佳.迭代訓(xùn)練50次后，再將模型進行解凍，經(jīng)過反復(fù)測試，這時的學(xué)習(xí)率設(shè)置為1×10-5效果最佳.模型解凍后繼續(xù)迭代訓(xùn)練50次.當(dāng)訓(xùn)練時，每隔3次當(dāng)?shù)Ｐ托阅懿辉偬嵘龝r，則會采用每次將學(xué)習(xí)率降低原來學(xué)習(xí)率的0.5倍的方法來優(yōu)化模型性能.

本文將實驗數(shù)據(jù)按照8∶2的比例劃分為訓(xùn)練集與驗證集，將29例數(shù)據(jù)作為測試數(shù)據(jù)，來預(yù)測模型的實際分割效果.

3.4 評價指標與結(jié)果

本文采用了語義分割中3個最常用的評價指標進行分析分割結(jié)果，它們分別為相似系數(shù)EDice，平均交并比U，以及類別平均像素準確率A.本文以二分類為例，計算公式為

(4)

(5)

(6)

其中：TP，F(xiàn)P和FN分別為模型預(yù)測的真正例、假正例和假反例；P1和P2分別為類別1和類別2的像素準確率，定義為

(7)

(8)

經(jīng)過所有網(wǎng)絡(luò)的訓(xùn)練，選擇兩張具有代表性的CSF細胞圖像進行預(yù)測分割展示.圖8為細胞大小一般且排列相對稀疏的數(shù)據(jù)圖，圖9為細胞較小且排列密集的數(shù)據(jù)圖.

圖8 細胞稀疏圖像

圖9 細胞密集圖像

如圖所示，使用傳統(tǒng)Otsu閾值分割時，屬于背景的部分像素會被判定為細胞，對噪聲極其敏感，對灰度差異不明顯及不同目標灰度值有重疊的分割不明顯.而PSPnet網(wǎng)絡(luò)對醫(yī)學(xué)圖像的細節(jié)還原性較差，無法很好地識別醫(yī)學(xué)圖像的小目標物體.Segnet與DeeplabV3+模型分割尺寸相對差不多大小的細胞時，會存在細胞黏連的問題，并且分割尺寸相對較小且數(shù)量較多的細胞時，存在細胞分辨不清晰、分割模糊的問題.U-Net分割模型基本上實現(xiàn)了對所有細胞的有效分割，但是依然存在不足之處，該網(wǎng)絡(luò)模型對細胞邊界與背景之間的區(qū)分度略有不足.而使用本文的分割模型做預(yù)測分割時，上述所有問題均得到了改善.

使用上述分割模型分別對CSF細胞圖像做出預(yù)測，EDice，U和A的具體結(jié)果如表1所示.

表1 CSF分割結(jié)果

由表1可看出，使用本文方法對CSF細胞圖像做細胞分割要優(yōu)于其他分割方法.

為了驗證本文方法的普適性，選用2018 Data Science Bowl公開數(shù)據(jù)集作為參照實驗數(shù)據(jù).該數(shù)據(jù)包含大量分割的核圖像，圖像是在各種條件下獲得的，細胞類型、放大倍數(shù)和成像方式各不相同.使用此數(shù)據(jù)得到預(yù)測分割的各項指標具體結(jié)果如表2所示.

表2 2018 Data Science Bowl數(shù)據(jù)集分割結(jié)果

由表2可看出，在2018 Data Science Bowl數(shù)據(jù)集分割中，本文方法要優(yōu)于其他分割方法.

4 結(jié) 論

本文以CSF細胞為分割研究對象，采用基于注意力機制的U-Net網(wǎng)絡(luò)，經(jīng)過數(shù)據(jù)擴充與預(yù)處理，引入遷移學(xué)習(xí)，提高資源利用率，并且與Otsu，PSPnet，Segnet，DeeplabV3+以及U-Net做了對比實驗，結(jié)果表明，本文方法要優(yōu)于其他分割方法.并且通過公開數(shù)據(jù)集2018 Data Science Bowl做了進一步對比實驗，實驗結(jié)果進一步驗證了本文分割方法在各項指標上均優(yōu)于其他分割方法.