点击下方卡片,关注 「AI视界引擎」 公众号
分段任何模型(SAM)已成为图像分割领域的一种变革性方法,因其强大的零样本分割能力和灵活的提示系统而受到赞誉。然而,其性能在处理质量退化的图像时受到挑战。
为了解决这一限制,作者提出了鲁棒分段任何模型(RobustSAM),它在提高低质量图像上SAM性能的同时,保留了其提示能力和零样本泛化能力。作者的方法利用了预训练的SAM模型,仅增加了少量的参数和计算需求。
RobustSAM的附加参数可以在8个GPU上30小时内优化,展示了其对于典型研究实验室的可操作性和实用性。作者还介绍了RobustSeg数据集,这是一个包含68.8万个图像- Mask 对的数据集,这些图像经历了不同的退化,旨在最佳地训练和评估作者的模型。
在跨越各种分割任务和数据集的广泛实验中,RobustSAM的性能得到了显著确认,尤其是在零样本条件下,突显了其在广泛实际应用中的潜力。
此外,作者的方法已被证明可以有效提高基于SAM的下游任务(如单图像去雾和去模糊)的性能。
1 Introduction
精确的图像分割对于机器人、增强/虚拟现实以及内容创作等众多下游应用至关重要。Segment Anything Model(SAM)[32]在野外图像分割领域开启了新篇章:利用包含超过十亿个标注 Mask 的全面SA-1B数据集,SAM能够泛化到大量物体,并且仅凭最小化的提示(从点标注到边界框)即可准确分割场景。这种创新方法通过无缝适应各种应用,彻底改变了零样本分割。
随着SAM在多样化分割任务中展示其多功能性,人们的关注转向了它在面对复杂和具有挑战性的情境时的鲁棒性和可扩展性。特别是,提高其在降质图像上的鲁棒性仍然是待探索的前沿。大量文献指出,随着成像降质(如低光照、噪声、模糊、恶劣天气和压缩伪影)的出现,SAM的性能会下降。这些降质显著影响了SAM生成的分割 Mask 的质量,直接影响依赖于这些 Mask 的下游任务。特别是,最近的图像恢复工作,如去雾[28]、去模糊[44]和超分辨率[53],一直在使用这些 Mask 或潜在特征作为结构先验,这种先验可以泛化到未见场景。然而,这些工作假设即使在降质条件下,SAM也能产生可靠和准确的 Mask 。如果SAM在这些情况下的鲁棒性受到损害,集成先验知识的益处将受到限制,从而限制它们在现实世界场景中的应用。
为了应对这一挑战,一种直观的方法是利用现有的图像恢复技术[10, 34, 59]在将图像输入SAM之前进行预处理。尽管这些方法能在一定程度上改善图像质量,但不能保证所选的图像恢复技术能提高图像分割[8, 13, 43, 68, 74]的性能。这是因为大多数图像恢复算法针对的是人类视觉感知,而不是SAM等分割模型的具体需求。
另一种策略涉及直接在降质图像上对SAM进行微调。然而,直接调整SAM解码器或整合新的解码器模块可能会深刻影响模型在零样本任务上的泛化能力。此外,盲目地用降质图像对SAM进行微调可能导致灾难性遗忘,即网络无意中失去了从原始清晰图像中学到的知识[18, 33]。
为此,作者引入了RobustSAM,它在处理降质图像的同时保持零样本功能,实现了鲁棒性。作者的方法提出了两个新颖的模块:抗降质标记生成模块和抗降质 Mask 特征生成模块。在由原始SAM从清晰图像对提取的特征一致性损失的监督下,这些模块旨在提取降质不变分割特征。作者还对SAM的原始输出标记进行了微调,使其适应作者的鲁棒分割方法。通过在训练过程中冻结SAM的原始模块,所提出的方法在提高处理降质图像的能力的同时,保持了其在零样本分割上的有效性。
此外,RobustSAM中提出的附加模块可以高效地训练。与需要在数百个GPU上训练的原始SAM相比,RobustSAM可以在八台A100上30小时内完成训练。这使得RobustSAM变得易于访问,准备好被集成到各种应用场景中。广泛的实验表明,作者的RobustSAM在清晰和降质情境下都表现出色。此外,作者发现RobustSAM通过提供更鲁棒的先验,增强了在降质情境下的基于SAM的下游任务,如单图像去雾和去模糊的效果。
为了增强RobustSAM的能力和鲁棒性,作者引入了Robust-Seg数据集。Robust-Seg结合了来自7个现有数据集的43K个精心标注的图像。每张图像都经过了15种精心建模的合成降质处理,从而在Robust-Seg中形成了包含688K图像的综合集合。这个广泛的数据集旨在推动图像分割的边界,并为未来的研究提供一个有价值的资源。
总结来说,作者的贡献如下:
- 作者提出了RobustSAM,这是一个建立在Segment Anything模型之上的零样本分割模型,它对各种图像降质具有增强的鲁棒性。这种增强的鲁棒性显著提高了下游应用的表现。
- 作者构建了Robust-Seg数据集,这是一个包含688K具有不同降质的图像- Mask 对的集合。作者希望Robust-Seg能为降质图像上的分割模型设立一个新的基准。
2 Related Work
Segment Anything Model
段 Anything Model (SAM) [32] 在图像分割方面取得了前所未有的性能,推动了计算机视觉[54]各个子领域的发展。SAM 接受直观的提示,如点或边界框,展示了在多种分割任务和新图像分布上的卓越零样本迁移学习能力。它的适应性在包括医学成像[95, 70, 20, 66]、伪装目标检测[70]和显著目标分割[54]在内的各个领域得到了证明。除了其分割能力之外,SAM 在增强计算机视觉任务方面发挥着基础性作用,包括语义分割[95, 66, 5]、图像编辑[89]和视频目标跟踪[82],[12]。尽管 SAM 展现出了巨大的潜力,但其性能受到图像质量差[26, 62, 65, 75]的挑战,影响了分割和下游任务的准确性。
Robust Segmentation
在自动驾驶和监控分析领域,许多研究发现当处理降质图像时,基于CNN的分割性能会下降,这导致了各种补救方法的开发。例如,QualNet[31]通过可逆编码方案探索了质量无关的特征提取,而URIE[68]处理多种图像损伤,通过分类约束提高分割稳定性。同时,FIFO[38]推动分割框架通过雾透过滤机制学习抗雾特征。然而,这些技术主要关注单一类型的降质,可能在对多种图像降质的情况下缺乏鲁棒性。此外,与下游任务联合训练的这种策略可能会稀释SAM的零样本优势。
Image Restoration
在图像恢复领域,诸如SRCNN [14] 等针对单一退化类型的方法开创了卷积神经网络在提高图像质量方面的应用。随后,不同领域的创新成果不断涌现,在超分辨率(SR)、去噪、去雾、去雨、水下增强[52, 81]和去模糊[17, 35, 36]等方面取得了显著的成功。MPRNet [91] 和 HINet [7] 等方法尝试通过单一网络解决多种退化问题。最近,基于 Transformer 的方法也在图像恢复任务中获得了关注[47, 76, 90]。然而,尽管像All-in-One [42],IPT [4],和AirNet [40] 这样的多退化方法提供了更大的灵活性和性能改进,但它们旨在提升人类的视觉质量,而不是改善下游任务(如分割)的性能。
3 Proposed Method
Preliminary: Segment Anything Model
作者提供了SAM框架[32]的简洁概述。如图2下半部分所示,SAM包含了三个关键组件:图像编码器、提示编码器和 Mask 解码器。图像编码器使用视觉 Transformer (ViT)处理输入图像。提示编码器处理稀疏提示输入(如点、框和文本)和密集输入( Mask ),将它们转换为适当的表示。Mask 解码器是一个修改后的Transformer解码块[73]。它将图像和提示嵌入与输出 Token 结合,生成 Mask 特征。这个过程包括提示自注意力以及提示和图像嵌入之间的双向交叉注意力。值得注意的是, Mask 解码器使用转置卷积来创建详细的 Mask 特征。每个 Mask 的输出 Token ,来自 Token 到图像的注意力,通过MLP转换为一个动态分类器。当这个分类器与 Mask 特征相乘时,得到最终的分割 Mask 。
Robust Segment Anything Model
作者提出了RobustSAM,它在解决图像退化的同时保留了SAM的零样本学习能力。与那些对SAM进行微调或联合训练复杂适配模块的标准方法不同,RobustSAM采用了一种简约而精心增强的方法。
3.2.1 Model Overview
训练。 为了训练RobustSAM,作者首先对清晰的输入图像应用退化增强,然后将退化的图像输入到RobustSAM中。最初,模型利用其图像编码器从这幅退化图像中提取特征。与原始SAM框架不同,作者对输出标记进行了微调,现在称为健壮输出标记(ROT)。这个ROT,连同提示标记和图像编码器提取的特征,通过原始SAM层处理,生成 Mask 特征(退化)和每个 Mask 的健壮输出标记。
AOTG块处理以提取抗退化的信息,将其转换为。同时,AMFG块精炼图像编码器早期和最终层的 Mask 和补充特征(和),去除与退化相关的信息,生成精炼特征(和)。遵循[30]中提出的架构,特征融合块将这些精炼特征结合到作者最终的健壮 Mask 特征中,以提高分割质量。
与此同时,原始清晰图像通过标准SAM处理,提取清晰版本的补充特征(), Mask 特征()和输出标记()。这些清晰特征与RobustSAM的精炼特征之间的连贯性损失确保了与未退化图像输出的对齐。然后,将退化输入的分割结果与使用分割损失函数的 GT 进行比较。
在作者的退化增强方法中,作者包括了15种退化和一个恒等映射。这确保了清晰图像保持其质量,避免了在非退化情况下的性能下降。
推理。 在推理过程中,仅使用RobustSAM(图2上半部分)来生成分割 Mask 。
以下,作者将详细讨论所提出的反退化输出标记生成和反退化 Mask 特征生成模块。
3.2.2 Anti-Degradation Mask Feature Generation
如图3所示,输入特征首先通过实例归一化(IN)处理。受到先前工作[72, 68, 25]的启发,IN的目的是标准化与图像退化相关的变化。直观地说,这移除了风格属性,同时保留了核心内容。这一步骤对于减轻个别图像失真的影响至关重要,确保在多种退化条件下内容的稳定性。与此同时,受到[68]的启发,作者增加了一个应用批量归一化(BN)的分支。BN至关重要,因为它解决了IN过程中可能导致的细节损失,如[68, 58]所指出的。
然后,作者将BN和IN单独生成的特征进行合并。一个注意力机制仔细检查合并后的特征以生成注意力图,这些图动态地衡量每种特征类型的重要性,从而合成一个结合了两种归一化技术优势的特征集[68]。为了补偿可能丢失的语义信息,这个增强的特征集沿着通道维度与原始输入特征进行拼接。此外,作者整合了通道注意力,类似于挤压和激励方法(SEC)[22, 46],以自适应地细化特征融合。
受到[24, 79, 84, 86, 88]的启发,作者引入了傅里叶退化抑制模块,通过使用傅里叶变换将特征从空间域转换到频率域来增强整合后的特征。这种技术利用振幅分量来捕捉图像退化的风格信息。通过应用1x1卷积,作者专注于隔离和去除退化元素。同时,保留相位分量以保持结构完整性。接下来,逆傅里叶变换将精细的特征带回到空间域。这个过程将退化视为图像风格,并生成鲁棒的分割不变的退化特征。该模块应用于先前模块生成的两个特征:互补特征(退化)和 Mask 特征(退化)。为了确保这些精细特征与使用清晰图像输入时SAM模型提取的相应特征(即和)保持一致,作者采用了 Mask 特征一致性损失()。
通过最小化的每一部分,作者确保精细特征与在清晰图像条件下提取的特征保持一致,从而保证在不同退化下特征的鲁棒性和一致性。
3.2.3 Anti-Degradation Output Token Generation
抗退化输出标记生成模块旨在细化每个 Mask 对应的鲁棒输出标记(),以去除与退化相关的信息。与传统的 Mask 特征不同,的主要功能是确保分类边界的清晰度,因此包含较少的纹理信息。因此,作者发现使用轻量级模块来过滤对退化敏感的信息是足够的。如图3右侧所示,该模块使用多层实例归一化,然后是一个MLP层。这种策略旨在保持计算效率,同时确保模型能够从受退化影响的输入中恢复鲁棒 Mask 信息。经过细化的标记与在清晰输入条件下由原始SAM提取的输出标记进行比较,以计算标记一致性损失,
这种损失确保了经过细化的标记与在清晰图像条件下提取的标记保持一致。经过MLP处理后,输出与鲁棒 Mask 特征结合,生成最终的 Mask 。
3.2.4 Overall Loss
整体损失函数整合了 Mask 特征一致性损失()、标记一致性损失()和分割损失(),形成对模型的综合惩罚。整体损失表达为:
这里, 是一个组合分割损失,它结合了Dice [50] 和 Focal损失 [69]:
其中 是预测的 Mask , 是真实 Mask ,- 是用于对不同损失加权的超参数。这个复合损失函数旨在确保提高分割质量的同时增强模型对抗退化的鲁棒性。
4 Implementation Details
Dataset
为了训练和评估RobustSAM,作者构建了一个全面的Robust-Seg数据集,包含68.8万个图像和 Mask 对。这个数据集由来自几个现有数据集的图像组成,具体包括LVIS [19],ThinObjectSk [48],MSRA10K [11],NDD20 [71],STREETS [67],FSS-1000 [45],以及COCO [51]。在这个数据集中,作者融入了原始清晰图像及其经过15种合成降质处理的版本,这些降质包括模糊、噪声、低光照、恶劣天气条件等等。这种方法确保了模型在广泛的图像质量上进行训练,并对各种图像质量具有鲁棒性。
在训练过程中,作者使用了MSRA10K、ThinObject-5k和LVIS的全部训练集(及其增强版本)。MSRA10k和LVIS的测试集(及其增强版本)用于验证模型的分割准确性。为了挑战模型的零样本泛化能力,作者将其测试在NDD20、STREETS、FSS-1000和COCO数据集的全范围图像(及其增强版本)上。
此外,作者还使用完整的BDD-100k [87]和LIS [6, 80]数据集进行了广泛的测试,这些数据集包含多种现实世界的降质,如低光照、模糊、雨和雪。这种方法确保了对RobustSAM在实际场景中的性能和在面对现实世界应用中通常遇到的恶劣环境条件下的鲁棒性进行了彻底的评估。
Training Detail
在RobustSAM的训练阶段,作者保持预训练的SAM参数冻结,只专注于优化提出的鲁棒性模块。作者使用基于点的提示来训练RobustSAM。
RobustSAM显著提升了分割质量,并旨在实现快速高效的训练。以0.0005的学习率训练40个周期,在8个Nvidia A100 GPU上完成130,000次迭代仅需30小时。《表1》详细列出了RobustSAM与SAM在训练和推理性能方面的比较。RobustSAM不仅提供了改进的分割结果,而且与SAM相比,训练效率也显著提高。
Evaluation Protocol
为了评估RobustSAM的性能,作者采用了几种指标:交并比(IoU)、Dice系数(Dice)、像素准确度(PA)以及在不同阈值水平下的平均精度(AP)。
Performance Evaluation
在具有挑战性条件的图像分割领域,作者的RobustSAM框架与几种现有方法进行了比较,以证实其有效性。作者将其与基础的SAM模型以及一种策略性的两阶段方法进行了基准测试,在这种两阶段方法中,图像首先通过通用的图像恢复技术以提升输入质量,随后再进行由SAM驱动的分割。为此,作者引入了AirNet [40],这是一种针对未知降质量身定制的最先进的通用视觉质量增强方法。此外,作者还整合了URIE [68],这是一种为更有效分割设置舞台的图像恢复方法。另外,作者还与HQ-SAM [30]进行了比较,后者是原始SAM的高质量迭代版本。
在已知数据集上的比较。作者在已知数据集LVIS [19]和MSRA10K [11]上评估了Robust-Seg框架下作者提出方法的性能。具体来说,作者是在这些数据集的测试集上评估这些方法。结果显示,作者的方法在处理由这些多样化场景提出的各种挑战时,表现出卓越的性能。此外,RobustSAM在显著提高降质场景下的分割效果的同时,在清晰场景中保持了或提升了性能。详细的 结果在表2和表3中呈现,证明了作者方法在不同分割场景下的有效性。
在图4中的零样本分割结果表明,RobustSAM在零样本分割中具有显著的鲁棒性,在不同降质情况下保持了高性能。
Ablation Study
为了进一步理解作者贡献的影响,作者进行了一项消融研究。所有实验都是在BDD-100k [87] 和 LIS [6, 80] 数据集上进行的。
微调SAM吗? 作者以各种配置微调了SAM模型:微调整个模型、解码器以及输出标记。结果展示在表7中。观察到微调整个SAM模型或其解码器会大大降低其零样本能力,导致性能显著下降。仅微调输出标记带来了性能提升;然而,与RobustSAM相比,它们仍然逊色。
所提模块的有效性。 此外,作者验证了每个所提模块的有效性,包括抗退化 Mask 特征生成模块(AMFG)、带傅里叶退化抑制的ADM模块(AMFG-F)、抗退化输出标记生成(AOTG)以及鲁棒输出标记(ROT)。表7中呈现的发现表明,每个引入的模块都积极影响了RobustSAM的整体性能,其中AMFG模块显示出最大的增强效果。
Improving SAM-prior Tasks
为了验证在图像质量下降的情况下,作者的RobustSAM是否能提升基于SAM先验的应用性能,作者选择了单幅图像去雾[28]和单幅图像去模糊[44]作为测试案例。遵循这些任务原始论文的设置,作者使用SAM和RobustSAM作为它们的先验,并在去雾任务的SOTS数据集[41]和去模糊任务的GoPro数据集[57]上评估了它们的性能。表8和图5的结果表明,在下游任务中使用RobustSAM可以获得更优的性能。这种提升可以归因于RobustSAM在退化图像上提高了分割准确性,为这些任务提供了更鲁棒的先验。
6 Conclusion
本文介绍了RobustSAM,该模型在各种退化情况下分割图像的能力卓越。
模型的优势源于其组件——特别是抗退化 Mask 特征生成模块、抗退化输出标记生成和鲁棒输出标记模块。
为了验证RobustSAM的有效性,作者提出了一个大规模的数据集,名为Robust-Seg。
此外,作者证明了RobustSAM的优势扩展到了改善基于SAM的任务,如去雾和去模糊,证实了其在退化条件下进行图像处理的可靠工具价值。
其性能为零样本分割的鲁棒性设立了新的标准,为未来的研究提供了有希望的方向。
参考
[1].RobustSAM: Segment Anything Robustly on Degraded Images.
点击上方卡片,关注 「AI视界引擎」 公众号