TMI2024 | 阿大、同濟等提出TraCoCo,用於3D醫學影像半監督分割

ScienceAI發表於2024-10-29

圖片

編輯 | ScienceAI

3D 醫學影像分割方法已經取得了成功,但它們對大量體素級標註資料的依賴是一個需要解決的缺點,因為獲取這些標註的成本很高。

半監督學習(SSL)透過使用大量未標註資料和少量標註資料進行模型訓練,解決了這一問題。

最成功的 SSL 方法基於一致性學習,即透過最小化從擾動檢視中獲得的模型響應之間的距離來實現的。這些擾動通常在檢視之間保持空間輸入上下文的一致性,這可能導致模型從空間輸入上下文中學習分割模式,而不是從前景物件中學習。

在最新的研究中,澳大利亞阿德萊德大學(The University of Adelaide)、同濟大學等機構的研究人員提出了 TraCoCo,這是一種一致性學習 SSL 方法,它透過改變輸入資料檢視的空間輸入上下文來進行擾動,使模型能夠從前景物件中學習分割模式。

此外,該團隊提出了一種新的“Confident Regional Cross entropy (CRC)”損失,該損失提高了訓練收斂性,並保持了對共訓練偽標籤錯誤的魯棒性。該方法在多個3D資料基準上達到了最先進的結果,例如左心房(LA)、胰腺CT(Pancreas)、和腦腫瘤分割(BraTS19)。該方法在一個 2D 切片基準——自動心臟診斷挑戰(ACDC)上也取得了最佳結果,進一步證明了其有效性。

圖片

文章地址:https://ieeexplore.ieee.org/abstract/document/10695462

程式碼地址:https://github.com/yyliu01/TraCoCo

背景

圖片

3D醫學影像分割方法雖然已取得成功,但它們依賴於大量體素級別標註資料,這是一個亟需解決的問題,因為獲得這些標註的成本非常高。為了克服這一挑戰,半監督學習(SSL)透過結合大量未標註資料和少量標註資料來訓練模型,從而減少對標註資料的需求。

最成功的 SSL 方法基於一致性學習,透過最小化模型在未標註資料的不同擾動檢視下的響應差異來實現一致性。然而,這些擾動通常保持檢視之間的空間輸入上下文較為一致,這可能導致模型從空間上下文中學習分割模式,而不是從前景物件中學習。

為解決這一問題,研究人員提出了一種新的 SSL 方法,稱為 Translation Consistent Co-training(TraCoCo)。TraCoCo 透過改變輸入資料檢視的空間輸入上下文來擾動檢視,使模型能夠從前景物件中學習分割模式。

同時,該研究還提出了一種新的 Confident Regional Cross-Entropy(CRC)損失,該損失旨在提高訓練的收斂性並保持對共訓練偽標籤錯誤的魯棒性。

圖片

Translation Consistent Co-training (TraCoCo)

TraCoCo是一種基於互助學習(Co-training)框架的方法,該框架透過翻譯一致性(Translation Consistency)來擾動輸入資料的空間上下文,從而減少模型對背景模式的“記憶”,確保模型專注於前景物件的分割。

模型框架

TraCoCo 包含兩個初始引數不同的獨立網路(通常為 VNet或 3D-UNet)。輸入資料透過隨機裁剪生成兩個子體積(sub-volumes),這兩個子體積的空間上下文不同,但在體素網格中存在重疊區域。

模型目標

為了確保模型在不同空間上下文下的前景分割結果一致,TraCoCo 引入了多種損失函式,包括監督學習損失、半監督學習損失和翻譯一致性損失。本文的方法實現透過最小化以下損失函式來進行訓練:

圖片

其中,分別由監督學習損失,半監督學習損失,和翻譯一致性損失組成。

監督學習損失(Supervised Learning Loss)

監督學習損失透過帶標註的資料集 D_L 進行計算,包括體素級交叉熵損失(Cross-Entropy Loss)和 Dice損失,用來最佳化模型的分割效能。損失函式定義為:

圖片

Translation Consistency Loss

Translation Consistency Loss 主要目標是確保模型在不同空間上下文下的分割結果一致。具體實現過程如下:從訓練體積中隨機提取兩個中心不同的子體積

損失函式定義為:

圖片

圖片
圖片

其中,公式 4 是 Kullback-Leibler(KL)散度,用於計算兩個子體積在重疊區域的分割結果之間的差異;公式 5 是基於負熵的正則化損失,用於平衡訓練體素中的前景和背景類。

半監督學習損失(Semi-supervised Learning Loss)

半監督學習損失用於強化兩個模型之間的分割一致性,Confident Regional Cross-Entropy (CRC) Loss, 其定義為:

圖片

圖片

3D CutMix

為了進一步提高訓練的泛化性,本文采用了 3D CutMix 技術。實現為隨機生成一個 3D 二值掩碼,該掩碼包含一個隨機定義位置和尺寸的“1”方塊。掩碼應用在未標註資料和偽標籤上,定義為:

圖片

實驗

本文在四個公開的 3D 醫學影像半監督分割資料集上進行了實驗,包括:

左心房(Left Atrium, LA):100 個 3D MRI 體積,其中 80 個用於訓練,20 個用於測試。

胰腺 CT(Pancreas-CT):82 個對比增強的 CT 掃描,採用特定的預處理步驟。

腦腫瘤分割 2019(BraTS19):包含 335 個腦部 MRI 樣本,每個樣本包含四種掃描型別(T1, T1-ce, T2, FLAIR)。

自動心臟診斷挑戰(ACDC):該資料集包含 100 個心臟 MRI 掃描,本文使用了其 2D 切片進行半監督學習。

實驗設定:採用了 VNet 和 3D-UNet 作為模型架構,實驗評估指標包括Dice、Jaccard、平均表面距離(ASD)和 95% Hausdorff 距離(95HD)。實驗結果顯示,TraCoCo 在這些基準測試上均優於現有的最先進(SOTA)方法,特別是在標註資料較少的情況下。

圖片

圖片
圖片
圖片

圖片

圖片

圖片

圖片

圖片

圖片

圖片

圖片

圖片

圖片

總結

本文提出了 TraCoCo 和 CRC 損失在 3D 和 2D 醫學影像分割任務中的有效性,證明了其在減少對背景模式的“記憶”以及提高前景物件分割精度方面的能力。

未來的工作計劃包括擴充套件 TraCoCo 到更多複雜的醫學影像任務,並在多模態醫學影像中探索其半監督學習能力 。

相關文章