医学图像和自然图像之间存在显著的差异,而且医学图像领域缺乏大规模的基准数据集,这是导致AI在医学领域进展缓慢的重要原因之一。通过构建大规模的基准数据集和可靠的基线模型,我们可以推动AI在医疗领域的快速发展,并加速医疗向更通用方向转变的进程。如果您对此话题感兴趣,欢迎加入群聊与我们一起探讨!(二维码见文末) 论文:https://arxiv.org/abs/2310.15161 开源代码:https://github.com/uni-medical/sam-med3d
最近提出的视觉基础模型 "Segment Anything Model"(SAM)是一种使用超过 10 亿个掩码(mask)进行训练的 ViT 模型,它在多个领域都展现出出色的零样本分割性能。SAM的兴起可以推动数据标注快速迭代,并为3D 医学图像分析的发展提供新的机遇。然而,多项研究表明,由于医学图像知识的严重不足,将 SAM 直接应用于医学领域的有效性有限。解决这个问题的一种直接的方法是:将医学知识融入到 SAM 中。比如,MedSAM 是一种典型示例,它通过使用110万个掩码(mask)对SAM 的解码器(Mask Decoder)进行微调,从而使 SAM 能够通过边界框(Bounding Box)作为提示来更好地分割医学影像;SAM-Med2D 则引入了适配器(Adapter)和约2000万个掩码(mask)对 SAM 进行了充分微调,从而在医学图像分割中表现出了卓越的性能。 然而,这些方法必须采用逐切片(slice)的方法来处理三维医学图像,也即,将三维数据从某个维度分解为二维切片,然后独立处理每个切片,最后将二维分割结果汇总为三维分割结果。这种方法忽略了切片之间的三维空间信息,因此在三维医学影像上表现不佳,这一问题可以从图1中的结果看出。
图1:SAM 相关的模型在三维医学图像数据上的表现,SAM 和 SAM-Med2D 在空间上都出现了断层的现象,而 SAM-Med3D 在空间上具有更好的连贯性。
除了将 SAM 直接应用于三维数据,一些研究人员希望通过引入二维到三维的适配器(Adapter)来捕捉三维空间信息。如图2所示,这些方法通常在保持编码器(Image Encoder)不变的同时引入了三维适配器(Adapter),以使模型能够从三维图像中学习到三维空间信息。然而,这些方法存在两个主要限制:
因此,解决这些限制将需要更大规模的数据集、更广泛的目标类型,并且可能涉及新的架构设计来更好地应对三维医学图像的挑战。
图2:和现有其他方法的对比,SAM-Med3D 使用了更多的训练数据,包含更多的类别,数据具有更强的多样性,而且其模型具备完整的三维结构。
首先,作者进行了三维医学图像数据集的广泛收集和标准化工作,创建了迄今为止规模最大的三维医学图像分割数据集。该数据集包含了 2.1 万个三维医学图像和 13.1 万个三维掩码(mask),共涵盖了 247 个不同的类别。这些数据来源于多个公开和私有的三维医学图像数据集。从图3可以清晰地看出,这一数据集的规模远远超过了现有最大的三维医学图像分割数据集,如 TotalSegmentator 和 BraTS21,其规模扩大了 10 倍以上。 这一数据集的创建为进一步推动三维医学图像分割研究提供了丰富的资源和机会,有望帮助研究人员突破现有技术的瓶颈。
图3:(a)所有训练数据的类别统计词云图,共有 247 个类别。(b)不同三维医学图像数据集的图像(image)和掩码(mask)数量比较,作者收集的数据包含2.1万张三维图像和相应的13.1万个三维掩码,而 AMOS 和 TotalSegmentator 的图像数量不足2千,具有4种不同模态的 BraTS21 图像数量也不足1万。
作者提出了 SAM-Med3D,这是一个所有参数均可学习的三维架构模型(无冻结的参数),如图3所示。此外,作者对 SAM-Med3D 的性能进行了全面评估。首先,作者使用了 15 个公共数据集来比较 SAM、SAM-Med2D 和 SAM-Med3D。然后,他们从解剖结构、模态和类别等不同角度进行深入分析,多维度评估了这些模型在三维医学影像分割中的性能。此外,作者还验证了 SAM-Med3D 的迁移能力:将其编码器用作预训练模型,在多个全监督分割任务中进行了验证。
综合全面的评估结果,SAM-Med3D 具有以下两个主要优势:
这些优势使 SAM-Med3D 有望成为一个在三维医学图像分割领域实现重大进展的模型,有望为医学专业人员和研究人员提供更强大的工具来处理三维医学图像。
图4:具有完整三维结构的 SAM-Med3D。
性能评估对于深入了解算法的能力以及提高算法的有效性和适用范围至关重要。在医学影像领域,SAM 相关模型的评估主要集中在二维医学图像上,如 X 射线或内窥镜图像,而在三维分割任务的评估方面存在明显的不足。考虑到在医学图像分析中,三维分割任务具有至关重要的作用,因此作者对 SAM、SAM-Med2D(用于医学图像的最先进的微调版SAM)以及所提出的 SAM-Med3D 进行了全面评估,旨在为三维医学图像上的可提示分割任务设定一个基准。 在数据方面,如上所述,作者基于 15 个不同的公共数据集构建了一个具有代表性的三维医学图像验证集。该数据集涵盖了各种不同的目标和模态。这项评估将有助于我们从多个维度(包括模态、解剖结构、器官和病变)研究SAM-Med3D 在三维医学图像分割任务中的性能。SAM-Med3D 的评估包括以下几个关键点:
这些评估角度提供了对模型全面的观察,有助于我们了解 SAM-Med3D 在三维医学图像分割任务中的性能、效率和潜力。
图5:SAM-Med3D在使用更少点击次数的情况下,获得了更好的性能。N表示待分割目标包含的切片(slice)数目,通常10 ≤ N ≤ 200。T_{inf}为N =100时所需的推理时间 (Inference time) 。
图6:从解剖结构和病变角度进行比较。A&T 表示腹部和胸部。SAM-Med3D 只需10个提示点(最后一行)即可取得比 SAM 和 SAM-Med2D 更好的性能,而后两者往往需要上百个提示点。
图7:三张图展示了不同模型在不同模态下的性能对比,其中SAM-Med3D在所有模态下均展现出优异性能。需要指出的是,SAM 和 SAM-Med2D 使用的提示点是基于切片的,而 SAM-Med3D使用的则是基于体素的,导致后者所需提示点远少于前两者。即使SAM-Med3D没有使用超声(US)图像训练,其性能仍与 SAM-Me相当。
图8:SAM-Med3D与性能最好的二维微调SAM模型SAM-Med2D在34个主要器官和5种病变上的Dice系数比较。∗和∗∗分别代表可见病灶和未见病灶。
作者将 SAM-Med3D 预训练的 ViT 图像编码器迁移到 UNETR 中进行使用,发现能够获得效果上的提升,证明了作者提出的 SAM-Med3D 具有迁移能力,这将能够对三维医学图像领域的发展提供帮助。
图9:全监督三维医学图像分割的可迁移性评估。作者利用SAM-Med3D的编码器作为预训练模型,在下游选择UNETR进行微调以评估预训练是否有效。
图10:在不同的解剖结构中,针对不同数量的点,对SAM、SAM-Med2D和SAM-Med3D进行可视化。作者同时展示了轴切片和冠状切片/矢状切片来全面说明三维结果。其中“Abd&Tho”表示腹部和胸部。
图11:在各种模式下,针对不同的点数,对SAM、SAM-Med2D和SAM-Med3D进行可视化。作者同时展示了轴切片和冠状/矢状切片来全面说明三维结果。
在这项研究中,作者提出了 SAM-Med3D,这是一种专门用于3D体素医学图像分割的三维 SAM 模型。SAM-Med3D 在大规模的三维医学图像数据集上从头训练,其在不同组件中都采用了三维位置编码,直接整合三维空间信息,这使得它在体素医学图像分割任务中表现出卓越的性能。具体而言,SAM-Med3D 在提供仅一个提示点的情况下,相较于 SAM 在每个切片上提供一个提示点来说,性能提高了32.90%。这表明它能够在更少的提示点的情况下,在体素医学图像分割任务中取得更好的结果,这证明了它出色的可用性。 此外,作者还从多个角度广泛评估了 SAM-Med3D 的能力。对于不同的解剖结构,如骨骼、心脏和肌肉,在提供有限提示点的情况下,SAM-Med3D 明显优于其他方法。在不同的图像模态下,特别是核磁共振图像,通常需要比CT图像更多的提示点才能达到相同的性能,但 SAM-Med3D 在各种模态(包括核磁共振图像)、器官和病变下始终表现出色。此外,SAM-Med3D 的可迁移性也在不同的基准任务上经过了验证,该模型表现出了很强的潜力,因此 SAM-Med3D 有望成为一种强大的三维医学图像 Transformer 的预训练模型。 需要强调的是,不仅仅在数值结果方面,在可视化的结果中,SAM-Med3D 模型也表现出了更好的切片间的一致性和可用性。然而,三维模型在体积图像中的提示点变得更加稀疏,这增加了训练的难度。因此,如何更好地训练三维SAM仍然是需要进一步探索的领域,但这项研究为这一领域的未来发展提供了有力的方向和工具。