基于跨任務一致性的半監督肝臟CT圖像分割

李明漾,王慶鳳,陳立偉,黃 俊,周 瑩

(1.西南科技大學 計算機科學與技術學院,四川 綿陽 621000;2.綿陽市中心醫院 放射科,四川 綿陽 621000)

0 引 言

肝臟作為人體非常重要的器官發病率非常高,肝臟是容易病變產生腫瘤的高發部位,肝癌屬于惡性腫瘤,死亡率極高。目前計算機斷層掃描(Computed Tomography,CT)已經成為了醫學成像的主要方法之一,快速精準地從CT圖像中分割出肝臟區域對醫生診斷肝癌、定性分析以及制定治療方案具有重要意義,肝臟分割結果可以為醫生臨床診斷提供可靠的診斷依據。目前臨床診斷中,基于CT圖像對肝臟的分割是由經驗豐富的醫生手動進行,分割結果容易受醫生主觀判斷和分析的影響,增加誤判或漏判的可能性。

近幾年,研究人員提出了一些基于深度學習的醫學圖像分割方法,在肝臟分割領域中取得了顯著成果,分割精度大大提高。現有大多數方法使用以深度卷積神經網絡為代表的深度學習方法,它們主要是采用全監督的設置來完成肝臟CT圖像分割任務,作為全監督學習的一個重要特性,其性能在很大程度上依賴于足夠數量的標記數據。然而在醫學領域標注可靠的注釋圖像非常耗時耗力。據統計,手動分割一個肝臟序列平均需要30 min[1]。標記數據的缺乏引導了許多研究者對注釋效率的研究,例如弱監督學習、無監督學習、半監督學習等。由于獲取少量帶有標記數據和大量未標記數據在臨床上是可行的,所以該文將基于半監督學習的方法分割肝臟CT圖像。

主要貢獻:為了解決醫學圖像中肝臟標注數據缺乏的問題,提出了一種基于跨任務一致性的半監督肝臟CT圖像分割方法。該方法具有多解碼器的網絡模型以增加模型擾動的多樣性,同時引入輔助預測任務學習肝臟的幾何信息提高模型分割性能。

1 相關工作

基于深度學習的圖像分割已經在醫學圖像領域取得了顯著的成功,所以首先需要回顧一些經典的醫學圖像分割方法和基于半監督學習的醫學圖像分割方法。

1.1 用于醫學圖像分割的深度學習方法

深度神經網絡對醫學圖像分割獲得了較好的效果,許多優秀的醫學圖像分割算法已經得到了廣泛的應用。全卷積神經網絡(FCN)[2]是醫學圖像分割領域中極其重要的一部分,FCN能夠提取圖像中目標區域的高層次特征與低層次特征并將兩者進行結合,實現了圖像的自動分割。U-Net[3]已經成為一種經典醫學圖像分割框架,具有編碼-解碼結構成對稱形狀,并能夠通過跳躍連接將編碼器和解碼器得到的不同分辨率的特征信息進行融合。He等人[4]引入了殘差結構思想,提出了ResNet用于醫學圖像分割,該結構將輸入特征圖與輸出特征圖進行特征融合,在增加網絡深度的同時避免了梯度爆炸或消失等問題。李秀華等人[5]將U-Net中編碼階段采用混合空洞卷積替換原始卷積并在跳躍連接中添加注意力機制,改進后提高了肝臟腫瘤分割精度。Cicek等人[6]將U-Net擴展到3D醫學圖像分割領域用于體積分割,其思想是采用3D卷積核提取特征,捕獲圖像豐富的空間信息。戴振暉等人[7]通過在3DU-Net添加殘差塊并結合Transformer模塊提出了Res-Swim-UNet模型,實現了肝臟及腫瘤的自動分割。Milletari等人[8]提出了V-Net,他們將殘差結構與3DU-Net結合并將下采樣的池化操作用步長為2的卷積代替,上采樣使用轉置卷積,并使用3D圖像輸入網絡模型。張宇嘯等人[9]提出了一種輕量級V-Net并結合深監督策略的肝臟及腫瘤分割算法,在減少模型的參數量和計算量的同時提高了分割精度。

1.2 半監督醫學圖像分割

半監督學習在醫學圖像分割領域已經得到了廣泛的研究和應用。半監督學習方法是需要將標記數據和未標記數據都用來訓練模型,目的是充分利用無標記數據來提高模型性能。Bai等人[10]采用了自訓練式半監督學習的思想,利用標記數據訓練初始模型,然后利用訓練好的初始模型預測未標記數據生成偽標簽,挑選出高質量偽標簽從而擴展帶標記的數據集。劉清清等人[11]提出的自訓練式半監督肝臟分割框架是通過改進的注意力模塊的3D scSE-Unet作為分割網絡進行循環迭代訓練。另外一種基于圖的算法假設兩個相似的樣本具有相同的輸出[12]。對抗性學習也是半監督學習的一種流行方法。Fang等人[13]通過結合對抗性訓練策略區分像素是來自預測分割結果還是標簽樣本。Zhang等人[14]采用對抗性學習使注釋圖像與未注釋圖像產生相似的分割輸出。

最近,越來越多的基于一致性正則化的半監督學習方法用來進行醫學圖像分割。其利用未標記數據的主要思想是基于假設,即通過將輸入數據添加小擾動,模型輸出產生的預測結果應該是一致的。例如,Sohn等人[15]在執行訓練樣本強弱條件下,檢驗數據擾動預測的一致性。Laine等人[16]提出的π模型方法簡單有效,為了執行可靠的預測,通過時間集合策略用前幾個時期的預測來擴展π模型。Tarvainen等人[17]提出了均值教師模型(MT),使用教師模型來產生學生學習時的目標,學生模型利用教師模型產生的目標學習,并討論模型參數一致性。Yu等人[18]進一步在MT模型中引入不確定性估計指導學習。楊星宇等人[19]提出了一種多時期一致的半監督模型(MPC-MT)實現無造影劑圖像的肝臟腫瘤分割,其思想是基于均值教師模型,該模型在多時期磁共振圖像之間建立一致性正則化。Li等人[20]提出的SASSNet模型引入了帶符號距離圖(SMD)[21],進一步學習器官的幾何形狀信息。Luo 等人[22]提出了基于雙任務一致性的半監督學習框架,通過在分割任務和水平集函數回歸任務之間建立一致性正則化約束,實現了任務級的正則化。與以往的方法不同,該文提出的半監督學習框架包含多個解碼器,同時添加輔助回歸任務學習肝臟的幾何形狀信息。該框架在主副解碼器的不同任務之間執行跨任務一致性正則化以利用大量未標記數據,提高模型分割肝臟的性能。

2 文中方法

2.1 模型框架

本節將介紹提出的基于跨任務一致性的半監督肝臟分割模型框架,總體結構如圖1所示。該文采用V-Net作為基礎網絡。它由一個共享編碼器、一個主解碼器和兩個副解碼器組成,每個解碼器有兩種預測任務。這三個解碼器在結構上是相同的,僅僅在上采樣方法上不同。定義副解碼器D1,D3的上采樣方法分別采用三線性插值、最近插值。主解碼器D2的上采樣方法與V-Net保持一致,采用轉置卷積。在主副解碼器中使用Dropout層以實現模型層面的擾動。

圖1 基于跨任務一致性的半監督學習肝臟分割框架

2.2 跨任務一致性正則化

在一般的基于一致性的半監督學習方法中,一致性損失設計用于數據級上的平滑預測。與之前方法不同,該文利用預測回歸任務得到的有符號距離圖(SDM)訓練網絡模型,該任務可以使模型學習目標特征的幾何輪廓信息,其中回歸任務預測的SMD可由解碼器后添加一個1×1×1卷積塊和Tanh激活函數實現。可以通過T(·)轉換函數將分割標簽轉換成SMD,轉換函數定義如下:

(1)

在SMD中的每個點(體素)表示其之間的距離和最近邊界,所以‖x-y‖2表示體素x,y之間的歐氏距離,?C表示目標對象的表面,SMD中點的值為0則表示位于目標對象的表面。Cin和Cout分別表示為目標對象的內部和外部。對于目標對象的內部區域體素取負值,外部區域取值反之。由于有符號距離圖(SMD)與像素分割概率圖不在同一個域內,且轉換式不可微的性質,為了建立主副解碼器之間的跨任務一致性,該文將SMD通過T-1(·)函數轉換成近似的像素分割概率圖。采用Luo等人[22]提出的方法中的平滑近似函數來實現轉換,轉換公式定義為:

(2)

其中,z是體素x的有符號距離圖的距離值,k是一個超參數盡可能地設置較大的值來近似變換接近分割真實值。此公式的目的是將SDM映射到與像素分割概率圖同一域內,SDM變換成近似分割概率圖使得建立跨任務一致性正則化有意義。

2.3 半監督訓練損失函數

在上述半監督學習方法描述中,如圖1所示,文中的損失函數由無監督損失Lmrc和監督損失Lsup組成。對于監督損失:文中方法的監督損失由標記圖像集合Dl上對應在主副解碼器的雙任務優化目標組成,采用Dice loss和Cross-entropy loss作為三個解碼器像素級分割預測任務的混合監督損失,計算如下:

(3)

式中,fseg表示像素級分割圖,Y表示分割標簽圖像。

該文引入的回歸任務作為模型中的重要部分,也將計算SDM產生的預測誤差,其監督損失定義如下:

(4)

式中,fsdm表示有符號距離圖(SDM),T(y)是分割標簽圖像Y通過T(·)函數轉換成的符號距離圖真實標簽。總的監督損失由以上兩個部分的損失加權求和所得。

Lsup=Lseg+λLsdm

(5)

式中,λ是平衡分割任務和回歸任務損失項中的權重系數,將其設置為0.3。對于整體訓練集D的無監督損失:模型在主副解碼器上采樣操作的不同以及Dropout層的作用,形成了模型擾動,在主解碼器與副解碼器之間建立跨任務一致性損失,使模型能夠利用大量未標記數據進行學習。在訓練過程中,在主解碼器D2的像素級分割圖和副解碼器D1,D3輸出的有符號距離圖,以及主解碼器的有符號距離圖和副解碼器D1,D3輸出的像素級分割圖之間建立一致性約束進行優化,具體一致性無監督損失計算如下:

(6)

Ltotal=Lsup+λ2Lmrc

(7)

繼前面的研究工作,該文采用高斯升溫函數策略λ2=0.1×e-5(1-t/tmax)來控制監督損失和無監督的一致性損失之間的平衡。t表示當前訓練步驟,tmax表示最大訓練步驟。高斯升溫函數可以使模型在訓練時更加高效,因為在訓練的前期階段模型預測結果相對來說可靠度低。

3 實驗結果與分析

3.1 數據集及預處理

該文采用肝臟腫瘤分割公開挑戰賽LiTS2017數據集進行方法驗證,該數據集包含131例含有肝臟部位CT掃描圖像,并且由3名放射科醫生手動注釋標簽。在數據集中隨機挑選104例CT圖像用于訓練,剩余27例CT圖像用于測試。最后,對數據集做預處理操作,將原始圖像的CT值(HU)調整到[-200,200]的范圍,并歸一化到[0,1]。為了降低計算機性能的要求,去除大量無關背景,將原始圖像進行中心裁剪,最終得到尺寸大小為256×256×256的CT圖像。

3.2 實施細節及評價指標

文中的實驗是基于Pytorch框架實現。實驗使用24 GB內存的NVIDIA RTX 3090 GPU進行訓練。該網絡框架采用SGD優化器經過6 000次的迭代訓練,將初始學習率設置為0.01,每經過2 500次迭代學習率衰減0.1,每次迭代的Batch size設置為4,其中包含兩個標記圖像和兩個未標記圖像。該文將采用以下評價指標來驗證提出方法的有效性:Dice相似系數(Dice Similarity Coefficient)、Jaccard系數(Jaccard Coefficient)和平均表面距離(Average Surface Distance,ASD)。各項評價指標計算公式如下:

(8)

(9)

(10)

式中,TP(True Positive)表示肝臟分割像素正確分類的預測結果;FP(False Positive)表示肝臟分割像素錯誤的預測結果;FN(False Positive)表示預測像素分割結果分類為肝臟實際為背景的像素;Aseg表示預測分割結果的邊界;Agt表示真實分割標簽的邊界;p1,p2表示邊界上的像素點。

3.3 對比實驗與定量分析

該文在LiTS2017數據集上驗證了該方法的有效性。與UA-MT[18],SASSNet[20],DTC[22],SMTL[23]等半監督方法進行比較。上述半監督方法都是以V-Net模型作為主干網絡,為了體現與其他半監督方法對比的有效性,控制變量,文中半監督框架也采用V-Net網絡模型作為基礎網絡。該框架在V-Net基礎上添加了兩個并行解碼器以增加模型擾動的多樣性。同時以上所述的半監督方法都是在標記數據和未標記數據設置數量相同比例下進行比較,在LiTS訓練集樣本中分別設置10%帶標注樣本和20%帶標注樣本的實驗組,如表1和表2所示,其中V-Net[8]方法采用全監督設置進行訓練作為基線對比。

表1 在10%標記數據訓練下與其他半監督方法的對比實驗

表2 在20%標記數據訓練下與其他半監督方法的對比實驗

首先,通過表1和表2的實驗結果可知,所有的半監督方法與利用相同標記數據數量設置下的全監督V-Net方法相比都得到了更優的分割結果,達到不錯的性能增益。其中表1顯示了使用10%標記數據訓練時, SASSNet,UA-MT,DTC的Dice系數、Jaccard系數、ASD比較接近,性能增益不明顯,其中SMTL在ASD評判標準中表現最好。文中方法在Dice系數和Jaccard系數評判標準上達到了93.13%和87.43%,相比于V-Net[8]方法全監督設置下,該方法通過利用未標記數據訓練模型將Dice系數提高了4.41百分點,Jaccard系數提高了6.34百分點。表2顯示了在使用20%標記數據訓練時,文中方法在Dice系數和Jaccard系數相比其他半監督方法產生分割精度增益更明顯。該方法相比SMTL將Dice系數提高了1.34百分點,Jaccard系數提高了2.38百分點,且在ASD評判標準下比UA-MT,SASSNet,DTC擁有更優的分割結果。

綜合兩個表格實驗結果表明,文中方法的分割結果和真實結果擁有更高的重疊率和相似度,在Dice和Jaccard評價標準上都優于其他半監督方法。增加了輔助任務使模型在肝臟的某些邊界區域能夠產生更好的分割結果,并且添加兩個并行解碼器建立一致性約束能充分利用未標記數據學習,提高分割結果質量。值得一提的是,文中方法在僅使用20%帶標記數據時就能達到接近100%帶標記數據全監督訓練時分割結果的Dice,Jaccard精度,證明了該方法能夠有效地利用大量無標記數據獲取信息,提高分割性能。

3.4 實驗結果可視化

最終,隨機選取兩名代表性患者的分割結果進行可視化分析。圖2展示了患者在V-Net網絡使用不同標記樣本數量全監督設置下的分割結果與文中方法分割結果的2D可視化對比。

圖2 與V-Net基礎網絡在不同標記數量的全監督設置下對比

其中(a)是原始圖像,(b)專家標簽,(c)和(d)分別是V-Net網絡使用20%和100%標記數據在全監督情況下的分割結果,(e)是文中方法可視化結果。在使用相同標記數據數量訓練網絡模型的情況下,提出的半監督方法明顯比全監督設置下的分割結果更好,尤其對肝臟的邊界區域分割更準確更接近于真實標簽,甚至部分細節更優于100%標記數據全監督設置下的分割結果。這是因為引入的輔助預測任務來學習肝臟的全局幾何信息,它能夠使模型學習到更多的肝臟邊界信息,從而提高模型的分割性能。

4 結束語

提出了一種基于跨任務一致性的半監督肝臟分割方法。該方法以V-Net作為基礎網絡,添加兩個副解碼器以形成一定程度上的模型擾動,同時在每個解碼器引入回歸任務,利用主副解碼器之間跨任務一致性約束建立損失函數,從而使模型能夠利用大量無標記數據學習。同時,該方法與其他半監督方法相比,在LiTS數據集上肝臟分割結果精度更優,并且在僅使用20%帶標記數據的情況下,肝臟分割結果在Dice和Jaccard的評價指標幾乎接近100%帶標記數據全監督設置下的分割結果。這意味著該方法在實用性方面具有一定優勢。