医学图像和视频中的任何内容分割:基准测试与部署
目录
关键字:SAM、分割一切基础模型、医学图像、视频、多模态
最近医学分割模型发展迅速,基于SAM的医学图像处理得到了进一步的发展。为了追踪医学图像处理的最新进展,我们跟踪了最新的SAM2模型,以及其用于医学图像分割的重要应用。
摘要:
最近在分割基础模型方面的进展,使得在广泛的自然图像和视频上能够进行准确和高效的分割,但它们对医学数据的实用性仍然不清楚。
在这项工作中,作者首先对Segment Anything Model 2(SAM2)在11种医学图像模式和视频上的全面基准测试,并与SAM1和MedSAM进行比较,指出其优势和不足。
然后,作者开发了一个迁移学习流程,并展示了通过微调可以快速将SAM2适应于医学领域。此外,我们将SAM2实现为一个3D切片插件和Gradio API,用于高效的3D图像和视频分割。
代码已在https://github.com/bowang-lab/MedSAM公开提供。
一、引言
分割是医学图像分析中的一项重要基础任务,它是许多下游任务的前提条件,比如早期癌症检测[7]和疾病诊断[49]。
在过去的十年中,从专家分割模型[13,43]到通用分割模型或分割基础模型[25][57]的方法论趋势非常明显。
- 专家模型,通常针对特定的解剖结构[32,34]、疾病[30,18,5]或医学成像模式[41,4,20]开发,它们只能对与训练集相同领域的数据进行分割。
- 相比之下,基础模型是在大规模和多样化的数据集上训练的,它们具有强大的能力,能够分割各种对象、图像类型和场景,甚至在未见过的领域的图像上也具有零样本泛化能力[6,50]。
Segment Anything Model(SAM)[25]是图像分割中的第一个基础模型,它在256个GPU上用1.1亿个掩码注释从1100万张图像中训练而来。
尽管SAM支持常用的交互提示(例如,点、边界框和掩码)以实现灵活的分割,并在广泛的图像和下游任务中采用[40],但由于训练数据(SAM-1B)只包含自然图像,它未能对医学图像进行分割[36][19]。
医学图像与自然图像有很大的领域差异,存在独特的挑战,比如低分辨率和有限的图像质量[46]。尽管如此,通过迁移学习[29],SAM可以轻松扩展到医学领域。
MedSAM[29]已经证明,通过对大规模医学图像进行微调,SAM可以实现卓越的2D医学图像分割性能,但其在分割3D医学图像和视频方面的能力仍然有限,因为在实际中通过逐图像的分割流水线为每个图像生成提示是低效的。许多最近的研究已经探索了将SAM扩展到体积图像分割,例如向图像编码器添加3D适配器[17,8],直接将图像编码器增强到3D[48,12],使用带有门控网络的混合专家(MoE)选择任务特定的微调模型[47],以及文本引导的3D分割[54]。
图 1. 数据集和评估协议。我们在各种2D和3D医学图像和视频上评估SAM2。2D图像和边界框提示直接传递给SAM2以生成分割结果。3D图像和视频分别在中间切片和第一帧上初始化一个边界框提示来获取2D掩码。然后,模型将2D掩码传播到剩余的切片/帧。
1.1 SAM2 在医学图像和视频中的应用
最近,SAM2 [10] 在不牺牲图像分割能力(方法)的情况下,进一步增强了SAM的可提示视频分割功能。该模型在一个前所未有的数据集上进行了训练,包含50.9K个视频,并在已建立的视频目标分割基准测试中超越了现有工作。
值得注意的是,它还在许多其他视频和图像分割基准测试中展现了强大的零样本性能,这些基准测试涵盖了各种分布。然而,SAM2在医学图像上的表现,尤其是3D医学图像和视频上的表现尚不清楚,因为在SAM2中只评估了一个医学数据集(内窥镜视频中的器械分割[3])。
此外,官方界面仅支持评估短视频且数据格式有限,这不能用于医学专业人士用自己的医学数据测试SAM2。
在这项工作中,我们通过在十种医学模式上对SAM2进行全面评估来解决上述问题,包括各种2D和3D图像和视频。我们还比较了SAM2与SAM1和MedSAM,以全面了解它们的优势和劣势。与同期工作[10,56,44]相比,我们开发了一个迁移学习流程,通过微调快速将SAM2适应于医学图像分割,并将SAM2进一步整合到3D Slicer插件[14]中,使用户能够轻松地使用SAM2标注官方SAM2界面不支持的各种3D医学数据(例如,CT、MR和PET)。我们还实现了Gradio [1]界面,以支持高效的医学视频分割。
二.结果
2.1 数据集和评估协议
基准测试数据集涵盖了11种常用的医学图像模式,包括
- 计算机断层扫描(CT)、
- 磁共振成像(MRI)、
- 正电子发射断层扫描(PET)、
- 超声波(US)、
- 内窥镜、
- 眼底摄影、
- 皮肤镜、
- 乳腺摄影、
- 光学显微镜
- 光学相干断层扫描(OCT)
由于SAM2是一个通用的图像和视频分割模型,我们评估了它在2D图像、3D图像和视频上的能力。2D数据集包含了所有模式,而3D数据集包括CT、MR和PET扫描。视频数据集由超声波和内窥镜视频组成。
除了最近的SAM2 [39],我们还将其与它的前身SAM [25]和通用医学图像分割模型(MedSAM [29])进行了比较,以便全面了解它们的性能。SAM和SAM2分别有三种和四种不同的模型尺寸,所有这些模型都在实验中进行了评估。与MedSAM类似,我们仍然提倡使用边界框提示,因为它不仅高效且无需尝试和错误,而且与点提示相比具有更少的歧义。对于2D图像,我们将图像和边界框提示传递给模型以获得相应的掩码。
我们将第一个分割任何内容的模型[25]称为SAM1,目的是将其与最近的SAM2 [39]区分开来。
对于3D图像,我们在中间切片初始化边界框提示,然后调用模型生成2D掩码。由于SAM和MedSAM只支持2D图像分割,剩余部分从中间向底部和顶部逐层分割。当前切片上的分割掩码被用来模拟下一片的边界框提示,以顺序方式进行。SAM2支持视频分割,3D医学图像可以被视为视频中的一帧[55]。因此,中间切片的分割掩码通过SAM2中的视频分割功能传播到剩余的图像。我们选择中间切片作为3D图像分割的起始帧,因为在传统的轴向视图中,它通常在所有切片中具有最大的物体尺寸。对于视频分割,我们在第一帧上初始化边界框提示以获得初始掩码,然后将掩码传递给视频分割模型。只有SAM2评估了视频分割,因为不同帧之间的物体位置变化很大,简单地逐帧应用SAM或MedSAM是行不通的。
2.2 二维图像分割的评估结果
图2和表1展示了在11种模式下的二维图像分割结果。与SAM1相比,SAM2在磁共振成像(MRI)、皮肤镜和光学显微镜图像中获得了更高的DSC分数,但在PET和OCT图像中的分数较低。它们在CT、X射线、超声波、内窥镜、眼底和乳腺摄影图像中的性能是可比的。
相比之下,MedSAM在11种模式中的9种中一致性地超越了SAM1和SAM2,除了PET和光学显微镜图像,因为其训练集不包含这些类型的图像。
表1. 八个SAM模型变体在11种模式的二维图像上的定量分割结果(平均DSC分数)。SAM1-Base: 93.7M; SAM1-Large: 312.3M; SAM1-Huge: 641.1M; SAM2-Tiny: 38.9M; SAM2-Small: 46.0M; SAM2-Base: 80.8M; SAM2-Large: 224.4M; MedSAM: 93.7M。
三 讨论
模型变体:SAM1与SAM2
当SAM2问世时,一个直接的问题出现了:“SAM2在医学图像分割上是否超越了SAM1?”对于2D MRI、皮肤镜和光学显微镜图像,答案可能是“是”。
对于3D CT和MRI扫描,SAM2-B也取得了最高分,这大幅度超越了所有SAM1模型。然而,对于2D OCT和PET图像,答案也可能是“否”,因为SAM1的表现更好。
SAM2的改进可以归因于它拥有更先进的网络架构,并且是在大规模图像和视频数据集上训练的。在某些模式上性能下降,可能是由于与SAM1相比模型尺寸较小。
尽管如此,考虑到复杂的模型开发过程,理解这种性能差距的原因需要访问模型训练代码并进行严格的消融研究。
SAM2在医学图像和视频标注中的部署。a),用于3D医学图像分割的Slicer插件;b),用于视频分割的Gradio API。
四 局限性和未来的工作
基准研究已经涵盖了常见的医学图像模式,但可以通过包括更多的3D模式,如3D超声和OCT图像,来进一步增强。
还有一些值得进一步探索的未解决问题。
- 首先,SAM2的视频分割能力在医学领域有广泛的应用,但当前模型经常无法分割没有清晰边界的目标。这个问题可以通过在医学数据集上进行迁移学习来解决。尽管微调视频分割模型要复杂得多,但我们的2D图像迁移学习流程为进一步开发提供了良好的基础。
- 其次,SAM2只支持点、框和掩码提示。相比之下,文本提示为复杂结构分割提供了更大的灵活性[53,54,52]。在SAM2中实现自然语言处理能力可以弥合复杂医学术语和模型理解之间的差距,促进更直观高效的用户体验。
- 另一个关键方向是使SAM2更加轻量化。SAM2-T比SAM1-B小得多,但仍然需要大量的GPU RAM来进行长视频分割。有必要进一步减小模型大小并提高推理效率,以在更广泛的临床环境中部署,而不牺牲性能。
总之,我们对SAM2在各种医学图像分割模式上进行了全面评估。不同模型大小和分割任务的性能差异显著,表明新的或更大的模型在分割准确性方面总是更好的。
此外,我们开发了一个微调流程,并证明了其在提高3D医学图像分割性能方面的有效性。我们还提供了3D Slicer插件和Gradio API,以促进SAM2在医学图像和视频分割中的部署。
尽管这类工程实现很少出现在基于深度学习的医学图像分割论文中,但我们认为它们对于临床实践的采用是重要的。作者已经公开了所有代码,以供进一步开发。
五、方法
5.1数据来源和预处理
所有评估的图像和视频均来自公共数据集[9,32,33,38,16,15,45,31,37,21,26,22,51,24,2],这些数据集用于2024年CVPR医学图像分割笔记本电脑挑战赛的验证集。它们都没有在MedSAM训练集中使用。
CT图像根据典型的窗口水平和窗口宽度进行强度截断预处理。MR和PET图像的强度被剪辑到非零区域强度的0.5和99.5百分位数。然后,通过最小-最大归一化将强度值归一化到[0, 255]。对于其余模式,强度不变。最后,所有图像都转换为npz格式,以便批量推理。
SAM1和SAM2在方法上的主要区别
SAM2[39]是SAM1[25]在视频分割上的自然延伸,具有统一的图像和视频分割任务框架。技术上,有两个主要的方法改进:
- 首先,视觉变换器(ViT)[11]被Hiera[42]替换,Hiera可以提取多尺度特征以产生高分辨率的分割细节。
- 对于视频的交互式分割,记忆注意力模块用于将当前帧的预测条件化到当前时间点的信息。变换器架构采用当前帧的图像和提示特征,以及之前帧的特征和预测的记忆。通过为记忆注意力模块配备记忆库,SAM2有效地扩展了SAM的提示编码器、图像编码器和掩码解码器模型架构,利用时间上下文实现视频帧间稳健一致的分割。
5.2 微调协议
我们在MICCAI FLARE22腹部CT训练数据集[33]上微调了预训练的SAM2-Tiny模型,数据分割比例为80%用于训练,20%用于验证。在微调过程中,提示编码器被冻结,因为它是领域不可知的,而图像编码器和掩码解码器被更新,允许模型适应CT扫描的具体特征。
输入图像被调整到[1024, 1024],并使用z分数归一化。对于这样的CT数据集,在将预处理的CT扫描切片图像输入模型之前,将其复制以创建3通道图像。从[1024, 1024]比例的真值掩码中获取边界框提示,并应用边界框坐标的随机扰动,最大坐标偏移为5像素,以提高模型的鲁棒性。真值掩码被调整到[256, 256]以匹配掩码解码器的输出尺寸。我们使用AdamW优化器[27]进行微调,学习率为6e-5,批量大小为16。微调进行了1000个周期,当训练损失稳定时手动提前停止。我们使用Dice损失和交叉熵损失的无权重和作为损失函数,因为这种复合损失在各种分割任务中被证明是稳健的[28]。
5.3 评估指标
我们遵循Metrics Reloaded[35]中的建议,使用Dice相似系数(DSC)和归一化表面距离(NSD)分别评估区域和边界重叠比率。
总结
Meta AI将SAM2公开提供给社区。感谢所有数据集的所有者公开提供宝贵的数据。论文解读感谢Guoqing同学。