SAMUS
期刊分析
期刊名： arXiv
期刊信息： 2023-9-13
其余信息： 代码
摘要
分割任意模型（SAM）是一种著名的通用图像分割模型，最近在医学图像分割领域引起了相当多的关注。尽管 SAM 在自然图像上表现出色，但在处理医学图像时**，尤其是涉及低对比度、模糊边界、复杂形状和小尺寸物体的图像时，它会遇到显着的性能下降和有限的泛化能力。** 在本文中，我们提出了 SAMUS，一种专为超声图像分割而定制的通用模型。与之前基于SAM的通用模型相比，SAMUS不仅追求更好的泛化性，而且还追求更低的部署成本，使其更适合临床应用。具体来说，基于SAM，引入并行CNN分支，通过跨分支注意力将局部特征注入ViT编码器，以实现更好的医学图像分割。然后，开发了位置适配器和特征适配器，以使 SAM 从自然领域适应到医学领域，从需要大尺寸输入 (1024×1024) 到小尺寸输入 (256×256)，以实现更适合临床的部署。收集了一个全面的超声数据集以供验证，该数据集包含约 30k 图像和 69k 掩模，涵盖六个对象类别。大量的比较实验证明了 SAMUS 在特定任务评估和泛化评估方面均优于最先进的特定任务模型和通用基础模型。此外，SAMUS 可以部署在入门级 GPU 上，因为它已经摆脱了长序列编码的限制。代码、数据和模型将在https://github.com/xianlin7/SAMUS发布。
引言
医学图像分割是识别和突出医学图像中特定器官、组织和病变的关键技术，是计算机辅助诊断系统的组成部分（Liu et al. 2021）。人们已经提出了许多用于自动医学图像分割的深度学习模型，展示了巨大的潜力（Ronneberger、Fischer 和 Brox，2015 年；Wu 等人，2022 年）。然而，这些模型都是针对特定对象量身定制的，应用于其他对象时需要重新训练，给临床使用带来很大不便。
分割任何模型（SAM）作为视觉分割的通用基础模型，由于其跨不同对象的卓越分割能力和强大的零样本泛化能力而赢得了相当多的赞誉（Kirillov et al. 2023）。根据用户提示，包括点、边界框和粗掩模，SAM能够分割相应的对象。因此，通过简单的提示，SAM可以毫不费力地适应各种分段应用。这种范例能够将多个单独的医学图像分割任务集成到一个统一的框架（即通用模型）中，极大地促进了临床部署（Huang et al. 2023）。
尽管构建了迄今为止最大的数据集（即 SA1B），但由于缺乏可靠的临床注释，SAM 在医学领域的性能迅速下降（Huang et al. 2023）。已经提出了一些基础模型，通过在医学数据集上调整 SAM 来使 SAM 适应医学图像分割（Ma 和 Wang 2023；Wu 等人 2023）。然而，与 SAM 相同，它们在特征建模之前对输入图像执行无重叠的 16× 标记化，这破坏了对于识别小目标和边界至关重要的局部信息，使得它们难以分割具有复杂/线状形状的临床对象，边界弱、尺寸小或对比度低。 此外，大多数需要1024×1024大小的输入，由于生成的长输入序列对GPU消耗造成很大负担。
在本文中，我们提出 SAMUS，将 SAM 卓越的分割性能和强大的泛化能力转移到医学图像分割领域，同时降低计算复杂度。 SAMUS继承了SAM的ViT图像编码器、提示编码器和掩模解码器，并对图像编码器进行了量身定制的设计。首先，我们通过减少输入大小来缩短 ViT 分支的序列长度，以降低计算复杂度。然后，开发了一个特征适配器和一个位置适配器来对 ViT 图像编码器从自然领域到医学领域进行微调。为了补充 ViT 图像编码器中的本地（即低级）信息，我们引入了一个并行 CNN 分支图像编码器，与 ViT 分支一起运行，并提出了一个跨分支注意模块，以使 ViT 分支中的每个补丁能够同化来自 CNN 分支的本地信息。 此外，我们构建了一个名为 US30K 的大型超声数据集来综合评估 SAMUS 的功效。实验结果表明，SAMUS 在特定任务和通用医学图像分割方面均优于最先进的方法。更重要的是，SAMUS 表现出卓越的泛化能力，同时与 SAM 相比大大降低了训练成本。这些贡献可总结如下：
- 基础模型 SAMUS 专为通用超声图像分割而设计，与 SAM 相比，需要更少的 GPU 资源。
- CNN 分支图像编码器和跨分支注意力模块，可有效补充 SAM 的 ViT 图像编码器的局部信息。
- 功能适配器和位置适配器，用于微调 ViT 分支图像编码器，进一步优化医疗领域的 SAM。
- 包含 30,106 个图像和 68,570 个掩模的大型超声数据集，可彻底评估 SAMUS 的有效性。
相关工作
视觉调整（Visual Tuning）
随着计算机视觉中基础模型的惊人发展，人们提出了一系列视觉调整方法来使这些基础模型适应下游任务。一般来说，最近的视觉调整方法可以分为五个主要类别，包括微调、参数调整、重新映射调整、提示调整和适应调整（Yu et al. 2023）。具体来说，微调方法涉及调整预训练模型的整个参数集或选择性地微调预训练模型的特定部分（Ma 和 Wang 2023）。参数调整方法直接修改模型参数的权重或偏差（Lian et al. 2022）。重映射方法通过知识蒸馏、基于权重的重映射或基于架构的重映射将学习到的信息从预训练模型转移到下游模型（Zhao et al. 2022）。提示调整通过将一组可学习参数与输入结合起来或设计一个子网络来生成视觉提示来引入下游任务的知识（Wang 等人，2023）。适配器调整是最广泛采用的策略，通过将额外的可学习参数与冻结的预训练模型结合起来，促进下游任务的学习（Wu et al. 2023）。
使 SAM 适应医学图像分割
SAM 在自然图像中表现出了卓越的性能，但在一些医学图像分割任务中表现不佳，特别是在形状复杂、边界模糊、尺寸小或对比度低的物体上（Huang et al. 2023）。为了弥补这一差距并使 SAM 能够有效地适应医学图像领域，人们提出了几种使用有限的下游医学数据集来调整 SAM 的方法。 MedSAM 通过冻结图像编码器和提示编码器，以可接受的成本在医学图像上训练 SAM，重点是调整 SAM 的掩模解码器（Ma 和 Wang 2023）。 SAMed 在图像编码器上应用基于低秩（LoRA）的策略，以较低的计算成本调整 SAM，使其更适合医学图像分割（Zhang 和 Liu 2023）。 MSA 在 ViT 图像编码器的每个变换器层上采用两个 downReLU-up 适配器来引入特定于任务的信息 (Wu et al. 2023)。如图1所示，与当前基于SAM的基础模型相比，所提出的SAMUS更注重补充局部特征和减少GPU消耗，这对于临床场景中准确且易于部署的医学图像分割至关重要。
方法
整体框架
如图2所示，SAMUS的整体架构继承自SAM，保留了提示编码器和掩码解码器的结构和参数，没有任何调整。相比之下，图像编码器经过精心修改，以解决局部特征不足和计算内存消耗过多的挑战，使其更适合临床友好的分割。主要修改包括减少输入大小、重叠补丁嵌入、向 ViT 分支引入适配器、添加 CNN 分支以及引入跨分支注意力 (CBA)。具体来说，输入空间分辨率从 1024 × 1024 像素缩小到 256 × 256 像素，由于 Transformer 中的输入序列较短，因此 GPU 内存成本大幅降低。重叠的补丁嵌入使用与 SAM 中的补丁嵌入相同的参数，而其补丁步长是原始步长的一半，很好地保留了补丁边界的信息。 ViT分支中的适配器包括一个位置适配器和五个特征适配器。由于输入大小较小，位置适配器用于适应较短序列中的全局位置嵌入。第一个特征适配器遵循重叠补丁嵌入，将输入特征与预训练 ViT 图像编码器所需的特征分布对齐。其余的特征适配器连接到全局变换器中前馈网络的剩余连接，以微调预训练的图像编码器。就CNN分支而言，它与ViT分支并行，通过CBA模块为后者提供补充的局部信息，CBA模块以ViTbranch特征作为查询，并与CNN分支的特征建立全局依赖关系。需要注意的是，CBA 仅集成到每个全局变压器中。最后，将两个分支的输出组合起来作为 SAMUS 的最终图像特征嵌入。
ViT 分支中的适配器
为了促进 SAM 的训练图像编码器（即 ViT 分支）泛化到较小的输入尺寸和医学图像领域，我们引入了一个位置适配器和五个特征适配器。这些适配器可以有效地调整 ViT 分支，同时只需要更少的参数。具体来说，位置适配器负责调整位置嵌入以匹配嵌入序列的分辨率。它首先通过步长和内核大小为 2 的最大池对位置嵌入进行下采样，实现与嵌入序列相同的分辨率。随后，应用内核大小为 3 × 3 的卷积运算来调整位置嵌入，进一步帮助 ViT 分支更好地处理较小的输入。所有特征适配器都具有相同的结构，包括三个组件：向下线性投影、激活函数和向上线性投影。每个特征适配器的流程可以表述为：
A ( x ) = G ( x E d ) E u \mathcal{A}(x)=\mathcal{G}\left(x E_{d}\right) E_{u} A(x)=G(xEd)Eu
CNN分支
CNN 分支由顺序连接的卷积池块组成。具体来说，输入最初通过单个卷积块，然后通过三个卷积池块进行处理。然后，CNN 分支中的特征图与 ViT 分支中的特征图共享相同的空间分辨率。在 CNN 分支的其余部分中，此类单个卷积块按顺序重复四次。更多细节如图2所示。CNN分支的这种极简、轻量级的设计是为了防止训练过程中的过度拟合。
Cross-branch 注意力
跨分支注意力（CBA）模块在 CNN 分支和 ViT 分支之间建立了一座桥梁，以进一步用 ViT 分支补充缺失的局部特征。对于来自 ViT 分支 F v F_v Fv 和 CNN 分支 F c F_c Fc的一对特征图，单头中的跨分支注意力可以表示为：
F ( F v , F c ) = ( S ( F v E q ( F c E k ) T d m ) + R ) ( F c E v ) \mathcal{F}\left(F_{v}, F_{c}\right)=\left(\mathcal{S}\left(\frac{F_{v} E_{q}\left(F_{c} E_{k}\right)^{T}}{\sqrt{d_{m}}}\right)+R\right)\left(F_{c} E_{v}\right) F(Fv,Fc)=(S(dm FvEq(FcEk)T)+R)(FcEv)
训练策略
在训练之前，SAMUS 使用在 SA-1B 上训练的权重初始化从 SAM 继承的参数。其余参数随机初始化。在训练过程中，仅更新适配器、CNN分支和CBA模块的参数，而其他参数保持冻结 。 训练过程使用组合损失函数进行监督，包括骰子损失和二元交叉熵损失。为了方便使用，SAMUS只使用最简单的正点提示。我们通过随机采样标签前景区域中的一个点来模仿专家提供提示的过程。 SAMUS 由 Adam 优化器进行训练，初始学习率为 0.0001，批量大小为 8，持续 200 个 epoch。
实验
数据集
为了全面评估 SAMUS 的有效性，我们构建了一个名为 US30K 的大型超声数据集，如表 1 所示，其中包含来自七个公开数据集的数据，包括 TN3K (Gong et al. 2023)、DDTI (Pedraza et al. 2015) 、TG3K（Wunderling 等人，2017 年）、BUSI（Al-Dhabyani 等人，2020 年）、UDIAT（Yap 等人，2020 年）、CAMUS（Leclerc 等人，2019 年）和 HMCQU（Kiranyaz 等人，2020 年）。 TN3K 和 TG3K 中的数据按照 TRFE 划分为训练集、验证集和测试集（Gong et al. 2023）。 BUSI 被随机分成 7:1:2 分别用于训练、验证和测试。 CAMUS 根据挑战首先分为训练集和测试集（Leclerc et al. 2019）。然后，我们从训练集中随机选择 10% 的患者来验证模型，其余数据作为最终的训练数据。为了评估不同模型的泛化能力，US30K 中的其他数据集在训练和验证过程中是看不见的。为了根据最先进的 (SOTA) 特定任务方法评估 SAMUS 的分割性能和泛化能力，在 TN3K、BUSI 和 CAMUS 数据集上分别实施和训练了几种 SOTA 方法以进行比较。此外，SAMUS 与其他基础模型的比较是通过在整个 US30K 数据集上训练它们并在单独的任务上评估它们来进行的。
与 SOTA 特定任务方法进行比较
比较方法：选择了 12 种 SOTA 特定任务方法进行比较，涵盖基于 CNN、基于 Transformer 和 CNN-Transformer 混合方法。基于 CNN 的方法包括 U-Net (Ronneberger, Fischer, and Brox 2015)、CPFNet (Feng et al. 2020)、CA-Net (Gu et al. 2020)、CE-Net (Gu et al. 2019) 和AAU-Net（Chen 等人，2022）。基于 Transformer 的方法包括 SwinUnet (Cao et al. 2022)、SETR (Zheng et al. 2021) 和 MISSFormer (Huang et al. 2022)。 CNNTransformer 混合方法包括 TransUNet (Chen et al. 2021)、TransFuse (Zhang, Liu, and Hu 2021)、FAT-Net (Wu et al. 2022) 和 H2Former (He et al. 2023)。
定量结果 ：表 2 总结了 TN3K、BUSI、CAMUS-LV、CAMUS-MYO 和 CAMUS-LA 上不同任务特定方法的定量结果。在这些最先进的方法中，H2Former 实现了最佳性能在 TN3K 和 CAMUSMYO 上，平均 Dice 得分分别为 82.48% 和 87.31%。 TransUnet、CA-Net 和 FATNet 分别在 BUSI、CAMUS-LV 和 CAMUS-LA 上取得最佳性能，平均 Dice 分数分别为 82.22%、93.59% 和 91.55%。相比之下，SAMUS 在 TN3K、BUSI、CAMUS-LV、CAMUS-MYO 和 CAMUS-LA 等五项任务上始终取得了较好的表现，平均 Dice 分数分别为 84.45%、85.77%、93.73%、87.46% 和 91.58% 。验证了SAMUS将SAM应用于医学图像领域的有效性。
定性结果 ：不同方法的定性分割结果，包括U-Net (Ronneberger, Fischer, and Brox 2015)、AAU-Net (Chen et al. 2022)、MISSFormer (Huang et al. 2022)、H2Former (He et al. 2022)、H2Former (He et al. 2022)。 2023）和 SAMUS，如图 4 所示。从视觉上看，由于超声图像对比度低、特征不均匀和对象边界模糊，分割超声图像构成了挑战。现有方法很难准确地区分目标与背景，导致大量的假阴性和/或假阳性。相比之下，SAMUS 在保留目标区域的完整性和减少误报方面表现出优越性。这得益于SAM框架的先天优势，以及SAMUS中引入的具体调整和设计。
泛化能力：不同特定任务方法泛化性能的定量比较如图3所示。在比较方法中，H2Former、TransUNet和TransFuse分别在DDTI、UDIAT和HMC-QU上取得了最好的性能。相比之下，SAMUS 在每个数据集上都超越了最佳比较方法，在 Dice 中分别平均提高了 7.06%、12.22% 和 7.42%。比较可见数据集和不可见数据集之间的性能，与跨三个不同分割任务的其他比较方法相比，SAMUS 遇到的性能下降最小。一个有趣的观察是，在乳腺癌分割任务中，SAMUS 在不可见数据集（即 UDIAT）上的性能甚至优于可见数据集（即 BUSI）上的最佳比较方法。它展示了 SAMUS 在处理未见过的领域方面的卓越泛化能力，展示了其在不同医学图像分割场景中的鲁棒性和适应性。
与 SOTA 基础模型比较
比较方法：选择四种SOTA基础模型进行比较，包括原始SAM（Kirillov等，2023）、MedSAM（Ma和Wang，2023）、SAMed（Zhang和Liu，2023）和MSA（Wu等，2023）。
定量结果：为了验证 SAMUS 作为各种下游任务的基础模型的通用性能，我们在 US30K 数据集上对基础模型进行了比较。如表 3 所示，在 SA-1B 上训练的模型 SAM 在未经调整的情况下在医学图像分割方面表现出显着的性能下降。 MedSAM 通过基于 US30K 数据集对 SAM 的掩模解码器进行简单的微调，显着提高了 SAM 的性能。 MAS是对比基础模型中表现最好的模型，有效提高了SAM的分割性能，在TN3K、BUSI、CAMUS-LV、Dice中平均提高了53.08%、27.65%、62.77%、53.05%和74.52%卡缪-MYO 和卡缪-LA。与 MSA 相比，SAMUS 在上述五个数据集上始终取得了显着的改进，Dice 得分分别为 83.05%、84.54%、91.13%、83.11% 和 92%。它验证了 SAMUS 中 CNN 分支和 CBA 模块的有效性，特别是对于补充对于医学图像分割至关重要的局部信息。
定性结果 ：不同基础模型（包括 SAM、MedSAM、SAMed、MSA 和 SAMUS）的定性分割结果如图 5 所示。如果不对医学图像进行调整，SAM 完全失去了分割所有内容的能力。通过对 SAM 应用调整方法，MedSAM、SAMed 和 MSA 可以在一定程度上恢复 SAM 的分割能力。然而，他们仍然难以准确地描绘超声图像中的分割边界，导致大量的假阴性和假阳性。相比之下，SAMUS 通过准确定位分割边界而表现出卓越的性能，即使对于低对比度的分割边界也是如此。与分析一致的是，用图像编码器补充局部信息是有帮助的，特别是对于医学图像分割中的边界/形状保留。
泛化能力 ：图6总结了不同基础模型在未见领域上的比较。总的来说，在医学图像分割任务中，在US30K上训练的基础模型的泛化性能远远优于原始SAM。就甲状腺结节分割、乳腺癌分割和心肌分割这三对分割任务而言，所有基础模型在心肌分割上都遇到了严重的性能下降，但在乳腺癌分割上泛化良好。 SAMUS 在所有三个未见过的数据集上始终保持最佳性能，在甲状腺结节、乳腺癌和心肌分割方面分别获得了 66.78%、78.06% 和 56.77% 的优异 Dice 分数。它强调了 SAMUS 卓越的泛化能力，在未见过的领域持续且大幅优于其他基础模型。
部署成本 ：我们对SAMUS和其他基础模型的部署效率进行了全面评估，包括GPU内存成本、模型参数、计算复杂度、推理速度、分割性能和泛化性能。为了便于比较，在训练时将batch size设置为1时测试GPU内存，并以千兆字节（即G）为单位进行测量。计算复杂度和推理速度通过每秒浮点运算数（即 GFLOP）和每秒帧数（即 FPS）来衡量。分割性能通过所有可见数据集的平均 Dice 分数来衡量，泛化性能根据所有不可见数据集的平均 Dice 分数来评估。所有上述指标均经过归一化并用雷达图描绘，如图 7 所示。在比较模型中，SAMed 表现出最低的 GPU 内存成本、模型参数、计算复杂度和最快的推理速度。然而，其分割和泛化性能均不如MSA和SAMUS。虽然 SAMUS 比其他模型拥有更多的参数，但其 GPU 内存成本和计算复杂度是第二低的，推理速度是第二快的，这表明 SAMUS 是一个更适合临床的模型。此外，SAMUS 的部署性能非常接近最易于部署的方法（即 SAMed），具有更好的分段和泛化性能。
消融实验
SAMUS中各个组件的有效性：SAMUS中的四个组件，包括CNN分支、CBA、特征适配器和位置适配器，依次引入原始SAM并在TN3K和BUSI数据集上进行训练以进行评估。如表9所示，耦合SAMUS的任何组件都可以有效提高SAM在医疗任务上的分割性能和泛化能力。即使是简单的位置适配器也可以将 TN3K、DDTI、BUSI 和 UDIAT 上的 Dice 中的 SAM 分别提高 50.6%、38.1%、26.77% 和 30.54%。通过引入局部特征，CNN分支比位置适配器获得了更显着的性能提升。此外，简单地将CNN分支和ViT分支的输出相加进行融合并不是最好的选择。引入CBA可以进一步促进局部特征的探索，从而实现Dice相对于CNN分支平均提升1.48%和2.11%。通过耦合所有四个组件，SAMUS 实现了最佳的分割性能和泛化能力。
不同提示对 SAMUS 的影响 ：为了分析提示的效果，我们评估了在不同点提示下在 US30K 上训练的 SAMUS 的性能。一般来说，SAMUS 对点位置和数量具有鲁棒性。对于类内表示变化较大的对象（例如甲状腺结节、左心室和心肌），在不同位置的单点提示下，Dice 中的分割性能变化在 1% 以内。而且，通过引入多点提示可以大大提高性能。相比之下，对于具有同质特征的对象（例如乳腺癌和左心房），不同单点提示下的性能变化在±0.3%的范围内，引入多点提示并不一定会带来性能提升。一个可能的原因是使用更多的点可能会产生信息冗余或排斥。
总结
在本文中，我们提出了 SAMUS，这是一种源自 SAM 的通用基础模型，用于临床友好且可推广的超声图像分割。具体来说，我们提出了一个并行 CNN 分支图像编码器、一个特征适配器、一个位置适配器和一个跨分支注意模块，以丰富小尺寸物体和边界区域的特征，同时减少 GPU 消耗。此外，我们构建了一个大型超声图像数据集 US30K，由 30,106 张图像和 68,570 个掩模组成，用于评估和潜在的临床用途。在可见域和不可见域上的实验都证明了SAMUS出色的分割能力和强大的泛化能力。此外，SAMUS 的 GPU 内存成本仅为训练整个 SAM 所需的 28%，并且 SAMUS 的推理速度比 SAM 快约 3 倍。
可借鉴参考
-
实验构成
-
常用结论