0-RevSAM2: Prompt SAM2 for Medical Image Segmentation via Reverse-Propagation without Fine-tuning

RevSAM2：通过反向传播无微调提示 SAM2 进行医学图像分割

arXiv:2409.04298v1 [cs.CV] 6 Sep 2024 V1 link_here

arXiv:2409.04298v2 [cs.CV] 25 Nov 2024 V2

背景

SAM2 在自然图像和视频零样本提示分割中表现优异，但应用于医学图像时，因训练数据缺乏医学图像，难以通过语义特征精确分割器官等结构。现有研究多采用大量标记数据微调 SAM2 部分组件以适应医学成像，但存在需大量数据和训练时间、依赖交互式提示及未充分利用 SAM2 架构增强功能等局限。标签不足的医学图像分割旨在解决医学领域标注数据有限问题，其中小样本医学图像分割（FSMIS）是主要研究方向，但多数方法将连续切片视为独立实体进行分割，忽略了切片间语义信息。

• 提出 RevSAM2 框架，在无需微调的情况下利用 SAM2 的记忆银行和传播能力，通过反向传播策略选择高质量查询信息，实现对未见 3D 医学图像的分割。
• 验证 RevSAM2 在有限标签场景下的有效性和优越性，与现有方法对比展示其在多器官数据集上的优势。

实验方法

整体框架

基于 SAM2 构建 RevSAM2 框架，给定 3D 查询医学图像和少量 2D 支持图像及对应分割标签，预测查询图像的体素级分割图。包括前向传播、反向传播和查询自传播三个阶段。

RevSAM2 的总体框架 （a） 和正向传播和反向传播的图示 （b）。为了评价 S 和 Y 正向传播到 qi 上得到的预测 pi 的质量，将 qi 和 pi 反向传播回 S 得到 Yi，并计算 Yi 和 Y 之间的平均骰子 πi 并将其作为评估 pi 准确性的指标。

前向传播

用图像编码器和提示编码器分别编码支持图像和标签，经记忆编码器融合信息存入记忆银行。对于查询图像的切片，通过记忆注意力机制基于支持图像信息获得融合视觉特征，再由掩码解码器预测分割掩码，但由于 SAM2 的空间位置扰动，分割质量可能不佳。

反向传播

以查询切片及其预测掩码为支持图像 - 标签对，反向计算支持图像的预测 Dice 分数，评估查询切片预测的准确性。具体通过类似前向传播的方式，从查询切片和预测掩码中获取特征，经记忆注意力模块处理后由掩码解码器生成支持图像的预测掩码，计算其与真实标签的平均 Dice 分数。根据分数选择高质量的查询切片作为条件切片用于后续传播。

查询自传播

将条件查询切片及其预测掩码的特征存入记忆银行，并维护一个 FIFO 队列存储最近的非条件查询切片及其预测掩码的特征。对于非条件查询切片，通过记忆注意力机制结合条件切片信息获得融合视觉特征，进而预测其掩码。最终将条件掩码和非条件掩码组合形成查询图像的最终预测掩码。

查询 self propagation 的图示。在查询自传播中，内存库持续存储通过反向传播选择的条件切片的特征，同时维护 FIFO 队列以在内部查询推理期间存储非条件切片的特征。

实验结果

在 BTCV、AbdomenCT - 1K、Synapse - CT 和 CHAOS - MRI 四个多器官医学图像分割数据集上进行实验，每个数据集进行单器官分割实验，随机选择少量切片作为训练集或支持图像模拟标签不足场景，采用平均 Dice 相似系数（mDSC）和平均归一化表面 Dice（mNSD）作为评估指标。

与多种方法对比，包括从头训练的方法（如 nnU - Net、Swin UNETR）、通用小样本分割模型（UniverSeg）、SAM - 基微调方法（如 SAMed、H - SAM 等）以及传统小样本方法（如 AAS DCL、SR&CL 等）。对 SAM2 进行不同提示方式（点、框、掩码）的测试。

BTCV 和 AbdomenCT 数据集：RevSAM2 在标签不足情况下（每个器官 10 个切片）表现出色，相比其他方法有显著提升，如在 BTCV 数据集上 mDSC 提升 10.03%，在 AbdomenCT 数据集上提升 12.18%。在域适应实验中，使用不同数据集的支持图像时，RevSAM2 也取得最佳性能。

Synapse - CT 和 CHAOS - MRI 数据集：RevSAM2 在这两个数据集上也达到了 SOTA 的 mDSC 性能，在脾脏分割上提升明显，分别提升 14.89%（Synapse - CT）和 9.45%（CHAOS - MRI）。

消融实验

验证方法中利用查询体积自身信息（查询自传播）和通过反向传播选择高质量查询预测的有效性。结果表明，仅使用前向传播（基线）或随机选择查询信息效果不佳，使用查询信息但不区分质量（前向传播加查询信息）会因低质量掩码影响后续分割；基于反向传播选择条件切片的方法效果最佳。还研究了支持图像数量和选择的条件切片数量 k 对性能的影响，发现支持图像增多时，基于 π 值选择高质量 P 更有效，当支持图像为 1 时，π 值与实际 Dice 分数相关性降低，k 值减小会使性能逐渐下降。

总结

提出 RevSAM2 框架，在数据稀缺场景下无需微调即可进行医学图像分割，通过反向传播选择高质量查询预测作为掩码提示在查询内传播，优于现有小样本算法，为 SAM2 在医学图像分割中的应用开辟新方向，有望为医学图像分割提供经济高效的解决方案。