编辑 | ScienceAI
上周,北京智源人工智能钻研院(BAAI)、北京大学和香港中文大学的钻研团队开源了 SegVol 医学通用联系模型。
与过去一些很棒的 Medical SAM 工作不同,SegVol 是第一个能够同时支持 box,point 和 text prompt 从事任意尺寸原分辨率的 3D 体素联系模型。作为一个便捷的通用联系工具,钻研职员将 SegVol 代码和模型开源到 GitHub:BAAI-DCAI/SegVol,欢迎大家应用。
目前开源的模型权重文件包括(1)应用 96k CTs 预训练 2,000 epochs 的 ViT模型,(2)在预训练基础上,应用 6k Masked CTs 在 A100 上训练 30×24×8 个 GPU 小时得到的 SegVol。最新进展请关注 GitHub 仓库的更新,如果有疑惑或建议可以写评论、开 issue 或私信,欢迎大家讨论 。
该钻研以《SegVol: Universal and Interactive Volumetric Medical Image Segmentation》为题,发布在预印平台 arXiv 上。
GitHub 地址: https://github.com/BAAI-DCAI/SegVol
论文链接:https://arxiv.org/abs/2311.13385
摘要
精确的医学图像联系为临床钻研提供了富有意义的结构信息。尽管深度学习在医学图像联系方面已经取得了显著的进展,但仍然缺乏一种能够通用联系各种剖解种别且易于用户交互的基础联系模型。
本文提出一种通用的交互式医学体素联系模型——SegVol。通过在 90k 无标注 CTs 和 6k 联系 CTs 数据上从事训练,该基础模型支持 point, box 和 text prompt,能够对 200 多个剖解种别从事联系。大量的实验证明,SegVol 在多个 benchmark 中表现出色。特别在三个具有挑战性的病变数据集上,SegVol 方法比 nnU-Net 的 Dice 得分高 20% 左右。SegVol 的代码和权重已经在https://github.com/BAAI-DCAI/SegVol 上公开。
核心贡献
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和 SegVol 方法的病变联系性能。这些例子包括肝肿瘤、结肠癌和肺肿瘤。可视化结果显示,与nnU-net产生的结果相比,SegVol重建的这些病变剖解结构更接近于Ground Truth。
(4)消融实验
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个数据集上训练的最终模型从事评估。详细的结果如图3 a和b所示。作为轻量级模型,当只应用一个数据集时,SegVol的性能不是最优的。然而,随着数据量的增加,SegVol的Dice score显著增加,特别是在应用text prompt从事联系的情况下。因为text prompt严重依赖带有语义信息的ground truth mask的数量。
总结
钻研职员提出了SegVol:一个交互式的通用医学体素联系的基础模型。该模型是应用90k无标注数据和25个开源联系数据集训练和评估的。与最强大的传统体素联系方法nnU-net(自动为每个数据集配置参数)不同,SegVol的目的是将各种医学体素联系任务统一到一个单一的架构中。SegVol作为一个通用的联系工具能够对超过200个剖解目标产生准确的联系响应。
此外,与传统方法相比,SegVol具有最先进或接近最先进的体素联系性能,特别是对于病灶目标。尽管具有通用性和精确性,但与其他体素联系方法相比,SegVol保持了轻量级架构。SegVol作为一个开源的基础模型,将很容易适用于广泛的医学图像表征和分析领域,可以很容易地被钻研职员和从业职员集成和利用。