[论文解读] Beyond Adapting SAM: Towards End-to-End Ultrasound Image Segmentation via Auto Prompting

Xian Lin, Yangyang Xiang+2 more|arXiv (Cornell University)|Sep 13, 2023

Advanced Neural Network Applications被引用 29

一句话总结

SAMUS 在 SAM 的基础上扩展了一个并行 CNN 编码器、适配器和跨分支注意力，以实现端到端的超声分割，具备更好的泛化能力和更低的部署成本，并在 US30K 上得到验证。

ABSTRACT

End-to-end medical image segmentation is of great value for computer-aided diagnosis dominated by task-specific models, usually suffering from poor generalization. With recent breakthroughs brought by the segment anything model (SAM) for universal image segmentation, extensive efforts have been made to adapt SAM for medical imaging but still encounter two major issues: 1) severe performance degradation and limited generalization without proper adaptation, and 2) semi-automatic segmentation relying on accurate manual prompts for interaction. In this work, we propose SAMUS as a universal model tailored for ultrasound image segmentation and further enable it to work in an end-to-end manner denoted as AutoSAMUS. Specifically, in SAMUS, a parallel CNN branch is introduced to supplement local information through cross-branch attention, and a feature adapter and a position adapter are jointly used to adapt SAM from natural to ultrasound domains while reducing training complexity. AutoSAMUS is realized by introducing an auto prompt generator (APG) to replace the manual prompt encoder of SAMUS to automatically generate prompt embeddings. A comprehensive ultrasound dataset, comprising about 30k images and 69k masks and covering six object categories, is collected for verification. Extensive comparison experiments demonstrate the superiority of SAMUS and AutoSAMUS against the state-of-the-art task-specific and SAM-based foundation models. We believe the auto-prompted SAM-based model has the potential to become a new paradigm for end-to-end medical image segmentation and deserves more exploration. Code and data are available at https://github.com/xianlin7/SAMUS.

研究动机与目标

• 推动普遍性同时对临床友好的超声分割。
• 在利用 SAM 的基础之上解决医学领域的局限性。
• 提升对超声目标的局部细节捕获与边界划分。
• 降低 GPU 需求并实现对入门级硬件的部署。

提出的方法

• 保留 SAM 的提示编码器和掩码解码器；为医学用改造图像编码器。
• 添加并行 CNN 分支，通过跨分支注意力注入局部特征。
• 引入位置适配器和五个特征适配器，将 ViT 适配到医学领域及更小的输入。
• 实现重叠的补丁嵌入以保留边界信息。
• 在冻结核心 SAM 组件的同时，使用带有适配器、CNN 分支和跨分支注意力的训练；使用 Dice + BCE 损失并结合单点正提示。

Figure 1: Structure comparison of different SAM-based foundation models for medical image segmentation.

实验结果

研究问题

• RQ1 在不对整个模型重新训练的情况下，是否可以有效地将 SAM 自适应用于超声分割？
• RQ2 局部特征集成和紧凑的适配器是否改善边界保留与小目标分割？
• RQ3 与任务特定模型及其他基础模型相比，SAMUS 在可见（训练）和不可见的超声数据集上的表现如何？

主要发现

• SAMUS 在 TN3K、BUSI、CAMUS-LV、CAMUS-MYO、CAMUS-LA 等数据集上，平均 Dice 分数高于最先进的任务特定方法（平均 Dice：84.45%、85.77%、93.73%、87.46%、91.58%）。
• 与基础模型的比较中，SAMUS 在可看数据上始终提供更高的 Dice 分数（例如 83.05–93.73 区间），并在不可见数据上具有更好的泛化性（例如在不可见数据上的基线平均提升）。
• 消融实验表明四个组件（CNN 分支、跨分支注意力、特征适配器、位置适配器）均对性能有贡献，其中仅位置适配器就显著提升了若干数据集的 Dice。
• SAMUS 将 GPU 内存成本降至全量 SAM 训练的约 28%，推理速度大约提高 3 倍，同时保持或提升分割与泛化。

Figure 2: Overview of the proposed SAMUS.