开源可试用！智源提出SegVol：通用且可交互的医学体素分割模型-CSDN博客

©作者 | 杜雨新

单位 | 智源人工智能研究院实习生

研究方向 | 医学图像分析

论文题目：

SegVol: Universal and Interactive Volumetric Medical Image Segmentation

论文链接：

https://arxiv.org/abs/2311.13385

代码链接：

https://github.com/BAAI-DCAI/SegVol

我们很高兴介绍在上周开源的 SegVol 医学通用分割模型。与过去一些很棒的Medical SAM工作不同，SegVol 是第一个能够同时支持 box，point 和 text prompt 进行任意尺寸原分辨率的 3D 体素分割模型。作为一个便捷的通用分割工具，我们已将 SegVol 代码和模型开源，欢迎大家使用。

目前开源的模型权重文件包括（1）使用 96k CTs 预训练 2,000 epochs 的 ViT 模型，（2）在预训练基础上，使用 6k Masked CTs 在 A100 上训练个 GPU 小时得到的 SegVol。最新进展请关注 GitHub 仓库的更新，如果有疑惑或建议可以写评论、开 issue 或私信，欢迎大家讨论。

摘要

精确的医学图像分割为临床研究提供了富有意义的结构信息。尽管深度学习在医学图像分割方面已经取得了显著的进展，但仍然缺乏一种能够通用分割各种解剖类别且易于用户交互的基础分割模型。

本文提出一种通用的交互式医学体素分割模型——SegVol。通过在 90k 无标注 CTs 和 6k 分割 CTs 数据上进行训练，该基础模型支持 point，box 和 text prompt，能够对 200 多个解剖类别进行分割。大量的实验证明，SegVol 在多个 benchmark 中表现出色。特别在三个具有挑战性的病变数据集上，我们的方法比 nnU-Net 的 Dice 得分高 20% 左右。SegVol 的代码和权重已经在 Github 上公开。

核心贡献

1. 在 96k CTs 上对模型进行预训练，并使用伪标签解耦数据集和分割类别之间的虚假关联。

2. 通过将语言模型集成到分割模型中，并在 25 个数据集的 200 多个解剖类别上进行训练，从而实现文本提示分割。

3. 协同语义提示（text prompt）和空间（point, box prompt）提示，实现高精度分割。

4. 设计了一种 zoom-out-zoom-in 机制，显著降低计算成本，同时保持精确分割。

▲ 图1（a, b）模型结构图。（c, d）zoom-out-zoom-in 机制图。

▲ 图2（a）联合数据集概览。（b）联合数据集中，掩码数量排名前 30 的标签，人体四个主要部位的掩码标签数量占比。（c）样例。（a）中人体图来自 brgfx on Freepik

实验

我们在多个分割数据集上充分评估了 SegVol。

1）19 种重要解剖结构的实验结果

在 prompt learning 的支持下，SegVol 能够支持 200 多个类别的分割。我们选择了 19 个重要的解剖目标来展示其强大的分割能力，如表 1 所示。肝脏的 Dice 得分高达 96.13%，而 19 个主要目标的平均得分为 83.02%。其强大的通用分割功能来自于 spatial 和 semantic 的复合 prompt。

一方面，spatial prompt 可以让模型理解分割目标的具体空间和位置。由表 1 可知，对于各种器官的平均分割结果，“box+text” prompt 的 Dice score 比 text prompt 高 5.85%。

另一方面，semantic prompt 分割目标的语义指代，消除了多种可能的结果。这反映在表 1 中，“point+text” prompt 的平均 Dice score 比单独使用 point prompt 高 4.62%。spatial prompt 和 semantic prompt 相互支持，最终赋予模型强大的分割能力。

2）对比实验

如表 2，我们将 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 样本的理解，显著改善了分割结果。

▲ 图3 数据集 scale 和病灶分割。（a）在不同数量的数据集中，CTs 和相应的 Ground Truth Mask 数量。（b）不同数据规模训练 SegVol 的 Dice Score。（c）病灶分割。

3）病灶分割能力

我们使用 nnU-net 作为基线模型，它在传统的医学体素分割模型中表现出最强的分割能力。如表 3 所示，SegVol 分割这些具有挑战性的病变的能力明显优于 nnU-net。在这三个病变数据集中，SegVol 的 Dice score超过 nnU-net 19.58%，这代表在复杂体素病灶分割方面 SegVol 的重大进步。

图 3c 给出了一系列示例，展示了 nnUnet 和我们的方法的病变分割性能。这些例子包括肝肿瘤、结肠癌和肺肿瘤。可视化结果显示，与 nnU-net 产生的结果相比，SegVol 重建的这些病变解剖结构更接近于 Ground Truth。

4）消融实验

Zoom-out-zoom-in 机制：我们在马上到！-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 个数据集上训练的最终模型进行评估。详细的结果如图 3a 和 b 所示。

作为轻量级模型，当只使用一个数据集时，SegVol 的性能不是最优的。然而，随着数据量的增加，SegVol 的 Dice score 显著增加，特别是在使用 text prompt 进行分割的情况下。因为 text prompt 严重依赖带有语义信息的 ground truth mask 的数量。

总结

我们提出了 SegVol：一个交互式的通用医学体素分割的基础模型。该模型是使用 90k 无标注数据和 25 个开源分割数据集训练和评估的。与最强大的传统体素分割方法 nnU-net（自动为每个数据集配置参数）不同，SegVol 的目的是将各种医学体素分割任务统一到一个单一的架构中。SegVol 作为一个通用的分割工具能够对超过 200 个解剖目标产生准确的分割响应。

此外，与传统方法相比，SegVol 具有最先进或接近最先进的体素分割性能，特别是对于病灶目标。尽管具有通用性和精确性，但与其他体素分割方法相比，SegVol 保持了轻量级架构。SegVol 作为一个开源的基础模型，将很容易适用于广泛的医学图像表征和分析领域，可以很容易地被研究人员和从业人员集成和利用。

更多阅读