智源杜雨新 投稿 凹非寺
量子位 | 公衆号 QbitAI
醫學領域的通用分割模型來啦，釋出即開源！
來自智源，模型名為SegVol，劃重點：
是第一次實作同時支援框（box）、點（point）和文本（text） prompt進行任意尺寸原分辨率的3D體素分割。
要知道，此前深度學習在醫學圖像分割方面已經取得了顯著進展，但仍然缺乏一種能夠通用分割各種解剖類别、易于使用者互動的基礎分割模型。
而SegVol正是一種通用的互動式醫學體素圖像分割模型，能夠對200多個解剖類别進行分割。
在實驗測試中，SegVol在多個benchmark中表現出色。特别在三個具有挑戰性的病變資料集上，SegVol比nnU-Net的Dice得分高20%左右。
目前，SegVol的代碼和權重已經在GitHub上公開。開源的模型權重檔案包括：
- 使用96k CTs預訓練2,000 epochs的ViT模型；
- 在預訓練基礎上，使用6k Masked CTs在A100上訓練30✖️21✖️8個GPU小時得到的SegVol。
更多細節，我們接着往下看。
一口氣分割200多個解剖類别
SegVol模型架構分為4個子產品：圖像編碼器、文本編碼器、提示編碼器、掩模解碼器。
支援文本、點、框三種互動提示，其中為防止過拟合文本編碼器當機，提示編碼器則是整合三種提示的嵌入，掩模解碼器使用自注意力和交叉注意力預測分割掩模。
研究人員在96k CTs上對模型進行預訓練，并使用僞标簽解耦資料集和分割類别之間的虛假關聯。
通過将語言模型內建到分割模型中，并在25個資料集的200多個解剖類别上進行訓練，進而實作文本提示分割。
△ (a)聯合資料集概覽，(b)聯合資料集中掩碼數量排名前30的标簽，人體四個主要部位的掩碼标簽數量占比，(c)樣例
同時協同語義提示(text prompt)和空間(point, box prompt)提示，實作高精度分割。
此外，研究人員還設計了一種名為zoom-out-zoom-in的機制，可以顯著降低計算成本，同時保持精确分割。
實驗結果
研究人員在多個分割資料集上充分評估了SegVol。
19種重要解剖結構的實驗結果
在prompt learning的支援下，SegVol能夠支援200多個類别的分割。
其中19種重要解剖結構的實驗結果如下：
如上表所示，肝髒的Dice得分高達96.13%，19個主要目标的平均得分為83.02%。
研究人員表示，其強大的通用分割功能來自于spatial和semantic的複合prompt。
一方面，spatial prompt可以讓模型了解分割目标的具體空間和位置。由上表可知，對于各種器官的平均分割結果，“box+text” prompt的Dice score比text prompt高5.85%。
另一方面，semantic prompt分割目标的語義指代，消除了多種可能的結果。
這也反映在上表中，“point+text” prompt的平均Dice score比單獨使用point prompt高4.62%。spatial prompt和semantic prompt互相支援，最終賦予模型強大的分割能力。
對比實驗
研究人員還将SegVol與五個重要資料集上的四種最先進的方法進行了比較。
對于體量在數十到數百個病例的醫學體素圖像資料集，由于SegVol能夠在25個資料集上聯合訓練，顯著優于在單個資料集上訓練的傳統分割模型。
從表2可以看出，SegVol在肝、腎、脾等easy類别上超過了傳統模型，平均Dice score達到了94.98%。
研究人員認為這主要是由于它從其他資料集的相同或相似類别中學到了更多的知識。
更重要的是，這種方法在肝惡性良性腫瘤、肺惡性良性腫瘤、腎上腺等hard類别的分割中也保持領先地位。
SegVol對hard類的平均Dice score比排名第二的nnU-net高14.76%。
原因是SegVol可以通過spatial prompt和semantic prompt獲得先驗資訊，進而增強對hard樣本的了解，顯著改善了分割結果。
病竈分割能力
此外，研究人員使用nnU-net作為基線模型，該模型在傳統的醫學體素圖像分割模型中表現出最強的分割能力。
如下表3所示，SegVol分割這些具有挑戰性的病變的能力明顯優于nnU-net。
在這三個病變資料集中，SegVol的Dice score超過nnU-net 19.58%，這代表在複雜體素病竈分割方面SegVol的重大進步。
下圖（c）給出了一系列示例，展示了nnUnet和本文方法的病變分割性能，這些例子包括肝惡性良性腫瘤、結腸癌和肺惡性良性腫瘤。
△病竈分割
可視化結果顯示，與nnU-net産生的結果相比，SegVol重建的這些病變解剖結構更接近于Ground Truth。
消融實驗
Zoom-out-zoom-in機制：
研究人員在MSD-Liver資料集上進行了消融研究，以評估Zoom-out-zoom-in機制的貢獻。
MSD-Liver資料集包括肝髒和肝惡性良性腫瘤兩個類别，允許研究Zoomout-zoom-in機制對“MegaStructures”和“MicroStructures”目标分割效果的影響。
如表4所示，将Zoom-out-zoom-in機制應用于SegVol模型使肝髒類别的Dice score提高了6.07%。這種提升在肝惡性良性腫瘤類别上更為明顯，Zoom-out-zoom-in機制将SegVol的肝惡性良性腫瘤Dice score提高了21.32%。
有趣的是，Zoom-out-zoom-in機制對point prompt分割肝髒結果的改善十分微小。這可能歸因于global一級的point prompt相對稀疏，當zoom in到local區域時，其稀疏性變得更加明顯，進而限制了該機制的潛力。
Dataset Scale：
資料規模是基礎模型建構的關鍵因素之一。研究人員進行了消融研究，以研究Image和Mask的數量對SegVol性能的影響。
他們将包含13個重要器官的BTCV資料集作為測試錨點，分别對1、2和8個資料集上訓練了500個epoch的模型，以及在25個資料集上訓練的最終模型進行評估。
△資料集scale。(a)在不同數量的資料集中CTs和相應的Ground Truth Mask數量，(b)不同資料規模訓練SegVol的Dice Score
詳細的結果如上圖a和b所示。作為輕量級模型，當隻使用一個資料集時，SegVol的性能不是最優的。
然而，随着資料量的增加，SegVol的Dice score顯著增加，特别是在使用text prompt進行分割的情況下。因為text prompt嚴重依賴帶有語義資訊的ground truth mask的數量。
總之，研究人員提出了SegVol，一個互動式的通用醫學體素圖像分割的基礎模型。
該模型是使用90k無标注資料和25個開源分割資料集訓練和評估的。與最強大的傳統體素分割方法nnU-net（自動為每個資料集配置參數）不同，SegVol的目的是将各種醫學體素分割任務統一到一個單一的架構中。
SegVol作為一個通用的分割工具能夠對超過200個解剖目标産生準确的分割響應。
此外，與傳統方法相比，SegVol具有最先進或接近最先進的體素分割性能，特别是對于病竈目标。盡管具有通用性和精确性，但與其他體素分割方法相比，SegVol保持了輕量級架構。
SegVol作為一個開源的基礎模型，将很容易适用于廣泛的醫學圖像表征和分析領域，可以很容易地被研究人員和從業人員內建和利用。
該研究論文一作杜雨新，作者Fan Bai同時來自港中文，作者Tiejun Huang同時來自北大，通訊作者為Bo Zhao。
論文連結：https://arxiv.org/abs/2311.13385
— 完 —
量子位 QbitAI · 頭條号簽約
關注我們，第一時間獲知前沿科技動态