Swin-LiteMedSAM:基于盒的轻量级片段任意模型，用于大规模医学图像数据集 !

未来先知

发布于 2024-09-24 18:29:32

1360

发布于 2024-09-24 18:29:32

文章被收录于专栏：未来先知

医学影像对疾病的诊断和治疗至关重要，其中医疗影像分割被视为一个重要子任务。然而，自动医学影像分割模型通常是任务特定的，难以应对多种场景，如不同的影像模式和感兴趣区域。随着Segment Anything Model（SAM）的引入，训练针对各种临床场景的通用模型已成为可能。最近，提出了几种Medical SAM（MedSAM）方法，但这些模型通常依赖于强大的图像编码器来实现高性能，这可能在计算需求高和推理速度慢的现实世界应用中实用性不强。为解决此问题，MedSAM（LiteMedSAM）的轻量级版本可以提供一种可行的解决方案，该模型在降低计算资源和时间的同时，实现了高性能。在本研究中，作者引入了LiteMedSAM的新变体Swin-LiteMedSAM。该模型将Swin Transformer作为图像编码器，并采用多种类型的提示，包括从一个给定的边界框生成的基于框的点和涂抹。此外，模型在图像编码器与 Mask 解码器之间建立了跳接。在 "在笔记本电脑上进行医学图像中的任意分割"（CVPR 2024）的挑战中，作者的方法在分割性能和速度之间取得了良好的平衡，显著提高了与其他模式相比的整体结果。与挑战组织者提供的LiteMedSAM Baseline 相比，提出的模型在验证集上取得了0.8678的DDCU和0.8844的NSD分数，最终测试集上分别取得了0.8193的DDCU和0.8461的NSD分数，取得了四等奖。代码和训练模型可在https://github.com/RuochenGao/Swin_LiteMedSAM。

1 简介

医学影像诊断对评估疾病至关重要，医学影像分割，即从医学影像中提取特定结构（如肿瘤和器官）也一直受到广泛关注。深度学习方法在该领域取得了显著的进步，催生了大量针对特定场景的模型。然而，每个场景通常都需要训练一个专门的分割模型，这需要大量的工作。近年来，在自然语言处理（NLP）领域的巨模型（LLMs）迅速发展的启发下，研究行人开始探索将巨模型应用于计算机视觉。Segment Anything Model（SAM）就是其中之一，它通过与大量数据进行训练，试图将一般的图像分割任务统一。然而，尽管SAM在医疗影像分割任务上取得了高绩效，但其在大规模参数量和医疗图像的高空间分辨率下的计算和处理时间需求巨大。这使得SAM模型在实际部署中面临挑战，甚至对于没有GPU依赖的笔记本电脑或边缘设备上的非工业学术团体 Conducting research on them。因此，作为解决这一问题的SAM模型越来越受到关注。

原始的SAM模型由三个主要组成部分：图像编码器、提示编码器和解码器。在这三个部分中，由于使用的ViT-H[3]，图像编码器是主要导致计算和内存成本高的因素。为了减少资源消耗并加速处理，各种研究都试图使图像编码器更轻量级。例如，FastSAM[15]引入了一个基于CNN的框架，而MobileSAM[13]通过将ViT-H图像编码器的知识提炼为一个小巧的ViT基编码器来解决这个问题。另外，EfficientSAM[11]使用Masked Autoencoders(MAE)[4]框架，从大规模图像编码器中有效地传输知识到一个小巧的编码器，从而实现更资源效率、更好的性能设计。EfficientViT-SAM[14]进一步增强了这种方法，通过结合EfficientViT[1]和融合的MBConv块[9]来创建一个轻量级的图像编码器。最近，CVPR 2024举办的"在医疗影像上进行Segment Anything"的挑战寻求普遍可提示的医疗影像分割模型，可以部署在无GPU依赖的笔记本电脑或边缘设备上。组织者开发了LiteMedSAM2作为基础，使用[13]中描述的稀释策略。尽管LiteMedSAM专注于优化图像编码器以降低资源使用，分割性能受到影响。因此，作者的目标是提高性能，同时不牺牲效率。为了实现这一目标，作者使用轻量级的Swin Transformer作为图像编码器，并引入了两个额外的提示，一个是基于 Box 的点，一个是基于涂鸦的 Box ，除了原始的 Box 提示外。为此，作者提出作者的模型，Swin-LiteMedSAM。该模型的主要贡献如下：

作者使用了一个名为Swin Transformer的小型神经网络作为图像编码器[6]。与ViT相比，Swin Transformer在处理大型图像时，无论是计算还是内存使用方面都更加高效。此外，作者在图像编码器和 Mask 解码器之间建立了跳接连接，以增强特征融合。

作者引入了一种名为box-based points和box-based scribble的额外类型的提示，这些提示是从给定的边界框自动生成的，有效地提高了模型性能，同时没有显著增加资源成本。

总的来说，Swin-LiteMedSAM在保持高推理速度的同时，相比于LiteMedSAM实现了显著的性能提升。

2 Method

Data preprocessing

为了加速模型训练和推理阶段并减少内存消耗，作者将输入图像缩小到。这是通过首先根据最长边保持原始宽高比对图像进行缩放，然后进行零填充以达到最终尺寸实现的。为了数据规范化，作者采用了方法[8]中的描述。请参阅[8]以获得更多信息。

请注意，CT、MR、US 和 PET 等灰度图像通常只有一个通道，而来自内窥镜、皮肤镜和视网膜成像等介质的 RGB 图像通常有三个通道。为了在模型训练过程中保持一致性，作者对灰度图像的通道数进行了复制，将它们从单一通道转换为三个通道。

Proposed method

如图1所示，作者模型的结构主要包括三个部分：图像编码器、提示解码器和 Mask 解码器。这三个部分的功能如下详细介绍。

图像编码器架构灵感来自LiteMedSAM原始的微小ViT设计。输入首先经过两个卷积层，捕捉到低层次的 spatial 特征并调整频道数为64。接着，编码器由四个阶段组成，其深度按照微型ViT配置（2, 2, 6, 2）排列。作者所使用的Swin块结构如图2所示。作者稍稍修改了标准的Swin块，在窗口化多头自注意力（W-MSA）模块和多层感知机（MLP）之间添加了一个带批量归一化卷积层，这样可以使得作者的编码器能够有效地捕捉到全局和局部特征。此外，四个阶段中通道数和空间分辨率始终保持一致于原始设计。最后，一个由几个卷积层和层归一化组成的头分支调整了频道数到256。

在提示解码器中，作者引入了两种额外的提示方式：基于框的点以及基于框的涂鸦，同时保留了原始的基于框的提示。基于框的点和框被组合成一个稀疏嵌入，基于框的涂鸦被用于密集嵌入。对于基于框的提示，参考了[10]和[2]中的洞见，证明使用多个点优于使用一个点，因此在提示解码器中作者使用了四个点。为了实现这一点，作者根据中心点将边界框区域划分为四个相等的子部分。在非零区域的每个子部分中随机生成一个点，从而得到四个分散在框内的点。如果一个子部分只包含零，则使用中心点。这种方法确保了点相对稀疏地覆盖了更多的区域。此外，一个基于框的涂鸦是在框内根据[16]中提供的算法随机生成的。涂鸦中的所有像素都设为1，并将其放入一个形状为（256，256）的全零矩阵的对应部分，以创建密集嵌入的 Mask 。同样，如果框中的所有像素都为零，则涂鸦被设置为一个形状为（256，256）的全零矩阵，以确保提示解码器专注于稀疏的提示嵌入部分，如图3（a）所示。

作者遵循SAM原设计，使用一个双向 Transformer 来处理从提示解码器和图像编码器得到的嵌入。此外，作者还建立了图像编码器与 Mask 解码器之间的 Short-Cut ，将最后三阶段的输出和一个由前三个阶段输出的拼接与多个卷积层连接并融合。这个输出然后与双向 Transformer 的输出以及一个上放大模块双倍放大图像分辨率。相似地，上放后的输出与图像编码器的第一阶段输出进行拼接，并进一步上采样以返回原始空间分辨率。损失函数包括 Mask 预测损失和IoU预测损失：

其中是预测 Mask 与地面 truth Mask 之间的Dice损失和二进制交叉熵之和，而是预测和实际IoU之间的均方误差损失。

Inference strategy for box-based points and box-based scribble

基于方框的点生成和基于涂抹的点生成的策略在训练和推理阶段之间存在一些差别。在训练阶段，生成四个点的范围在整个边界框内，这使模型接触到各种情况，并提高其泛化能力。然而，在整个框内随机生成点可能不太理想，因为在推理阶段，物体通常位于框的中心部分。随机生成可能容易将一些点放在边界附近，这些点的效果不好，甚至可能对性能产生负面影响。此外，对于单点提示的情况，框的中心点始终是首选 [2]。同样，框的两个角点已经提供了物体的部分外部和周围信息。因此，点应在相对中心的位置更好地分布。对于给定的框框，由其左上角点 (, ) 和右下角点 (, ) 表示，作者引入两个变量，offset 和 offset，来调整高度和宽度的坐标，使得四个点不会出现在边缘区域，如图3(b) 所示。这种调整可以用以下公式表示：

在这里，新的左上角点 (, ) 和右下角点 (, ) 形成一个新的方框，并在其中随机生成四个点。offset 的范围为，offset 的范围为，其中和分别表示图像的宽度和高度。在本文中，作者将offset的范围调整为，将offset的范围调整为，以确保点分布更接近中心。此外，offset 和 offset 随机调整每个样本在其范围以内，以实现更好的整体性能。作者还遵循训练阶段的点分布策略，以确保四个点位于图像的四个象限。

然后，图3(c)说明了在推理阶段生成涂抹的策略。考虑点生的经验分布，作者认为将涂抹放置在边缘更比点更能捕捉轮廓信息有效。因此，作者将offset的范围调整为，将offset的范围调整为，以扩展生成涂抹的区域。请注意，作者根据先验知识在非零像素值区域内生成涂抹，人们通常避免在零像素值区域绘制涂抹。

3 Experiments

Dataset and evaluation metrics

3.1.1 Training and validation dataset

作者仅使用提供的挑战数据集，不使用其他公共数据集。该数据集包括11个模态：CT、MRI、PET、X射线、超声、乳腺摄影、OCT、内窥镜、眼底、皮肤镜和显微镜，总计超过100万个2D图像- Mask 对。

测试数据集 本挑战中的测试集是隐藏的，所有测试图像均来自全世界20多个不同机构的新收集。

评价指标 评价指标包括准确性的Dice相似度系数（DSC）和归一化表面Dice（NSD）以及效率的Docker容器运行时间。这些指标一起决定了排名。请注意，仅提供了平均结果。评价平台环境见表1。

Implementation details

以下是关于实现细节部分的开头。

3.2.1 Training environment settings

训练环境如图2所示。

3.2.2 Training protocols

作者的训练策略分为两个阶段。在第一个阶段，作者利用知识蒸馏将大型ViT-B图像编码器中的信息转移到微小的Swin Transformer作为图像编码器。需要注意的是，作者从ViT-B编码器预保存了输出图像嵌入以加速蒸馏过程。在第二个阶段，作者使用第一阶段训练好的图像编码器来训练整个模型。这两个阶段的训练细节如表3所示。

数据采样策略 在训练中，作者从数据集中随机采样图像案例。如果是3D的案例，例如CT、MR或PET扫描，作者从3D图像中随机选择一片。如果是2D的案例，例如X光或显微图像，作者直接使用图像。这种策略大大降低了训练时间，并确保了在不同模态上训练样本的分布更加平衡。

数据增强 作者对图像应用垂直和水平翻转，每个方向的概率均为50%。

推理环境设置 在推理阶段，运行环境与训练阶段不同。作者构建了一个使用'python:3.10-slim'镜像的docker容器，并安装了PyTorch 2.2.2的CPU版本。其余方面与训练阶段相同。

4 Results and discussion

Quantitative results on validation set

表格4显示，Swin-LiteMedSAM在平均DSC（86.70%）和NSD（88.55%）方面高于LiteMedSAM，后者在DSC和NSD上的得分分别为83.81%和83.26%。总的来说，与LiteMedSAM相比，Swin-LiteMedSAM在九种模态上的表现更平衡且全面。它在PET和显微镜上的表现显著提高，同时还在大多数模态上保持强大的性能。然而，模型在美国模态上的DSC和NSD得分明显下降。

然后，表格5进一步突显了作者所提出的每种组件的重要性，尤其是在实现优越分割性能方面，尤为明显的是跳过连接的引入，以及对基于框的点和涂鸦的采用。在这里，引入了两个额外的基于框的提示，其提高有限。这可能是由于以下两个因素：首先，一些提示可能由于随机放置在次优位置，从而影响整体性能。其次，训练不到位也是影响因素。尽管数据采样策略有助于平衡模态分布并加快训练速度，但它大大减少了训练样本的数量。这种减少可能会阻碍提示的有效训练，后者需要大量多样化的案例才能发挥最佳效果。

Quantitative results on testing set

表6表明，作者的方法在大多数成像模态（包括CT图像、PET和X射线图像）上显著优于LiteMedSAM，这些方法不仅提高了DSC，还减少了所有模态的运行时间。具体来说，与LiteMedSAM相比，作者的方法在CT图像上的绝对DSC提高了17.15%，相对提高了30.76%，NSD提高了18.51%，相对提高了31.75%。同样，所有模态的运行时间都更快。对于PET和X射线模态，作者的方法不仅展示了良好的DSC和NSD结果，而且与LiteMedSAM相比具有竞争力。在PET中，虽然NSD提高了，但保持了相似的DCS表现，显著减少了运行时间。在X射线方面，尽管与LiteMedSAM相比，DSC略低，但差异微小，仍具有竞争力。此外，作者还观察到LiteMedSAM的原始稳定性存在显著不稳定现象。以CT模态为例，LiteMedSAM在验证集上表现极佳，超过了Swin-LiteMedSAM。然而，在测试集上评估CT时，CT的性能出现了显著下降，DSC从92.26%降为55.75%，NSD从94.90%降为58.48%。尽管Swin-LiteMedSAM在CT模态上也遇到了类似的问题，但性能下降要温和得多。此外，这个问题在其他模态中也得到了证实，进一步证实了Swin-LiteMedSAM模型在稳定性方面具有更好的表现，这对于现实世界的应用是至关重要的。

Qualitative results on external public dataset

由于挑战验证和测试集的真实 GT 不可用，作者选择SegRap2023 [7]，一个包含多个器官标注的公共头颈CT数据集来进行模型性能验证。

如图4所示，作者从SegRap2023中挑选了三个代表性的例子来 visually 检查作者的模型的性能。在第一种情况中，作者的模型在脑分割上表现出了强大的性能。这主要归功于大脑的巨大尺寸和与周围组织显著的对比度。进入第二种情况，作者观察到即使对于较小的目标，如脊髓、食道和气管，作者的模型也保持了良好的性能。然而，在第三种情况中，作者的模型与 GT 值相比性能有所不足。主要问题是医学图像中的语义存在歧义。例如，在进行口腔分割时，作者的方法只识别出了牙齿。这种差异源于口腔的边界提示框也可以被解释为只分割牙齿。在这种情况下，很难提供更精确的提示信息来指定分割目标。### 限制与未来工作

作者方法的主要限制是在这个挑战中，作者的模型使用2D图像进行训练和验证，而医学成像数据，如CT、MRI和PET，通常为3D格式。目前，作者通过在单个2D切片上进行预测来处理这些3D图像，这没有充分利用3D解剖信息，并可能阻碍性能的改进。问题的关键在于模型接收的输入提示信息通常基于2D信息，例如边界框和点。因此，在未来的工作中，作者计划探索如何在3D中提供有效的提示信息，并适当地调整模型以处理3D图像。

此外，作者为控制基于框的点的分布和涂鸦应用了一些手动规则，这无法找到最佳设置，如果设置不当，可能会对整体性能造成伤害。另外，由于医学成像方式和分割目标的形状各异，点的分布和涂鸦的分布应相应地调整。因此，为基于框的点生成学习和基于涂鸦生成涂鸦的方法将极大地有益于改进模型性能。