CVPR 2024｜RobustSAM：在低质量图像上鲁棒地分割一切

科研机构：Snap Research, 台湾大学

论文：https://arxiv.org/abs/2406.09627

项目：https://robustsam.github.io/

简介

Segment Anything Model (SAM)是专用于图像分割的基础模型。SAM能快速识别和分割照片中的元素，支持用户通过输入提示（prompt）的方式，进而选择感兴趣区域进行分割。SAM借鉴了自然语言处理领域的Foundation Model理念，使用提示学习来适应不同的分割任务，展示了零样本学习（zero-shot learning）的能力。然而，虽然SAM有强大的图像切割能力，但其在处理恶劣天气或低光源下所拍摄的图像时，性能却会有显著的下降：

因此，本论文的目的在于提升SAM模型面对低质量图像时的鲁棒性。我们提出了RobustSAM的模型，在维持SAM模型强大的零样本学习能力的前提下，以少量的额外参数量显着提升了在低质量图像上的切割表现。以下表格显示了SAM和RobustSAM在可学习参数量及运算资源上的差异：

方法

Overview of RobustSAM

上图显示了大致的训练流程，总共可分为两部分：SAM和RobustSAM。RobustSAM包含原始SAM模型的模块和预训练参数，在上图中以灰色模块表示，在训练过程中不会被更新。同时，在RobustSAM中引入了新的网络模块（以紫色模块表示），这些模块则会在训练过程中被更新。

在训练的一开始，会先将清晰（clear）图像通过人工合成的方式生成一张低质量（degrade）图像。接着，将清晰图像输入SAM并获得mask feature以及token feature，它们将用来生成最终的分割结果。此外，本篇的方法也同时提取了图像编码器的low-level feature，并且会与前述的mask feature以及token feature一起参与consistency loss的计算。

接下来，会将一开始合成的低质量图像输入到RobustSAM模型，并同样获得对应的feature。由于此时的输入为低质量图像，因此可以预期输出feature中会包含许多影响分割任务的degradation信息。为了去除这些degradation信息，本文使用了两个模块：AMFG模块和AOTG模块。通过降低SAM所输出feature和RobustSAM所输出feature之间的一致性损失，这两个模块将会得到有效的训练。

Anti-Degradation Mask Feature Generation (AMFG)

该模块的目的在于移除mask feature中有关degradation的信息。我们可以将该阶段的处理流程分成两部分来说明。

• 第一部分称为适应性的特征整合（Adaptive Feature Integration）。由于可以将degradation信息视为风格信息，因此会使用实例正则化（Instance Normalization, IN）来处理输入特征并移除degradation信息。为了补偿IN可能造成的细节损失，作者还使用了批次正则化（Batch Normalization, BN）。不同正则化所输出的特征会经过一个注意力模块来决定其重要性，进而融合出最终的输出特征。
• 为了进一步精炼图像特征，AMFG模块的第二部分使用了傅里叶变换来将特征从空间域投射至频率域。有鉴于风格或degradation信息通常隐藏在振幅分量中，因此此处使用了卷积网络来进行处理。经过卷积处理的振幅分量会与原始的相位分量融合在一起，并通过逆傅里叶变换得到输出特征。

Anti-Degradation Output Token Generation (AOTG)

为了有效的还原输出token的特征，AOTG模块使用了由IN层和多层感知器（MLP）层所组成的轻量模块。由于token feature的维度远远小于mask feature的维度，因此该轻量模块已足以有效的移除token中的degradation信息。

实验结果

此篇论文做了许多实验来证明其有效性。

量化比较（Quantitative Comparison）

Seen dataset with synthetic degradation

Unseen dataset with synthetic degradation

Unseen dataset with real degradation

从以上的结果可以得知，无论处理合成的degradation又或是现实世界中的degradation，RobustSAM都取得了优于其它方法的表现。

可视化结果分析（Qualitative Comparison）

上述结果清楚显示了比起其它现有的图像还原方法，RobustSAM在分割表现上有着显着的提升。

基于SAM模型的下游任务

本文中同时也比较了SAM及RobustSAM在下游任务（除雾及去模糊）上的增强效果：

由此可见，在使用分割结果作为下游任务的额外先验信息时，RobustSAM对于最终的增强效果是高于SAM的。

消融实验(Ablation Study)

SAM模块的微调分析 & RobustSAM 架构有效性分析

从消融实验结果中可以得到以下结论：

1. 直接微调SAM模型中的模块参数反而会导致表现比原本来得差，微调的参数量越大，表现就降低得越多。这是由于灾难性遗忘（catastrophic forgetting）所致，在学习新数据的同时，之前学到的有用信息也却被覆盖了。
2. RobustSAM中新增的不同模块对最终表现均有一定的提升帮助，通过这种冻结原有SAM模型参数，仅更新额外模块参数的训练方式，成功进一步提升分割表现。

结语

本文提出了RobustSAM模型用以改善原本SAM在低质量图像上，分割表现不佳的问题。通过额外新增小型的网络模块，RobustSAM大大提升了其分割能力的鲁棒性，同时也保留了SAM零样本分割的能力。

公众号后台回复“数据集”获取100+深度学习各方向资源整理