超越SOTA：PP-SAM 在有限数据集上的图像分割突破，简化采样 SA M 过程，仅需最小的标注！

AIGC 先锋科技

发布于 2024-07-08 14:19:36

1050

发布于 2024-07-08 14:19:36

分段任何模型（SAM）最初是为通用分段任务设计的，最近被用于息肉分割。然而，用来自新成像中心或诊所的数据对SAM进行微调带来了重大挑战。这是因为这需要创建一个昂贵且耗时的标注数据集，同时推理过程中用户提示的变异性也带来了可能性。为了解决这些问题，作者提出了一种健壮的微调技术，PP-SAM，它允许SAM在有限图像的情况下适应息肉分割任务。为此，作者利用变量扰动边界框提示（BBP）来丰富学习上下文，并在推理过程中增强模型对BBP扰动的鲁棒性。在息肉分割基准上的严格实验表明，作者的变量BBP扰动显著提高了模型的韧性。值得注意的是，在Kvasir数据集上，1次射击微调分别在使用50像素和100像素BBP扰动进行推理时，将DICE分数提高了20%和37%。此外，作者的实验显示，在推理过程中使用50像素扰动的1次射击、5次射击和10次射击PP-SAM分别比最近的一项尖端（SOTA）息肉分割方法提高了26%、7%和5%的DICE分数。作者的结果激励了作者的PP-SAM在有限样本的其他医学成像任务中的更广泛应用。作者的实现可在https://github.com/SLDGroup/PP-SAM获取。

1 Introduction

基于深度学习算法在结肠镜检查过程中检测癌前病变已显示出巨大潜力。近期，一种基础模型，即Segment Anything Model（SAM），被引入用于通用语义分割。一些研究探索了其在息肉分割中的零样本推理或微调[17, 9]潜力。

然而，当仅使用一个成像中心/诊所的数据对SAM进行微调时，由于其泛化能力有限，对于可能存在分布外数据的不同中心/诊所进行微调至关重要。此外，由于在推理过程中用户提示可能不准确，复杂性进一步增加。

由于内镜医师使用的提示具有主观性，存在因人为因素（如疲劳、经验以及当天检查的病例数量）导致的变异性。当内镜医师使用的边界框提示大于息肉感兴趣区域的大小时，SAM表现不佳。

因此，开发一种对不准确（受干扰）边界框提示具有鲁棒性的适应方法至关重要。

为了解决这些问题，作者研究了针对息肉分割的SAM微调，即_PP-SAM_，用于处理_变量边界框提示干扰_。通过对结肠镜图像上的SAM进行微调，作者展示了其在息肉分割方面的卓越性能，证明了其有望在多样化临床环境中提升结直肠癌筛查和诊断。作者的方法简化了微调过程中数据标注所需的时间、成本和资源，使其对多中心息肉分割有效。

作者的主要贡献如下：

PP-SAM框架：作者引入了_PP-SAM_，这是一个在有限数据情况下用于息肉分割的新的基于SAM的鲁棒适应框架。作者还探索了SAM不同模块的迁移学习能力。
变量提示干扰：作者提出了一种非常简单但有效的策略，即在微调过程中使用_变量边界框提示干扰_，以使模型对提示干扰更鲁棒。
鲁棒性分析：作者对零样本和少样本SAM在推理过程中对边界框提示干扰的_鲁棒性_进行了严格分析。作者的实验结果表明，零样本SAM对提示干扰非常敏感，因此作者提出的基于干扰提示的适应策略的PP-SAM可以显著提高模型在推理过程中的鲁棒性。

本文的其余部分组织如下：第3部分描述了作者的方法论。第2部分解释了相关的前期工作。第4部分介绍了实验设置和结果。最后，第5部分总结了本文。

2 Related Work

在本节中，作者分为三个部分描述相关工作：分割任何模型（SAM）、SAM在医学图像分割中的应用，以及SAM在息肉分割中的应用。

Segment anything model

在[8]中，SAM的作者提出了一种基础模型，通过引入一个可提示的分割任务、一个分割模型以实现零样本迁移到各种任务，以及一个新的图像分割数据集。SAM的想法来源于自然语言处理（NLP）领域，在该领域，预训练在大型数据集上的大型语言模型（LLM）已经显示出强大的零样本性能[4]。这类大型模型已经显示出在训练期间未见过的任务和数据集上的泛化能力[3, 8]。SAM的工作[8]表明，这种在极大规模数据集上的训练同样可以转化为计算机视觉领域，以分割各种不同的图像类型，包括医学图像。

SAM in medical image segmentation

尽管SAM的表现强劲，但由于其在自然图像上的预训练，它在非分布领域（如医学成像）上的表现并不理想。通过在来自不同解剖学和模态的19个医学成像数据集上评估SAM[12]，参考文献[12]的作者进行了一项实验研究，以确定SAM在医学成像中应用的可能性。作者发现，基于单一提示的SAM的性能高度依赖于数据集和任务，因此得出结论，SAM在一些医学成像数据集上表现出色的零样本分割性能，而在其他数据集上表现不佳[12]。

Auto-SAM[17]用一个在相同输入图像上操作的编码器替换了SAM对 Mask 或一组点的条件作用。引入这个编码器使得Auto-SAM[17]在多个医学图像分割基准测试中获得了最先进（SOTA）的结果。MedSAM[11]则通过设计一个用于医学图像分割的基础模型，采用了一个包含超过一百万张图像的精选数据集[11]。MedSAM在医学图像分割方面也超越了现有的SOTA基础模型，甚至胜过了一些专业模型[11]。

Polyp segmentation

结肠息肉是结肠癌的重要前驱病变，因此正确分割结肠息肉可以减少结肠癌的误诊[9]。已经提出了许多息肉分割的方法，但有限的结肠镜图像数量仍然是一个主要挑战。像SAM这样的方法，可以在没有大量数据的情况下进行分割，因此看起来非常有吸引力。

在[9]中，作者提出了Polyp-SAM，这是针对息肉分割的SAM模型的微调版本。Polyp-SAM在五个数据集上达到了SOTA或接近SOTA的性能，从而显示了SAM在医学图像分割任务中的有效性。在[2]中，作者使用了名为Polyp-SAM++的文本提示辅助的SAM，与未经提示的SAM相比，它被证明更稳健、更精确[2, 11]。在SAM模型的前期工作中（医学成像中的SAM模型和息肉分割中的SAM模型），基于原始SAM模型，并使用不同的方法比原始SAM取得了更好的分割结果，以实现更专业任务的SOTA结果，而仅靠自然图像的预训练可能不足以实现这一点。然而，这些方法中没有考虑到在将SAM应用于息肉分割时，现实生活中的提示固有不准确（扰动）的问题。

3 Methodology

图1展示了作者提出的PP-SAM框架。首先，作者输入一个小的带标签图像数据集。然后，作者从相应的真实分割 Mask （GT）中提取边界框。接着，作者使用作者的_可变边界框提示扰动_方法对边界框进行扰动。最后，作者使用带有GT Mask 和扰动边界框提示的数据集来微调SAM [8]。下一部分将描述主要组件。

Prompts

在本小节中，作者描述了作者提出的用于微调的可变扰动提示，以及在推理过程中用于鲁棒性分析所使用的扰动。

3.1.1 Variable perturbed prompts for fine-tuning

虽然SAM可以使用各种提示，但作者主张采用边界框提示，因其简单性。作者的方法包括使用_可变（扰动）边界框提示_对SAM进行微调，以进行息肉分割。作者的可变扰动将边界框随机地从0扩展到n个像素，在所有四个方向上，如图1中的框（2）所示。当作者使用这种策略对SAM进行微调时，扰动的变化增强了模型在现实生活推理过程中对提示扰动的鲁棒性。

3.1.2 Prompts during inference

为了评估作者方法的鲁棒性，作者在推理过程中对不同 Level 的固定扰动（0、5、10、...、95和100像素）在边界框（所有边）上评估模型的性能。例如，10像素的扰动意味着在推理过程中将边界框在所有边上扩展10像素。

SAM architecture

SAM [8] 是一个基础的图像分割模型，可以响应各种提示（例如，点、框、 Mask ）。尽管在广泛的SA-1B数据集上训练，SAM显示出强大的零样本泛化能力。SAM包含三个关键组件：图像编码器、 Mask 解码器和提示编码器。下面将描述这些组件。

3.2.1 Image encoder

SAM图像编码器（图1中（3）框）基于视觉 Transformer （ViT）主干网络[5]；它接收高分辨率（即）的图像作为输入，并生成16倍下采样的图像嵌入（即）。

3.2.2 Prompt encoder

SAM提示编码器（图1中的(4)框）利用了两套提示：稀疏的（边界框、点、文本）和密集的（ Mask ）提示。同时，它还使用位置编码和学习到的嵌入来编码点和框。

3.2.3 Mask decoder

SAM使用一个轻量级的 Mask 解码器（图1中的(5)框），它包括一个动态 Mask 预测和一个交并比（IoU）得分回归头。

Transfer learning

作者研究了SAM中不同组件的迁移学习能力：图像编码器、提示编码器和 Mask 解码器。图2中的实验结果表明，微调图像和提示编码器就足够了；因此，在作者的所有实验中，作者都保持 Mask 解码器冻结不变。

Limited data settings

在本工作中，作者随机选择不同的较小数据集（图1中的(1)框）来对SAM进行息肉分割的微调。具体来说，作者用随机选择的图像微调PP-SAM，其中和，以及整个数据集。

4 Experiments

在本节中，作者描述了数据集、评估指标、实现细节以及实验结果。

Datasets

作者使用Kvasir [7] 数据集对SAM进行少样本息肉分割的微调。这个数据集包含1,000张息肉图像。遵循[6]，作者采用Kvasir中的相同900张图像作为训练集，剩下的100张作为测试集。为了评估作者提出的解码器的泛化能力，作者使用了三个未见过的测试数据集，分别是 ClinicDB [1]，EndoScene [19]，和 ColonDB [18]。ClinicDB 包含从31个结肠镜视频提取的612张图像。EndoScene 和 ColonDB 分别包含60张和379张图像。作者三个未见过的测试集的图像与训练集显著不同，因为它们是从不同的医院/诊所/中心收集的，这些中心的采集设备和程序各不相同。### 评估指标

在所有实验中，作者使用DICE相似度分数作为评估指标。DICE相似度分数衡量的是重叠准确性，适用于不平衡数据的二值分割。由于息肉分割是一个不平衡的息肉（病变）和背景区域的二值分割任务，作者倾向于使用DICE分数来评估PP-SAM在息肉分割上的性能，其中对解剖结构的精确描绘对于诊断和治疗计划至关重要。 GT Mask 和预测 Mask 的DICE分数在公式1中定义：

Implementation details

作者使用Pytorch 1.11.0实现并微调了作者的PP-SAM，在单个配备48GB内存的NVIDIA RTX A6000 GPU上运行。在实验中，作者将_可变边界框提示扰动_的最大长度设置为。作者将图像调整为，并将边界框重新缩放以匹配新的图像分辨率。作者使用AdamW优化器[10]，学习率和权重衰减率均为0.0001。作者不使用任何数据增强和学习率调度器。

在微调过程中，作者优化了加权交叉熵和平均交并比（mIoU）的复合损失函数。除非另有说明，否则使用ViT-B的SAM（SAM-B）以批量大小为1微调100个周期；作者根据在所有边上有30个像素推理边界框扰动情况下的DICE分数保存最佳模型。作者报告了五次运行的平均DICE相似度分数，以评估作者的微调性能。作者使用测试图像的原始分辨率来计算DICE分数。

Results

为了评估作者提出的PP-SAM的性能和鲁棒性，作者进行了如下六组不同的实验。

4.4.1 Transfer learning capabilities of different modules of SAM

作者实证探索了SAM在不同四种实验配置下的迁移学习能力。关于推理过程中边界框提示扰动的影响，这些调查的结果如图2所示。在作者的分析中，策略性地冻结 Mask 解码器（即仅微调图像和提示编码器）被证明是最有效的方法，获得了最高的DICE分数。这种卓越性能可能源于避免了过拟合，这种情况在用有限数据集微调 Mask 解码器时可能会发生。相反，保持图像编码器冻结会使模型更容易受到提示扰动的攻击。当图像和 Mask 解码器都冻结时，性能的下降更为显著，这强调了它们在模型适应性中的集体重要性。由此，作者坚定地提倡仅冻结 Mask 解码器，选择性地微调图像和提示编码器，以优化迁移学习的效率和模型的鲁棒性。

4.4.2 Effectiveness of variable bounding box prompt perturbations during fine-tuning

在图3中，作者展示了在微调过程中不同边界框提示扰动的影响。作者评估了0（无扰动）、10、20、30、40、50像素的扰动以及原始边界框提示四周0-50像素范围内的随机扰动下的DICE得分。结果显示，在训练期间没有提示扰动的模型容易受到更大的推理扰动影响。随着训练扰动的增大，对这些扰动的韧性有所提高。然而，使用可变扰动（0-50像素）进行微调的模型在小型和大型推理扰动下都表现出更好的整体性能。作者认为，在训练期间对不同侧面施加可变扰动可以提高模型针对各种边界框提示扰动 Level 的鲁棒性。

4.4.3 Learning ability of PP-SAM for polyp segmentation on Kvasir dataset

在图4中，作者展示了在测试过程中将零样本和少样本微调技术应用于Kvasir数据集的结果。如图所示，随着推理过程中边界框扰动大小的增加，DICE分数明显下降，这一趋势与作者的预期一致。值得注意的是，在整个推理过程中，作者微调后的模型在面对这些即时扰动时显示出增强的鲁棒性。实施了作者随机1次和50次微调与_变量扰动边界框提示_显著提高了模型的鲁棒性，相比于零样本微调，在所有边的100像素扰动情况下，分别将DICE分数提高了37%和60%。作者还可以得出结论，DICE分数从1次微调提高到50次微调，而在50次微调之后改进效果差异不大。

4.4.4 Generalizability of PP-SAM for unseen polyp segmentation

图6：在未见过的EndoScene测试集上的实验结果。作者使用从Kvasir训练集中随机抽取的图像训练的模型进行这些实验。

图5展示了在ClinicDB测试集上未见过的息肉分割评估结果，其中作者的微调方法显示了显著的性能提升，即相对于零样本学习，1次样本和50次样本的DICE分数分别提高了24%和43.5%。在图6中，作者可以看到在EndoScene数据集上未见过的息肉分割的DICE分数也有类似的提升。具体来说，作者的1次样本和50次样本微调分别比在推理时带有100像素扰动的零样本学习提高了19%和50%的DICE分数。

图7展示了在未见过的ColonDB数据集上的结果，作者的1次样本和50次样本微调在推理时带有100像素扰动的情况下，分别比零样本学习提高了27%和45%。作者观察到，在所有未见过的测试集上，经过20到50次样本微调后，DICE分数的提升最小。

4.4.5 Scalability of SAM with ViT encoders

图8报告了带有ViT-B（基础）和ViT-L（大型）图像编码器的SAM模型在零样本和微调方面的结果。从这个图中，作者可以得出结论，带有ViT-L的SAM在零样本息肉分割方面超过了ViT-B。然而，两种ViT编码器微调后的模型显示出相似的性能。因此，这些结果可以用来论证在现实世界场景中部署模型时需要谨慎考虑，在这些场景中可能会出现提示的扰动，并强调模型对这些变化的鲁棒性的重要性。

4.4.6 Performance comparison with SOTA

作者报告了在图9中作者微调后的SAM和一个最新方法PVT-CASCADE [13]的结果。从条形图中可以看出，采用可变边界框扰动进行微调的作者的PP-SAM在最多50次射击（即，PP-SAM需要更少的标注数据就能达到接近最优性能）上显著优于PVT-CASCADE。更准确地说，使用作者的方法对SAM进行1次射击（74.5%）、5次射击（77%）和10次射击（81.6%）的微调，比在推理时使用50像素边界框扰动的PVT-CASCADE分别高出26%、7%和5%的DICE分数。当使用25像素扰动时，作者的方法性能进一步提高（即1次、5次和10次射击分别提高32%、11%和9%）。因此，对于涉及有限精确标注 GT 分割 Mask 的多息肉分割，作者的PP-SAM优于非提示型方法。