[论文解读] MaskSAM: Towards Auto-prompt SAM with Mask Classification for Volumetric Medical Image Segmentation

Bin Xie, Hao Tang+4 more|arXiv (Cornell University)|Mar 21, 2024

Advanced Neural Network Applications被引用 4

一句话总结

MaskSAM 提出了一种无需提示、具备掩码分类感知能力的分割一切模型（SAM）在3D体素医学图像分割中的适配方法。通过引入可学习的提示生成器以生成辅助掩码、边界框和分类器标记，并将3D深度卷积与深度MLP适配器集成到SAM的Transformer模块中，该方法在AMOS2022数据集上实现了90.52%的Dice分数，相比nnUNet提升了2.7%。

ABSTRACT

Segment Anything Model (SAM), a prompt-driven foundation model for natural image segmentation, has demonstrated impressive zero-shot performance. However, SAM does not work when directly applied to medical image segmentation, since SAM lacks the ability to predict semantic labels, requires additional prompts, and presents suboptimal performance. Following the above issues, we propose MaskSAM, a novel mask classification prompt-free SAM adaptation framework for medical image segmentation. We design a prompt generator combined with the image encoder in SAM to generate a set of auxiliary classifier tokens, auxiliary binary masks, and auxiliary bounding boxes. Each pair of auxiliary mask and box prompts can solve the requirements of extra prompts. The semantic label prediction can be addressed by the sum of the auxiliary classifier tokens and the learnable global classifier tokens in the mask decoder of SAM. Meanwhile, we design a 3D depth-convolution adapter for image embeddings and a 3D depth-MLP adapter for prompt embeddings to efficiently fine-tune SAM. Our method achieves state-of-the-art performance on AMOS2022, 90.52% Dice, which improved by 2.7% compared to nnUNet. Our method surpasses nnUNet by 1.7% on ACDC and 1.0% on Synapse datasets.

研究动机与目标

• 解决SAM在医学图像分割中的局限性，包括缺乏语义标签预测能力以及对人工提示的依赖。
• 通过预训练的2D SAM实现零样本、无需提示的3D医学图像分割推理。
• 通过将SAM架构适配至3D空间理解，弥合2D自然图像与3D医学体数据之间的领域差距。
• 提出一个统一框架，无需用户提供的提示即可同时预测精确掩码及其对应的语义标签。
• 在无需从头开始进行任务特定微调的情况下，实现在主要3D医学分割基准上的最先进性能。

提出的方法

• 设计一种提示生成器，从SAM图像编码器中多层级特征图生成辅助二值掩码、边界框和分类器标记。
• 在掩码解码器中引入可学习的全局分类器标记，将其与辅助分类器标记相加，以预测每个掩码的语义标签。
• 将3D深度卷积适配器（DConvAdapter）注入图像编码器中每个Transformer模块，以利用3D归纳偏置处理3D图像嵌入。
• 在掩码解码器中引入3D深度MLP适配器（DMLPAdapter），用于处理3D提示嵌入，采用倒瓶颈结构与跳跃连接以提升特征精炼能力。
• 在图像编码器和掩码编码器中均应用深度位置嵌入（DPosEmbed），以增强3D空间建模能力。
• 使用预测结果与真实标签之间的二分图匹配计算损失，支持具有可变长度多类别掩码的端到端训练。

实验结果

研究问题

• RQ1 无需人工点或框提示，能否实现对SAM的无提示适配，以实现高性能的3D医学图像分割？
• RQ2 在原始设计仅支持二值掩码预测的前提下，如何扩展SAM以预测分割掩码的语义标签？
• RQ3 为使预训练的2D SAM有效提取并利用体素医学图像中的3D空间信息，需要哪些架构修改？
• RQ4 与标准适配器相比，3D感知适配器（DConvAdapter与DMLPAdapter）在提升3D泛化能力与分割精度方面的表现如何？
• RQ5 具备可学习提示与分类器标记的统一端到端框架，在3D医学图像分割基准上能否显著超越现有SOTA方法（如nnUNet）？

主要发现

• MaskSAM在AMOS2022数据集上达到90.52%的Dice分数，相比nnUNet提升2.7%。
• 在ACDC数据集上，MaskSAM的Dice分数比nnUNet高出1.7%，表明其在心脏MRI分割任务中具备强大泛化能力。
• 在Synapse数据集上，MaskSAM性能相比nnUNet提升1.0%，证实其在多种解剖结构上的鲁棒性。
• 消融实验表明，采用平均边界框的提示生成器（MaskAvgBBoxPG）性能最佳，相比仅使用掩码的生成器提升1.9%。
• DMLPAdapter与DConvAdapter的集成使性能相比基线提升0.3%，其中DMLPAdapter在特征精炼方面优于其他深度MLP的放置方式。
• 引入深度位置嵌入（DPosEmbed）相比基线性能提升超过0.15%，验证了其在3D空间建模中的关键作用。