牛津大学提出 MedSAM-2 ,刷新图像分割 SOTA 榜 !

图像处理机器学习关系型数据库

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

在本文中,作者介绍了一种先进的分割模型Medical SAM 2(MedSAM-2),该模型利用SAM 2框架同时解决2D和3D医学图像分割任务。

采用将医学图像视作为视频的理念,MedSAM-2不仅适用于3D医学图像,还解锁了新的One-prompt Segmentation能力,使得用户可以为一个或特定目标目标的某个图像提供提示,然后模型可以自主地后续图像中分割相同类型的目标,无论这些图像的时间关系如何。

作者在多种医学影像模式下评估了MedSAM-2,包括腹部器官、视网膜、脑肿瘤、甲状腺结节和皮肤病变,并与传统的和交互式分割设置下的最先进模型进行了比较。

作者的研究结果表明,MedSAM-2不仅在性能上超越了现有模型,还显示出在各种医学图像分割任务上的高级泛化能力。

作者将在https://github.com/MedicineToken/Medical-SAM2中发布作者的代码。

1 Introduction

大型基础模型由于其强大的零样本学习能力在各种应用领域都展示了巨大的潜力。 Wei 等人(2022年);Kirillov 等人(2023年);Wang 等人的贡献(见下文)。特别是在医学图像分割领域,Kirillov 等人的 Segment Anything Model (SAM) 在零样本图像分割任务中取得了显著的成功。在 SAM 的基础上,提出了增强版的 Meta Segment Anything Model,即 SAM 2,由 Ravi 等人在2024年提出,作为一项既先进又实时适用于静态图像和视频流的物体分割模型。这种统一的方法显著减少了用户的交互时间,提高了视频分割的效果。配备了实时处理功能,擅长处理涉及物体运动和遮挡的复杂场景,SAM 2 准备改变各种领域的物体分割实践。

作为视觉分割的关键和独特组件,医学图像分割需要将图像分割为不同感兴趣的区域(ROI)。这个过程对于诊断和引导手术等应用来说至关重要。近年来,由于自动化分割方法能够同时提高一致性和准确性而引起了注意。

深度学习技术的进步促使神经网络模型在医学图像分割任务中获得成功。此外,这些模型包括传统的卷积神经网络(CNNs)以及较新的视觉 Transformer (ViTs)。 医学图像分割中的一个显著挑战就是模型泛化。特别是,训练在特定目标(如器官或组织)上的模型通常很难适应其他目标。因此,针对每个分割目标的独特模型通常需要使用。以前解决这个问题的方法包括采用交互式分割范式 Wu 等人(2023);Cheng 等人(2023);Ma 和 Wang(2023),例如用户可以点击所需的器官,而模型在不需要针对该特定器官进行训练的情况下预测其分割。然而,这种方法需要用户为每个图像提供提示,这可能是相当费力的。在医学图像分割中还存在另一个独特的挑战,即许多为计算机视觉设计的标准深度学习架构都是针对2D图像的。然而,在医学成像中,数据通常以 CT,MRI 和超声图像的 3D 格式存在。当将通用的计算机视觉模型应用于 3D 医学成像数据时,这种不匹配会产生显著的差距。

在这项工作中,作者证明了最初作为图像和视频分割的通用视觉模型SAM 2,可以独特地适应于解决医学图像分割面临的两项重要挑战。这种适应性导致了医学图像分割的高效模型,称为Medical SAM 2 (MedSAM-2)。具体而言,作者采用将医学图像视为视频的高级概念,并设计了一个独特的模块和训练推理 Pipeline ,使SAM 2可用于2D和3D医学图像分割。因此,该模型不仅在小叶病图像等多选题上实现了优越的性能,还揭示了以前模型难以解决的能力,即One-prompt Segmentation Wu和Xu (2024)。在这种场景下,用户只需提供一个提示,例如在视网膜图像中的视盘。随后,模型可以自主分割后续图像中的视盘,即使这些图像之间没有时间关系。这个困难的设置,几乎没有任何一个现有模型有效应对。MedSAM-2,也就是作者称之为能够处理这种One-prompt Segmentation设置的能力,可以应对这一挑战。尽管One-prompt Segmentation很复杂,但这种方法对临床应用非常有益。当分析多个图像时,临床医生只需要选择第一个作为提示,而在所有后续图像中,所需区域会自动分割,无需进一步输入。这大大降低了临床医生的工作量。此外,MedSAM-2是一个通用模型,能够对任意物体在任何图像中进行分割,这是零样本泛化的一个特征。这意味着它可以应用于无需进行定制适应的视觉新内容。另外,MedSAM-2在用户端提供了很大的灵活性,让用户可以在模型处理数据时随时、任意地调整分割目标。这对医学图像分析中的临床医生来说大大增强了便利性。

作者在15个不同的基准测试上评估了MedSAM-2,涵盖26项独特的任务进行验证。与先前的完全有监督分割模型和基于SAM的交互模型相比,MedSAM-2在所有测试方法上都表现出优越的性能,并实现了2D和3D医学图像分割任务的顶尖结果。具体而言,在某提示分割设置下,MedSAM-2超过了所有之前的少样和单样分割模型,从而展示了其卓越的泛化能力。

综上,作者的贡献如下:

  1. 作者是首先提出基于SAM-2的医学图像分割模型MedSAM-2。
  2. 作者采用了一种新颖的医学影像-视频的理念,这种理念启示作者设计了一种独特的流水线,该流水线可以实现 MedSAM-2 中的一个提示分割能力,这种能力是以前的方法难以实现的。
  3. 作者开发了独特的模块和流水线,其中集成了信心记忆库和加权拾取,以技术上实现这一能力。
  4. 作者在15个不同的基准测试上评估了 MedSAM-2,包括26个不同的任务,在这些任务上,模型表现出色。

2 Method

SAM 2 最近作为一种高性能图像和视频分割的通用模型被引入。通过将医学影像视为视频,并以此设计一套独特的流水线,作者提高了 3D 医学图像分割性能并解锁了一项独特的单提示分割能力。作者将这种创新模型命名为 MedSAM-2。在以下部分中,作者将讨论将 MedSAM-2 应用于 3D 医学影像的动机(参见第 2.1 节),使用 2D 医学影像的合理解释(参见第 2.1 节),以及新模块和流水线(参见第 2.3 节)以及模型架构的简要概述(参见第 2.4 节)。

Applying SAM 2 on 3D medical images

处理3D医学图像的直接方法是将它们处理成一系列的2D切片。然而,相邻切片之间的内在关联可以用于解决单切片2D图像处理中遇到的许多挑战。例如,由于患者或医学设备运动,医学图像中的器官或组织可能呈现模糊的边界。利用相邻切片的上下文信息可以实现即使切片质量受损也能准确分割一个切片。

SAM 2是一个预训练的基础网络,用于视频分割,擅长处理3D视频数据并设计用于跨越视频帧跟踪和分割目标。因此,利用这样一个模型处理3D医学图像很直接。3D医学图像的连续性类似于自然视频——不仅在尺寸上,而且在帧之间的紧密关联。这种相似性使得SAM 2可以在整个3D图像中跟踪并增强目标分割,无需为每个切片提供额外的提示。

From Universal iVOS to One-prompt Segmentation

尽管采用SAM 2进行3D医学成像非常简单,但它在2D图像分割方面的独特优势,超越了其固有的分割能力,并非立即明显。然而,作者发现,当一组包含相同类型器官但内容各异且缺乏时间联系的2D医学图像被视为视频序列时,SAM 2可以解锁一种“一次提示分割”的能力,这种能力其他方法很少实现。

一次提示分割,正如Wu和Xu(2024年)所定义的方法一样,是一种方法,在这种方法中,用户只需要向训练好的模型提供一个未见过的示例。模型则可以有效地执行这个任务,无需重新训练或微调,因为得益于其泛化能力。在本论文中,作者认为将SAM 2的视频分割能力转移到处理2D图像序列是一种实现一次提示分割的有效途径。为了正式描述:

考虑一个涵盖各种不同医学图像分割任务的集合。每个任务都标注为图像和标签对。通常,全监督分割方法通常会学习一个函数来估计从输入图像的分割图。然而,这种函数针对特定任务。

图1:MedSAM-2能力的示意图。当给定一个提示在一个3D切片中的一维时,MedSAM-2可以分割所有后继的时空3D帧。当给定一个提示在一个2D图像中的一维时,MedSAM-2可以使用相同的标准准确分割其他不相关的2D图像(它们没有时间关系),从而体现了一次提示分割能力的产生。

picture.image

在普遍交互式视频目标分割(iVOS)的情况下,其目标是学习一个普遍的函数,其中和分别表示时间顺序的连续视频帧,表示图像包含未见物体,而是目标未见物体的提示集合,提示模型可以在任何包含该目标的上进行泛化。提示集合中的多个提示可能与单个特定图像相关联,而多个图像可能需要一个提示。

另一方面,在一提示分割设置中,方法学习一个适用于任何医学任务的通用函数。在这里,包括一个固定的模板图像和一个相关的提示。如果作者考虑iVOS,包含未见物体的视频作为任务集中的新任务,那么对于iVOS,作者有。比较一次提示分割函数与iVOS中的函数,作者可以明显看出在一提示分割和iVOS下,存在以下两个条件:

  1. iVOS中的图像-标签对并非严格的时间顺序,允许一次提示分割在即使它们缺乏时间关系的情况下分割相同的目标跨不同的图像。
  2. iVOS中的提示局限于一个,应用于处理图像集中的随机图像。因此,当一组包含相同物体的图像被视为视频序列时,从iVOS转移到一次提示分割是轻而易举的。

尽管这两个范式在技术上可能类似,但在各种应用中具有不同的意义。一次提示分割在临床环境中特别易用,因为它只需要提供一个带有提示的样本,使得模型可以适应新的任务,仅需一维前向传播。这种简单性对于没有计算机科学背景的医生来说是一种显著的优势,因为它消除了复杂的训练或微调过程的需要。

Segment Medical Images as Videos

作者认识到视频分割和医学单张图像分割之间的固有相似性,作者将已设计的视频分割模型SAM 2(Self-Attentive Memory Bank Model)用于医学图像分割,通过将医学图像视为视频实现增强记忆的医学图像分割。这种方法不仅提高了3D医学图像分割的性能,还解锁了2D医学图像流的“一次提示”分割能力。通过将作者提出的信心记忆库和加权抓取策略融入其中来实现。

在SAM 2模型中处理3D医学图像与处理视频数据类似,因为三维医学图像中的相邻切片具有强大的时间关联。具体而言,作者使用SAM 2中使用的记忆系统来检索先前的切片及其相应的预测值进行连续切片分割。核心概念是使帧嵌入、后图像编码器与过去的预测和提示图像的条件相关。这些记忆,由基于神经网络的记忆编码器编码,存储在作为缓冲区的记忆库中。在条件过程中,输入图像嵌入通过记忆嵌入增强,主要基于 Transformer 注意力机制。一旦输入嵌入通过记忆注意力进行条件化,它将向前解码器预测分割。分割结果然后添加回记忆库以协助跟踪下一个切片。

作者从不同SAM 2模型处理2D医学图像的方式以便实现 MedSAM-2 的一次提示功能。首先,作者将包含相同器官或组织的医疗图像集归为一组,称为“医学图像流”。虽然这些图像没有时间关联,但作者将其视为视频序列以进行分割。作者的目标是,一旦用户提示此流中的随机图像,MedSAM-2将不仅为该图像预测分割,还会将此预测传播到流中的所有其他图像。为了实现这一点,作者采用了独特的记忆机制。

与SAM 2中使用的顺序输入一进一出队列不同,作者引入了一种“信心优先”记忆库,用于存储模型的确信模板。具体而言,在推理阶段,MedSAM-2将存储在内存库中的最确信的结果(即信心最大的模板)以第一进一出队列的方式存储,同时存储提示模板。作者通过计算模型的预测概率来得出信心。这种方法确保了内存库中的模板是模型最认可的准确样本,从而减少了噪音模板的影响。作者还在添加图像到内存库时实现了图像多样性约束,以确保内存库涵盖各种各样的图像,以更好地匹配传入的输入图像; 多样性是通过图像嵌入的相似性进行评估的。

当将输入图像嵌入与内存银行的信息合并时,作者与SAM 2中的方法不同,作者采用了加权抓取策略,赋予与输入图像更相似的图像更高的权重,以方便传播。在训练阶段,作者使用校准头进行预处理以确保模型在其预测上具有更高的信心和更少的信心,从而增强信心记忆库的有效性。

MedSAM-2 Architecture

Med-SAM2的特点在于其具有以下几个组件:一个图像编码器,可以将输入图像抽象成嵌入;一个记忆编码器,可以将预测的帧嵌入抽象化;以及一个记忆注意力机制,它通过在记忆库中存储的记忆条件输入嵌入。该基本架构类似于SAM2中Ravi等人(2024年)所使用的架构。对于 unfamiliar 与此架构的读者,作者将简要介绍每个组件。

网络中的编码器和解码器类似于SAM Roy等人(2023年)中使用的那些。编码器由一个分层的视觉 Transformer 组成,解码器包括一个轻量级的双向 Transformer ,该 Transformer 集成提示嵌入与图像嵌入。提示嵌入由提示编码器生成,该编码器处理用户的提示,从中提取对应的嵌入。

记忆注意力组件由一系列堆叠的注意力块组成,每个块都包含自注意力块和交叉注意力机制。这种结构将记忆库中存储的帧和目标指针集成在一起。目标指针由一系列向量组成,这些向量表示要被分割目标的语义信息,这些向量是根据每个帧的遮挡解码器输出标记推导出来的。这种架构允许用户在3D图像分割过程中任何帧提示,因为记忆注意力机制包含用户提供的中间提示,以改进分割结果。

3 Experiment

Dataset

作者通过在不同五个独特的医学图像分割数据集上进行实验,使用自动生成的 Mask 提示,开始了作者的研究。这些数据集可分为两类。第一类旨在评估一般的分割性能,作者选择了腹部多器官分割任务。这个任务在领域中具有最大的挑战性,作者使用了BTCV数据集Fang和Yan(2020),这是一个广为人知且公开的基准,包括12个解剖结构。

剩下的四个数据集被选择用于评估模型在不同成像模式下的一般化能力。这些包括在视网膜图像上的黄斑盘和黄斑杯分割,使用REFUGE2数据集Fang等人(2022);在MRI扫描上的脑癌分割,使用BraTs 2021数据集Baid等人(2021);在超声图像上的甲状腺结节分割,使用TNMIX标准,该标准结合了TNSCUI Ma等人(2017)的4554张图像以及DDTI Pedraza等人(2015)的637张图像;以及在皮肤镜图像上的黑色素瘤或痣分割,使用ISIC 2019数据集Milton(2019)。

作者还使用10个额外的2D图像分割任务,采用不同的提示来评价模型的单一提示分割性能,具体包括KiTS23Heller等人(2023)、ATLAS23Quinton等人(2023)、TDSCtds(2020)和WBCZheng等人(2018)数据集采用的_point_提示技术。对于SegRap Astaraki等人(2023)、CrossM23 Cro(2023)和REFUGE Fang等人(2022)数据集,作者使用_BBox_(边界框)提示。最后,_mask_提示被应用于CadVidSet Wang等人(2020)、STAR Hoover等人(2000)和ToothFairy Bolelli(2023)数据集。

Evaluation Metrics

作者使用Intersection over Union (IoU)和Dice Score来评估模型在医学图像分割中的性能。

Intersection over Union (IoU)Intersection over Union (IoU),也称为Jaccard指数,是一种用于评估特定数据集上目标检测器准确性的度量方法。它通过将预测分割和真实值的区域内重叠部分面积除以它们的并集面积来量化两个数据集之间的重叠。IoU的公式如下:

IoU提供了在目标 Level 上清晰的可视化,同时评估预测相对于实际数据的尺寸和位置精度,这对于理解检测模型性能特别有用。

图3:在连续3D医学图像分割中,比较了MedSAM,MedSAM-2和地面实物的性能。

picture.image

Dice ScoreDice Score,或Dice系数的统计工具,用于比较两个样本之间的相似性。由于它在检测被检查目标的尺寸方面的敏感性,因此在医学图像分析中特别常见。Dice Score的计算方法是:

这个分数范围从0到1,1表示模型预测和实际地面实物的预测完全一致。Dice Score对于在目标大小具有显著变化的场景中具有很高的参考价值。

IoU和Dice Score这两项指标都能全面提供模型准确的洞察,特别是在涉及目标大小变化显著的场景中,Dice Score尤其有效。

Hausdorff Distance (HD95) 度量Hausdorff Distance (HD95)是一种用于确定两个点集之间差异程度的度量,通常用于评估图像分割任务中目标边缘的准确性。它对于量化预测分割和实际地面实物品边缘之间的最坏情况距离特别有用。

Hausdorff Distance通过计算一集合中到另一个集合中最近一点的距离的最大值来确定最大距离。在图像分割中,这意味着计算预测边界上一点到地面实物品中最接近一点的距离,反之亦然。Hausdorff Distance的公式如下:

其中和分别表示地面实物品和预测边界上的边界点集合,表示点和之间的欧几里得距离。

尽管Hausdorff Distance提供了严格的可视化度量,但它可能对异常值过于敏感。为了减轻这一缺陷,使用HD95度量,该度量仅考虑距离的95百分位数而不是最大值。这种调整使HD95对异常值不那么敏感,为实际应用提供了更稳健的度量:

在医学图像分析中,这个指标特别重要,因为在精确分割解剖结构方面具有关键性,而异常值可能会扭曲分割性能的评价。

Results

3.3.1 通用医学图像分割性能

在此部分,作者将在广泛认可的社区2D和3D医学图像分割任务中,对所 Proposal 的MedSAM-2模型与现有的最先进(SOTA)医学图像分割方法进行比较。对于3D医学图像,提示以0.25的概率随机提供,而对于2D图像,概率为0.3。3D医学图像分割的初始结果如表1所示,重点关注多器官分割任务。

picture.image

对3D医学图像在BTCV多器官分割数据集上的性能评估:评估通用性能提出的模型与已有最先进分割方法。

表1展示了MedSAM-2在BTCV多器官分割数据集上的性能,该数据集是在广泛认可的社区中广泛接受的。表中的结果表明,MedSAM-2相对于以前的方法取得了显著的进步。值得注意的是,在BTCV数据集上,MedSAM-2在多器官分割方面创造了新的国际领先水平,取得了总体更好的性能。具体而言,MedSAM-2达到的最终Dice分数为88.57%。这一分数不仅远 surpassbaseline的无监督模型SAM-2的零样本,而且也超过了之前完全监督的最新国际领先水平,即MedSegDiff,距离前者达0.70%。此外,在交互式模型中,MedSAM-2保持了领先优势,比之前的交互式模型Med-SA提升了2.78%。值得注意的是,所有这些竞争交互模型都需要在每一帧上提供提示,而MedSAM-2则可以在用户提示大大减少的情况下实现更好的结果。

在2D医学图像上,作者还评估了MedSAM-2与针对不同图像模式特定任务的最新(SOTA)分割方法,详细见表2。对于视网膜杯分割,与Yu等人(2019年)的ResUnet进行比较,与Wang等人(2019b)的BEAL进行比较;对于脑肿瘤分割,与Wang等人(2021b)的TransBTS和Wang等人(2021b)的SwinBTS进行比较;对于甲状腺结节分割,与Gong等人(2021年)的MTSeg和Chu等人(2021年)的UltraUNet进行比较;对于皮肤病变分割,与Wu等人(2022a)的FAT-Net和Wang等人(2021a)的BAT进行比较。此外,作者还将MedSAM-2与需要用户提示每个图像的交互式模型进行了比较。

picture.image

表中的结果说明了,MedSAM-2在五项不同任务上都超过了所有与它比较的方法,展示了其在不同医疗分割任务和图像模式上的优越泛化能力。具体来说,MedSAM-2通过Dice分数实现了2.0%的视网膜杯改进,1.6%的脑肿瘤改进和2.8%的甲状腺结节改进。即使在与交互式模型(需要对每个图像进行提示)进行比较时,MedSAM-2仍保持领先地位,证实了所提出的信心记忆库在提高性能方面的有效性和优势。

3.3.2 在不同提示下的单提示分割性能

作者进一步评估了MedSAM-2在单提示分割设置下的性能,通过将其与使用不同提示的各种少/单一次学习基础进行比较。这个实验评估了MedSAM-2在一个在只有一个提示,且后续图像之间没有明显联系的挑战场景中的性能。

作者将MedSAM-2与一些少/单次学习模型进行了比较:PANetWang等人(2019年)、ALPNetOuyang等人(2020年)、SENetRoy等人(2020年)和UniverSegButoi等人(2023年),所有这些模型在一致性方面都提供了相同的模板。为了进行公平的比较,所有少/单次学习方法在测试时都只使用了一个模板和一个提示。此外,作者还与单次学习模型进行了比较:DATZhao等人(2019年)、ProbOne Ding等人(2021年)、HyperSegNasPeng等人(2022年)和One-prompt Wu和Xu(2024年)。这些模型使用与作者的相同的数据集进行训练,并由于它们不能除One-prompt Wu和Xu(2024年)外使用稀疏提示,因此它们在训练和测试中都得到了「提示」作为分割标签。

图4展示了每个方法在每个任务上的平均Dice分数。值得注意的是,与作者的方法相比,比较的少/一次性的模型需要提示(mask),这使他们相对于作者的方法具有潜在的优势。然而,MedSAM-2始终在这些模型上明显优于它们,在各种任务上表现出强大的泛化能力。即使与在64个不同的医学数据集(远多于作者)上进行高度训练的One-prompt Segmentation模型Wu和Xu(2024)进行比较,MedSAM-2仍然优于它,在10个任务中的9个任务上表现出色。在所有方法都提供提示(如图4中所示)的场景下,MedSAM-2的优势更为明显,通常比第二个最佳模型平均优势3.1%,这是所有提示设置中最大的优势。

picture.image

4 Conclusion

总之,MedSAM-2在医学图像分割领域代表了重要的技术进步。通过将SAM 2框架集成并采用类似于医学视频的处理方法,MedSAM-2成功地扩展了高级分割技术在2D和3D医学图像上的应用。One-prompt分割的新兴能力显著降低了持续用户交互的需求,有助于实现更高效、更用户友好的工作流程,这在临床环境下尤为有益。作者对各种成像模式进行全面评估的结果表明,MedSAM-2始终优于当前最先进的技术,具有卓越的泛化能力和在多个分割任务上的强大性能。

这项研究的深远影响表明,在医学影像中使用视频分割原则可以显著提高诊断工具的精度和效用。展望未来,作者预计在实时处理和集成更多种类的数据类型等方面,作者将对方法进行进一步的改进。未来的工作也将探索MedSAM-2在其他非医学领域的复杂分割场景下的应用潜力,拓宽了这项强大技术的应用范围。

参考

[1].Medical SAM 2: Segment medical images as video via Segment Anything Model 2.

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
CV 技术在视频创作中的应用
本次演讲将介绍在拍摄、编辑等场景,我们如何利用 AI 技术赋能创作者;以及基于这些场景,字节跳动积累的领先技术能力。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论