爱笑的小姐姐 · 2023年09月14日

上海AI实验室通用视觉团队 | SAM-Med2D:高精度医学图像分割一切

image.png

由于医学图像和自然图像之间存在较大差异,以及缺少大规模医学图像基准数据集,这是导致AI在医学领域进展缓慢的原因之一。构建大规模基准数据集和可靠的基线模型,能够推动AI在医疗领域的快速发展,加速医疗向更通用的方向转变。欢迎感兴趣的读者加入群聊与我们讨论!

论文: 

https://arxiv.org/abs/2308.16184

开源代码: 

https://github.com/OpenGVLab/SAM-Med2D

Gradio demo: 

https://openxlab.org.cn/apps/detail/litianbin/SAM-Med2D

文章和代码公开3天内,获得多家媒体报道,在推特阅读量突破十万。

image.png

image.png

image.png

image.png

image.png

image.png

自然图像VS医学图像

众所周知,ImageNet的出现极大推动了基于深度学习的人工智能的发展。它为计算机视觉领域提供了一个大规模的数据集和强大的基线模型,使得研究者们能够在自然图像分类、分割、检测等任务上取得了突破性的成果。Segment Anything Model (SAM)代表了自然图像分割领域最先进的研究进展。然而,直接将预训练的SAM应用于医学图像分割,效果并不理想。这种限制主要源于自然图像和医学图像之间存在显著的领域差异。SAM的训练数据(图1a.上)完全来自于高清自然图像,导致其缺乏医疗领域的特定知识。

医学图像和自然图像存在显著差异本质就是成像方式完全不同:自然图像是以自然光为光源,通过手机、相机等设备转换为RGB格式的图像,像素值范围也是0-255之间;医学图像是通过特定的协议和扫描仪收集的,其展现方式因其特定的临床目的而各异,包括电子、激光、X射线、超声波等,成像都不是RGB图,其像素值范围从几千到几十万不等(见图1a右侧柱状图)。

image.png
图1a. 自然图像(上)与医学图像(下)对比

另外,由于标注成本高且标注质量参差不齐,导致医学图像在不同医院和临床研究中的质量存在很大差异。上述挑战导致医学图像和自然图像在数据量、数据质量上存在显著差异。图1b比较了公开的自然图像数据集和医学图像数据集的数据量规模差异。

image.png
图1b. 医学图像和自然图像数据量规模差异

基于以上问题,医学领域缺乏大规模基准数据集和可靠的基线模型的支持。我们的目标是填补这一空白,推动AI在医疗领域的快速发展,加速医疗向更通用的方向转变。为了实现这一目标,我们研发了SAM-Med2D。首先,团队收集了迄今为止最大最全的医学图像分割数据集,它由 4.6M 张图片和 19.7M masks组成,涵盖了当前公开可用数据集中几乎所有对象类型,弥补了SAM缺乏医学领域知识的问题。在此数据基础上,作者团队全面微调了SAM:在同等分辨率时,FT-SAM在Bbox提示模式下实现了11.93%的提升;SAM-Med2D实现了17.67%的提升。在单点提示模式下,SAM-Med2D相较于SAM表现出了压倒性的优势 (18.94% vs. 70.01%)。在单点提示模式下,SAM-Med2D表现出了压倒性的优势 (18.94% vs. 70.01%)。

SAM-Med2D三大亮点

1.最大规模的医学图像分割数据集:作者团队收集并整理了一个庞大而全面的医学图像数据集,涵盖了多种临床分割任务和图像模态。这使得SAM-Med2D在训练过程中能够获得更准确和具有代表性的医学图像信息,弥补了 SAM 在医学领域数据不足的问题。

2.更全面的微调方案:相对于现有的医学SAM方法,作者对SAM的三个重要组成部分都进行了微调,使SAM-Med2D能够更好地适应医学图像的特殊特征和变化,提高分割结果的准确性和稳定性。图2对比了最近基于SAM的微调方法。

image.png
图2: SAM-Med2D是一种全面的微调方法,支持对医学图像进行多种提示来生成mask

3.优异的医学图像分割性能:对SAM-Med2D从不同维度进行了全面的性能评估。结果表明,SAM-Med2D在各种医学图像分割任务中表现出色,在同等分辨率时,FT-SAM在Bbox提示模式下实现了11.93%的提升,而SAM-Med2D实现了17.67%的提升。在单点提示模式下,SAM-Med2D表现出了压倒性的优势 (18.94% vs. 70.01%)。

迄今为止最大最全的医学图像数据集

01 数据准备

大模型成功的经验告诉我们,数据量对于模型的学习至关重要。通过学习更多的数据,模型可以获得更加丰富的领域知识,并更好地适应不同的应用场景。尽管SAM已经在超过1B的masks上进行了训练,但其训练集与医学领域的数据相比存在较大的差距,这也是SAM在医学领域表现不理想的原因之一。

image.png

图3:本研究所使用的数据集概述。(a) 总共 31 个主要器官及其相应的解剖结构,* 表示数据集中存在病变标签。(b) 呈现了模态的分布及其在数据集中的相应比例 (以对数缩放)。(c) 按解剖结构分类的图像和masks的数量,以及包含数据集的总数

为此,作者团队收集整理了迄今为止最大的医学图像分割数据集,它由4.6M张图片和19.7M masks组成。图3(b)显示了该数据集10种不同的成像模式和对应的数据占比。根据解剖学结构和是否包含病灶将数据集分为5个大类,即Head and Neck、Thorax、Abdomen、Pelciv、Contains lesions,图3(c)显示了这些类别的图片和masks数量。我们从这些数据集中的271种标签中整理和汇总出31种主要的器官(图3 (a)),这涵盖了当前公开可用数据集中几乎所有对象类型,弥补了SAM缺乏医学领域知识的问题。此外,引入了9个MICCAI 2023数据集(包含约0.52M图片和1.31M masks)验证模型的泛化性。我们相信,通过更加全面、多样化的训练数据,SAM将更好地适应医学图像领域的复杂性和特殊性,为医疗健康领域的应用提供更加准确和可靠的技术支持。这也将为医学图像分割领域的研究和发展带来新的机遇和挑战。

image.png
图4: SAM-Med2D概览。我们冻结图像编码器,并将可学习的适配器层合并到每个 Transformer 块中,以获取医学领域的特定领域知识。我们使用Point、Bbox 和Mask信息微调提示编码器,同时通过交互式训练更新解码器的参数

02 从SAM到SAM-Med2D

在这项工作中,作者对SAM进行了微调,微调后的方法称为SAM-Med2D,它可以有效地扩展到医学图像领域。如图4所示,SAM-Med2D由三个主要组件组成:image encoder、prompt encoder和mask decoder。该框架允许在不同提示下为同一图像生成不同的mask。对于image encoder,在微调时冻结了原始图像编码器中的所有参数,并为每个Transformer块部署了一个Adapter层。对于prompt encoder,考虑三种提示模式进行微调 (Point、Bbox和Mask)。与MedSAM,MSA等基于SAM的微调方法相比保留了更完整的提示功能,增强了其在医学成像领域的适用性。对于mask decoder,没有对其结构做任何更改,在训练期间保持每一次迭代都更新其参数。为了使模型具有歧义感知能力,每个prompt会同时预测多个mask(默认为三个)。

与SAM和其他交互式分割方法相同,作者模拟交互式分割的方式对SAM-Med2D进行训练。每个批次的数据需要迭代9次,在第一次迭代时等概率选择一个前景点或Bbox作为稀疏提示,前景点从GT中随机样点,Bbox为GT的最大外接矩形框,对每个坐标中进行最大值为5个像素的偏移。需要注意的是,除第一次迭代时模型同时更新Adapter层、prompt encoder和mask decoder的参数外,后续的迭代仅更新mask decoder的参数。从第二次迭代开始,将上一次迭代生成的mask和GT之间的误差区域随机选择1,3,5或9个点作为后续的稀疏提示,并将上一次迭代生成的低分辨率特征图映射到[0,1]的范围内作为当前迭代的密集提示。同时,在最后一次迭代和任意一次中间迭代只输入密集提示以鼓励模型从所提供的Mask中受益。

03 评估SAM-Med2D

进行全面的性能评估能助力科研社区更深入地理解影响算法感知医疗对象能力的因素,进而改进方法并提升其在实际应用中的效能。然而,以往的评估工作受制于数据量有限以及缺乏通用医疗图像分割方法的瓶颈,导致评估几乎都局限在小规模数据集和有限的类别上,无法充分揭示算法的优势和适用范围。为填补这一研究空白,作者全方位、多维度地对SAM-Med2D进行了评估,为未来的研究者提供交互式分割方法的基准标准。

在模型方面,选择了SAM作为基础模型,并从其交互方式中选择了两种直观的提示模式(Bbox和Point)对SAM-Med2D进行评估。这样的选择是因为Bbox和Points是常用的交互方式,同时可以作为简单高效的标注手段。通过评估SAM-Med2D在这两种提示模式下的性能,我们可以深入研究它们在医学图像分割任务中的优势和不足。

在数据方面,作者评估了SAM-Med2D在MRI、CT、超声等10种不同模态医学图像上的表现。这样的综合评估将使研究者能够了解SAM-Med2D在特定模态下的优势和挑战,并揭示其在多模态图像中的潜在应用能力。然后,考虑到不同解剖学结构和器官具有各自独特的形态、特征和变化模式,作者对SAM-Med2D在4种解剖学结构和31种主要器官上的表现进行了评估。这样的评估有助于深入了解SAM-Med2D在不同场景下的性能差异,并能够有针对性地改进方法以应对特定结构和器官的挑战。最后,作者也高度重视SAM-Med2D的泛化能力,在9个MICCAI 2023数据集上进行了测试。这些数据集代表了来自不同来源、不同机构或不同设备采集的医学图像,具有多样性。通过在这些数据集上评估SAM-Med2D,能够验证其在泛化到新数据上的能力。

实验结果

01 定量评估

image.png
图7: FT-SAM 和 SAM-Med2D 在 31 个器官中的分割性能比较
![Image]image.png
图8: 在 9 个 MICCAI2023 数据集上进行泛化验证,其中“ * ”表示不加adapter层参数的 SAM-Med2D

02 定性评估

image.png
图9: 定性比较SAM-Med2D和SAM。前三行描述了不同模态的分割结果,而后三行描述了不同解剖结构的分割结果

image.png
图10: 单个图像多目标分割结果融合, 仅可视化 1 Point提示的结果。右边为SAM的分割结果

总结

SAM-Med2D通过在大规模医学图像数据集上微调来适应各种医学图像场景。与预训练的SAM相比取得了令人满意的性能改进和泛化能力。在同等分辨率时,FT-SAM在Bbox提示模式下实现了11.93%的提升,而SAM-Med2D实现了17.67%的提升。在单点提示模式下,SAM-Med2D表现出了压倒性的优势 (18.94% vs. 70.01%)。

作者还对不同模型在解剖学结构,不同成像模式和主要器官维度进行了综合评估,SAM-Med2D在各方面优于预训练的SAM和FT-SAM。此外,在9个MICCAI 2023数据集的泛化实验证明了在大规模数据集上预训练的模型具有强大的领域可迁移性。尤其是使用点提示模式时,SAM-Med2D可以更快地生成所需的mask,甚至优于其他方法使用 Bbox 提示模式的结果。未来作者团队将继续扩大数据规模、增加提示策略使SAM-Med2D具备医疗领域的“everything”能力。同时,SAM-Med2D的代码和预训练模型将开放给社区使用,希望这项工作能让医学计算机视觉领域的研究人员受益。

作者:通用医疗GMAI
文章来源:书生 OpenGVLab

推荐阅读

更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
18756
内容数
1315
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息