爱笑的小姐姐 · 2023年09月15日

UNet与SAM结合的正确的道路SAMUS,一路SOTA没对手!

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

image.png

SAM是一种杰出的通用图像分割模型,在医学图像分割领域近来引起了相当大的关注。尽管SAM在自然图像上表现出色,但在处理医学图像,特别是涉及低对比度、模糊边界、复杂形状和小尺寸目标的医学图像时,其性能明显下降且泛化能力有限。

在本文中,作者提出了SAMUS,这是一种专为超声图像分割定制的通用模型。与以前基于SAM的通用模型不同,SAMUS不仅追求更好的泛化性能,还降低了部署成本,使其更适用于临床应用。

具体而言,基于SAM,引入了一个并行的CNN分支,通过跨分支的注意力将局部特征注入ViT编码器,以实现更好的医学图像分割。然后,开发了Position Adapter和Feature Adapter,将SAM从自然领域适应到医学领域,并从需要大尺寸输入(1024×1024)适应到小尺寸输入(256×256),以更适合临床应用。

作者收集了一个包括约30,000张图像和69,000个Mask,涵盖6个目标类别的全面超声数据集进行验证。广泛的比较实验在任务特定评估和泛化评估下展示了SAMUS相对于最先进的任务特定模型和通用基础模型的优越性。此外,SAMUS可以部署在入门级GPU上,因为它已经摆脱了长序列编码的限制。

代码、数据和模型:https://github.com/xianlin7/SAMUS

1、简介

医学图像分割是一项关键技术,用于识别和突出显示医学图像中的特定器官、组织和病变,是计算机辅助诊断系统的重要组成部分。已经提出了许多深度学习模型,用于自动医学图像分割,展示出巨大的潜力。然而,这些模型都是为特定目标量身定制的,应用于其他目标时需要重新训练,给临床使用带来了很大的不便。

SAM作为视觉分割的通用基础模型,因其在各种目标上出色的分割能力和强大的零样本泛化能力而广受好评。根据用户提示,包括点、边界框和粗略Mask,SAM能够分割相应的目标。因此,通过简单的提示,SAM可以轻松地适应各种分割应用。这一范 paradigm 使多个个体医学图像分割任务集成到一个统一的框架中(即通用模型),极大地促进了临床部署。

尽管构建了迄今为止最大的数据集(即SA-1B),但由于可靠临床注释的稀缺性,SAM在医学领域遭遇了快速的性能下降。一些基础模型已经提出,通过在医学数据集上调整SAM,将其适应医学图像分割。然而,与SAM一样,它们在进行特征建模之前在输入图像上执行无重叠的16×Tokenization,这破坏了用于识别小目标和边界的局部信息,使其难以分割具有复杂/丝状形状、弱边界、小尺寸或低对比度的临床目标。此外,它们中的大多数需要尺寸为1024×1024的输入,导致GPU消耗巨大,因为生成了长输入序列。

在本文中,作者提出了SAMUS,旨在将SAM的出色分割性能和强大的泛化能力转移到医学图像分割领域,同时降低计算复杂性。SAMUS继承了SAM的ViT图像编码器、提示编码器和Mask解码器,具有针对图像编码器的定制设计。首先,作者通过减小输入尺寸来缩短ViT分支的序列长度,以降低计算复杂性。然后,作者开发了一个Feature Adapter和一个Position Adapter,用于将ViT图像编码器从自然领域微调到医学领域。

为了补充ViT图像编码器中的局部(低级)信息,作者引入了一个并行的CNN分支图像编码器,与ViT分支并行运行,并提出了一个跨分支注意力模块,使ViT分支中的每个块可以从CNN分支 assimilate 局部信息。此外,作者构建了一个名为US30K的大型超声数据集,包括30,106张图像和68,570个Mask,以全面评估SAMUS的有效性。实验结果表明,SAMUS在任务特定和通用医学图像分割方面均优于最先进的方法。更重要的是,与SAM相比,SAMUS具有显着的泛化能力,同时大大降低了培训成本。

贡献总结如下:

  • 一种基础模型SAMUS,专为通用超声图像分割而设计,与SAM相比,需要更少的GPU资源。
  • 一种CNN分支图像编码器和跨分支注意力模块,可以有效补充ViT图像编码器中的局部信息。
  • 一种Feature Adapter和Position Adapter,用于微调ViT分支图像编码器,进一步优化SAM以适应医学领域。
  • 一个包含30,106张图像和68,570个Mask的大型超声数据集,用于全面评估SAMUS的有效性。

2 相关工作

2.1 视觉调优

随着计算机视觉中基础模型的惊人发展,提出了一系列视觉调优方法,以将这些基础模型适应下游任务。通常,最近的视觉调优方法可以分为五大类,包括微调、参数调优、重新映射调优、提示调优和自适应调优。

具体来说,微调方法涉及调整预训练模型的整个参数集或有选择地微调预训练模型的特定部分。参数调优方法直接修改模型参数的权重或偏差。重新映射方法通过知识蒸馏、基于权重的重新映射或基于架构的重新映射将从预训练模型中学到的信息传递给下游模型。提示调优通过将一组可学习参数与输入进行结合或设计一个子网络来生成视觉提示的方式,引入了下游任务的知识。Adapter调优是最广泛采用的策略,通过将额外的可学习参数与冻结的预训练模型相结合,促进了下游任务的学习。

2.2 将SAM适应到医学图像分割

SAM在自然图像中表现出色,但在某些医学图像分割任务中遇到困难,尤其是在具有复杂形状、模糊边界、小尺寸或低对比度的目标上。

为了弥补这一差距,使SAM能够有效地适应医学图像领域,已经提出了几种方法,使用有限的下游医学数据集来调整SAM。MedSAM通过冻结图像编码器和提示编码器,重点调整SAM的Mask解码器,以可接受的成本在医学图像上进行训练。SAMed采用低秩(LoRA)策略对图像编码器进行调优,以更低的计算成本调整SAM,使其更适用于医学图像分割。MSA在ViT图像编码器的每个变换层上采用两个downReLU-upAdapter来引入任务特定信息。

image.png

如图1所示,与当前基于SAM的基础模型相比,提出的SAMUS更注重补充局部特征并减少GPU消耗,这对于在临床场景中进行准确且易于部署的医学图像分割至关重要。

3 方法

3.1 概述

image.png

如图2所示,SAMUS的总体架构继承自SAM,保留了提示编码器和Mask解码器的结构和参数,没有进行任何调整。相比之下,图像编码器经过精心修改,以解决局部特征不足和计算内存消耗过多的挑战,使其更适合于临床友好的分割。主要修改包括减小输入尺寸,重叠的Patch嵌入,向ViT分支引入Adapter,添加CNN分支以及引入跨分支注意力(CBA)。

具体来说,输入的空间分辨率从1024×1024像素缩小到256×256像素,由于Transformer中较短的输入序列,GPU内存成本大幅降低。重叠的Patch嵌入使用与SAM中的Patch嵌入相同的参数,但其Patch跨度减半,与原始跨度保持良好的信息一致性。

ViT分支中的Adapter包括Position Adapter和5个Feature Adapter。Position Adapter用于适应较短序列中的全局位置嵌入,因为输入尺寸较小。第一个Feature Adapter遵循了重叠的Patch嵌入,以使输入特征与预训练的ViT图像编码器所需的特征分布相一致。其余的Feature Adapter附加到全局Transformer中的前馈网络的残差连接上,以微调预训练图像编码器。

在CNN分支方面,它与ViT分支并行,通过CBA模块向后者提供互补的局部信息,该模块以ViT分支特征作为Query,并与CNN分支的特征建立全局依赖性。值得注意的是,CBA仅集成到每个全局Transformer中。最后,两个分支的输出被合并为SAMUS的最终图像特征嵌入。

3.2 ViT分支中的Adapter

为了促进SAM的训练图像编码器(即ViT分支)在更小的输入尺寸和医学图像领域的泛化,作者引入了一个Position Adapter和5个Feature Adapter。这些Adapter可以有效地调整ViT分支,同时只需要更少的参数。

具体来说,Position Adapter负责调整位置嵌入以匹配嵌入序列的分辨率。它首先通过Stride和Kernel-Size为2的最大池化对位置嵌入进行降采样,实现与嵌入序列相同的分辨率。随后,应用Kernel-Size为3×3的卷积操作来调整位置嵌入,进一步帮助ViT分支更好地处理较小的输入。

所有Feature Adapter具有相同的结构,包括下投影、激活函数和上投影3个组成部分。每个Feature Adapter的过程可以表示为:

image.png

3.3 CNN分支

CNN分支由顺序连接的卷积-池化块组成。具体来说,输入首先通过一个单一的卷积块,然后通过3个卷积-池化块进行处理。

然后,CNN分支中的特征映射与ViT分支的特征映射具有相同的空间分辨率。在CNN分支的其余部分,这种单一的卷积块被连续重复4次。CNN分支的这种简约和轻量级设计是为了防止训练过程中的过拟合。

3.4 跨分支注意力

跨分支注意力(CBA)模块在CNN分支和ViT分支之间建立了一个桥梁,以进一步通过ViT分支补充缺失的局部特征。

对于来自ViT分支Fv和CNN分支的特征映射对,单一Head的跨分支注意力可以表示为:

image.png

3.5 训练策略

在训练之前,SAMUS使用在SA-1B上训练的权重来初始化从SAM继承的参数。其余参数被随机初始化。在训练过程中,只有Adapter、CNN分支和CBA模块的参数会被更新,而其他参数将保持冻结。训练过程通过组合损失函数进行监督,包括Dice损失和二元交叉熵损失。

为了方便使用,SAMUS只使用最简单的正点提示。作者通过在标签的前景区域中随机采样一个点来模拟专家提供提示的过程。SAMUS使用Adam优化器进行训练,初始学习率为0.0001,批大小为8,共进行200个epochs的训练。

4、实验

4.1 与SOTA任务特定方法的比较

比较方法:选择了12种SOTA任务特定方法进行比较,涵盖了基于CNN、基于Transformer和CNN-Transformer混合方法。

  • 基于CNN的方法包括U-Net、CPFNet、CA-Net、CE-Net和AAU-Net
  • 基于Transformer的方法包括SwinUnet、SETR和MISSFormer
  • CNN-Transformer混合方法包括TransUNet、TransFuse、FAT-Net和H2Former

定量结果:

image.png

在TN3K、BUSI、CAMUS-LV、CAMUS-MYO和CAMUS-LA上,不同任务特定方法的定量结果总结在表7中。

image.png

在这些最先进的方法中,H2Former在TN3K和CAMUS-MYO上取得了最佳性能,分别获得了82.48%和87.31%的平均Dice分数。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在内的所有5个任务上都取得了更好的性能,平均Dice分数分别为84.45%、85.77%、93.73%、87.46%和91.58%。这验证了SAMUS将SAM适应医学图像领域的有效性。

定性结果:

不同方法的定性分割结果,包括U-Net、AAU-Net、MISSFormer、H2Former和SAMUS,如图4所示。

image.png

从视觉上看,超声图像的分割因其低对比度、不均匀特征和模糊的物体边界而具有挑战性。现有方法难以准确区分目标和背景,导致了大量的假阴性和/或假阳性。

相比之下,SAMUS在保持目标区域的完整性和减少假阳性方面表现出优越性。这归功于SAM框架的固有优势,以及SAMUS引入的具体调整和设计。

泛化能力:

image.png

不同任务特定方法的泛化性能定量比较如图3所示。在比较方法中,H2Former、TransUnet和TransFuse分别在DDTI、UDIAT和HMC-QU上取得了最佳性能。

相比之下,SAMUS在每个数据集上都超过了最佳比较方法,并分别提高了Dice分数,平均增加了7.06%、12.22%和7.42%。在可见和不可见数据集之间的性能比较中,与其他比较方法相比,SAMUS在3个不同的分割任务中都遇到了最小的性能下降。

一个有趣的观察是,在乳腺癌分割任务上,SAMUS在不可见数据集(即UDIAT)上的性能甚至优于可见数据集(即BUSI)上最佳比较方法。这显示了SAMUS在处理未知领域方面的出色泛化能力,展示了它在各种医学图像分割场景中的稳健性和适应性。

4.2 与SOTA基础模型的比较

比较方法:

选择了4种SOTA基础模型进行比较,包括原始的SAM、MedSAM、SAMed和MSA。

定量结果:

为了验证SAMUS作为一个基础模型在各种下游任务上的通用性能,作者在US30K数据集上进行了基础模型的比较。

image.png

如表8所示,SAM,即在SA-1B上训练的模型,在未经调整的情况下在医学图像分割上表现出显著的性能下降。通过在US30K数据集上简单调整SAM的Mask解码器,MedSAM显著改善了SAM的性能。

在比较的基础模型中,MAS是表现最佳的模型,有效提高了SAM在TN3K、BUSI、CAMUS-LV、CAMUS-MYO和CAMUS-LA上的分割性能,平均Dice分数分别增加了53.08%、27.65%、62.77%、53.05%和74.52%。

与MSA相比,SAMUS在上述5个数据集中始终取得了显著的改进,Dice分数分别为83.05%、84.54%、91.13%、83.11%和92%。这验证了SAMUS中的CNN分支和CBA模块的有效性,特别是在补充对于医学图像分割至关重要的局部信息方面。

定性结果:

image.png

图5展示了不同基础模型的定性分割结果,包括SAM、MedSAM、SAMed、MSA和SAMUS。在医学图像中没有经过调整的情况下,SAM完全失去了分割一切的能力。通过将调整方法应用于SAM,MedSAM、SAMed和MSA可以在一定程度上恢复SAM的分割能力。

然而,它们仍然难以在超声图像中准确勾画分割边界,导致大量的假阴性和假阳性。相比之下,SAMUS表现出卓越的性能,能够准确地定位分割边界,即使是低对比度的分割。这与使用图像编码器补充局部信息对于医学图像分割中的边界/形状保持特别有帮助的分析是一致的。

泛化能力:

image.png

不同基础模型在未知领域上的比较总结在图6中。总的来说,在US30K上训练的基础模型在医学图像分割任务中的泛化性能要远远优于原始的SAM。在三对分割任务中,即甲状腺结节分割、乳腺癌分割和心肌分割,所有基础模型在心肌分割上都遇到了严重的性能下降,并且在乳腺癌分割上有很好的泛化性能。

SAMUS在所有三个未知数据集上始终表现出最佳性能,分别为甲状腺结节、乳腺癌和心肌分割的Dice分数为66.78%、78.06%和56.77%。这强调了SAMUS在未知领域的卓越泛化能力,始终明显优于其他基础模型。

部署成本:

作者对SAMUS和其他基础模型在部署效率方面进行了全面评估,包括GPU内存成本、模型参数、计算复杂性、推理速度、分割性能和泛化性能。

为了便于比较,GPU内存在训练期间将批处理大小设置为1时进行测试,并以千兆字节(即G)为单位进行测量。计算复杂性和推理速度以每秒浮点运算(即GFLOPs)和每秒帧数(即FPS)来衡量。分割性能以在所有可见数据集上的平均Dice分数来衡量,泛化性能则基于在所有未知数据集上的平均Dice分数进行评估。所有上述指标都经过归一化,并在雷达图中呈现,如图7所示。

image.png

在比较模型中,SAMed的GPU内存成本、模型参数、计算复杂性和推理速度都最低。然而,它的分割和泛化性能都不如MSA和SAMUS。尽管SAMUS拥有比其他模型更多的参数,但其GPU内存成本和计算复杂性仍然是第二低的,而推理速度是第二快的,表明SAMUS是一个更适合临床的模型。

此外,SAMUS的部署性能与最易于部署的方法(即SAMed)非常接近,并且在分割和泛化性能方面要好得多。

4.3 消融实验

SAMUS中每个组件的有效性:

SAMUS中的4个组件,包括CNN分支、CBA、Feature Adapter和Position Adapter,依次引入原始SAM并在TN3K和BUSI数据集上进行评估。

image.png

如表9所总结的,SAMUS的任何一个组件都能有效提高SAM在医学任务上的分割性能和泛化能力。即使是一个简单的Position Adapter也可以在TN3K、DDTI、BUSI和UDIAT上分别提高SAM的Dice分数分别为50.6%、38.1%、26.77%和30.54%。通过引入局部特征,CNN分支的性能提升明显高于Position Adapter。

此外,仅仅将CNN分支和ViT分支的输出进行融合并不是最佳选择。引入CBA可以进一步促进对局部特征的探索,从而在Dice上相对于CNN分支实现平均增长1.48%和2.11%。

通过结合所有四个组件,SAMUS实现了最佳的分割性能和泛化性能。

不同提示对SAMUS的影响:

为了分析提示的效果,作者评估了在不同点提示下在US30K上训练的SAMUS的性能。总的来说,SAMUS对点的位置和数量具有很强的鲁棒性。对于具有大的类内表征变化的目标(例如甲状腺结节、左心室和心肌),在不同位置的单点提示下,Dice分数的性能变化在1%范围内。

此外,通过引入多点提示,性能可以得到很大的提升。相比之下,对于具有均匀特征的目标(例如乳腺癌和左心房),在不同单点提示下的性能变化在±0.3%的范围内,引入多点提示未必会带来性能提升。一个可能的原因是使用更多的点可能会产生信息冗余或排除。

5、参考

[1]. SAMUS: Adapting Segment Anything Model for Clinically-Friendly and Generalizable Ultrasound Image Segmentation.

作者:AI视界引擎
文章来源:AI视界引擎

推荐阅读

更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
18838
内容数
1371
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息