SAM-Med3D论文阅读笔记

前言

CodeLink: https://github.com/uni-medical/SAM-Med3D.

Motivation:目前医学领域的SAM采用逐切片方式处理3D医学图像，忽略了重要的3D空间信息，导致性能受限。此外也尝试通过2D到3D的适配来捕获3D空间信息, 但受限于较小的数据规模和固有的2D架构，限制了其全面建模3D空间信息的能力。因此为解决2D到3D适应方法的局限，我们提出SAM-Med 3D，采用全新可学习的3D架构，并在大规模、多类别的医学图像数据集上训练

Contribution:        

       1) 本文介绍了一种名为SAM-Med 3D的三维空间注意力模型。该模型通过从头开始训练大规模的三维医学图像数据集。

        2)此外，该SAM-Med 3D模型还经过验证，可作为3D医学图像Transformer的强大预训练模型。虽然3D模型增强了可用性，但体积图像内的提示稀疏性对模型捕获空间信息和有效利用稀疏提示的能力提出了挑战。

不足之处：尽管SAM-Med 3D在医学图像分割任务中取得了显著成果，但未来的研究仍需要探索新的3D提示形式和更适合3D环境的训练策略。

        一、Abstract

        尽管SAM在2D自然图像分割上取得了令人惊叹的成果，但是其在3D体素医学图像的分割上的应用仍然存在重大的缺陷，也就是次优性能以及不稳定的预测。该问题很难通过在医学图像数据集中微调SAM来解决，因为原有的SAM 2D结构本身就忽略了3D空间信息。因此本文提出了SAM-Med3D, 用与3D医学图像分割的修改版SAM的最全面研究。首先，通过全面地将SAM重新制定为在全面处理的大规模体积医学数据集上训练的彻底的3D架构。其次，提供了关于其性能的全面评估。 在15个数据集上评估了其性能并且分析它从多个角度，包括解剖结构，模态，目标和泛化能力。该方法与SAM相比，展示了显著增强的效率和用于广泛的三维体积医学图像分割的能力。

        二、 Introduction

        医学图像分割是医疗保健领域的关键任务，但现有技术往往局限于特定器官或病变的分割，这限制了其临床适用性。由于3D解剖结构复杂性和注释数据有限，开发通用性强且精确的分割方法仍是一大挑战。

        SAM作为一种强大的视觉基础模型(VFM)在多个领域展现了出色的zero-shot分割性能，但其对医学领域的适用性受限。通过微调，如MedSAM和SAM-Med2D，可以提升其医学图像分割能力。然而，这些方法采用逐切片方式处理3D医学图像，忽略了重要的3D空间信息，导致性能受限。此外，近些年来一些研究人员为克服逐切片方法的局限性，研究人员尝试通过2D到3D的适应来捕获3D空间信息。但这些方法受限于较小的数据规模和固有的2D架构，限制了其全面建模3D空间信息的能力。

        上图描述了SAM， SAM-Med2D与我们的方法SAM-Med3D在3D医学图像上的比较

        为解决2D到3D适应方法的局限，我们提出SAM-Med 3D，采用全新可学习的3D架构，并在大规模、多类别的医学图像数据集上训练。此数据集远超现有最大医学图像分割数据集，为模型提供了丰富的3D空间信息学习机会。

        三、 Method

        3.1 模块改动

        

        上图表示我们所修改后的SAM框架，把原生的2D组件都用对应的3D组件进行了替换，即3D Image Encoder， 3D Prompt Encoder与 3D Mask Decoder。

        3D Image Encoder:

                patches首先使用了一个 16*16*16的3D卷积进行了embedding, 然后用3D 可学习的绝对位置编码与其进行配对(3D ABS PE) 。这一部分就是将SAM中的2D PE扩展了一个维度，让其变成3D的绝对位置编码。然后将其输入到3D 注意力模块中进行对应Q, K, V 的计算。
         

              3D Attention Block:

                        我们把一个3D相对位置编码(3D relative PE)嵌入到了SAM中的多注意力(MHSA)机制中，让其能够直接捕获空间信息的细节。除了扩展了一个维度外，其余与SAM中的多注意力模块一致。

                

        3D Prompt Encoder:

        对于稀疏prompt, 使用了3D Abs PE来让其表示3D 空间信息， 同时密集prompt使用3D卷积及逆行处理。

        

        3D Mask Decoder:

        在这部分中， 我们利用3D的反卷积技术， 将3D的上采样过程融入到了SAM的mask decoder中。

        3.2 初步实验

                我们设计了一个初步实验来比较预训练好权重的SAM在我们SAM-Med3D上的效果。

                在3D模型中重用SAM权重的直接解决方案是对这些层采用权重复制策略，其中权重的形状表现出差异。以Conv为例，我们将2D卷积的内核复制D次，并将它们堆叠以形成3D卷积，其中D表示第三维中内核的大小。

              使用和不使用预先训练的权重， 我们在AMOS数据集上训练了SAM-Med 3D，进行了25万次迭代。

                

                从该表格可以看出， 使用SAM中预训练好的权重并未带来显著的性能提高， 因此我们在一个大规模的3D医学数据中从头开始训练我们的SAM-Med3D网络。

        3.2.1 大规模3D医学数据集构建

              训练集：  我们构建了一个基于大量公共以及私人的3D医学图片数据集用于训练。

                验证集：

                我们收集了13个公用数据集用来审查多临床场景下的情况，并且也引入了来自MICCAI 2023的两个私有数据集来验证不同模型的性能。 评估数据集包含了七个关键的解剖结构(胸部和腹部器官、大脑结构、骨骼等)以及还包括在医学领域具有重要意义的五种病变类型(MRI, CT , US..)。此外，它还包含了具有挑战性的、之前未知的目标，总共有153个不同的目标，分属不同的类别。我们的验证集可以分为三大部分：

                器官分割基准：我们使用了包括Totalseg-Test、AMOS-Val、BTCV和HaN-Seg在内的四个数据集，涵盖了CT和MRI模式下的100多种器官掩模。其中，HaN-Seg数据集提供了器官风险注释，增加了验证集的挑战性。我们通过利用这些数据集，进行了全面的测试和验证。

                脑部结构分割基准：我们综合使用了Totalseg-Test、AMOS-Val、BTCV和HaN-Seg四个数据集，这些数据集提供了100多种器官的掩模，包括CT和MRI模式。其中，HaN-Seg的器官风险注释增加了验证集的难度。我们利用这些数据集进行了全面测试，以验证相关模型或算法的性能。

                病灶分割基准：我们使用了包括FeTA 21、FeTA 22、iSeg 17、iSeg 19、MRBrains 13、MRBrains 18和cSeg 22在内的七个数据集，涵盖了MRI图像中大脑、脑干和小脑结构的注释。同时，结合Totalseg-Test和HaN-Seg的CT图像大脑结构掩模，这些不同数据源和MRI序列差异的数据集共同构成了一个具有挑战性的测试基准。

                3.2.1 数据处理

        为了标准化这些多样的数据， 我们按照一下四个步骤来清洗和处理所有收集到的数据：

        1)基于Meta信息的数据清洗。

                

        2)基于连通域的掩码清理。

                

        3)基于连通域的标签质量改进。

        

        

4)基于对称性的标签质量改进。

      

        3.3 评估基于SAM的模型

                我们全面评估了SAM、SAM-Med 2D及我们提出的SAM-Med 3D，旨在为3D医学图像的可验证分割任务设定基准。

                我们构建了一个包含13个不同公共数据集的3D体积分割验证集，该数据集具有多样性和代表性，包含各种目标和模态。这种评估方式将全方位地研究SAM-Med 3D在3D体积分割任务中的性能，从而加深我们对该算法的理解。

                我们评估了SAM-Med 3D等模型的效率，并模拟了临床场景中的互动分割与点提示模式。SAM-Med 3D使用基于块的推断方法，与目前先进的医学图像分割方法相比具有推理时间上的优势。在3D医学图像推断中，2D方法缺乏切片间交互，导致性能上限，而3D方法则能进行全局交互，提高分割准确性。

四、 Experiments

        4.1 定量评估

                

                该表列出了SAM、SAM-Med 2D和我们提出的SAM-Med 3D在验证集上的性能指标。

               实验显示，SAM-Med 2D(嵌入了医学领域知识)在性能上明显优于SAM，SAM-Med 3D则表现更为出色，Dice评分提高了32.90%，且在不同提示点计数上均超越SAM和SAM-Med 2D。此外，SAM-Med 3D在计算效率上也具有显著优势，推理时间仅为SAM的15%，同时保持较高的Dice评分。

        

        该表是从解剖结构和病变角度进行比较。我们根据解剖结构对SAM-Med 3D的性能进行了深入评估，发现其在各种器官和病变上均表现出色。相比SAM，SAM-Med 3D在仅提供有限提示点的情况下，更能准确定位目标，尤其是在复杂解剖结构上。在所有解剖结构的评估中，SAM-Med 3D与SAM-Med 2D性能相当，甚至在某些情况下Dice评分提升了17.71%。 我们发现，腹部和胸部目标不遵循SAM-Med3D的这种领先趋势（与SAM-Med2D相比）。这种情况可以归因于这些目标的更简单的3D结构，这减少了对3D空间细节的依赖。

        在评估中，我们还发现SAM-Med 3D在可见病变的分割上表现优越。对于不可见病变，SAM-Med 3D在有限提示下表现良好，性能随着提示增多而提升。在10个提示下，SAM-Med 3D的性能达到48.44%，超过SAM-Med 2D在200个提示下的47.66%。

        除此之外，我们还对比了SAM, SAM-Med2D与我们的方法SAM-Med3D在三个主要模态上的DICE得分情况，如下所示。

        我们发现SAM在三种模态下的DICE得分（1 point prompt情况下）都低于20%。 对于CT图像来说，随着point点数的增加情况有所提升， 5 point时会提升到50%左右。而对于MRI以及US图像则情况就不容乐观了。

        SAM-Med2D在MRI图像上当给定1 point的情况下，性能与我们的方法相比处于绝对的劣势。而在CT与US的模态想，随着point点数的增加， SAM-Med2D与SAM-Med3D的性能所差无几。

        4.2 定性评估

        为了保持比较的公平性，我们在每个轴上将3D图像均匀地重采样到1.5 mm的间距，并使用每种方法的最佳分辨率（SAM为1024*1024，SAM-Med 2D为256*256，SAM-Med 3D为128*128 ）进行测试。

        对于每个可视化的情况，描绘了两个视图：（1）轴向：该水平横截面将身体或器官分成上段和下段。它是观察器官解剖学最常用的视角。(2)矢状面/冠状面：表示垂直横截面，该视图将身体或器官分为左/右或前/后部分。

       

        基于上述两幅图我们发现：1）SAM-Med 3D需要的提示明显更少。2) SAM-Med 3D具有更好的切片间一致性。