同时支持框、点、文本提示

智源杜雨新 投稿 凹非寺

量子位 | 公众号 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