CVPR2024｜迈向分割的大一统！OMG-Seg：一个模型搞定所有分割任务（南洋理工）

写在前面&笔者个人思考

图像分割已经从单任务分割走到了语义分割、实例分割、全景分割三种分割任务的统一；大模型以及多模态的发展又带来了文本和图像统一，使得跨模态端到端成为可能；追求更高级、更全面任务似乎成了CV/NLP领域的共识。听起来很玄乎，对于本文的统一所有分割来说，做法也很好理解，本质上就是在Mask2former的基础上增加了SAM的提示来支持prompt驱动以及交互式分割，同时增加CLIP encoder来支持开放域，给query加ID使其同时支持图像和视频的分割；

PS：不得不说，Mask2former还是一个很强的基线，当前SOTA的大一统分割模型Oneformer是在这个基础上改进的，本文的超大一统方案也基于这个base进行升级；

论文的主要思路

本文统一解决了各种分割任务，提出的OMG-Seg，即一种模型，足够出色，能够高效而有效地处理所有分割任务， 包括图像语义分割、实例分割、全景分割，以及它们的视频对应任务、开放词汇设置、由提示驱动的交互分割（如SAM）和视频目标分割 。据本文所知，这是第一个能够处理所有这些任务并取得令人满意性能的模型。本文展示了 OMG-Seg，这是一个基于Transformer的编码器-解码器架构，具有任务特定的查询和输出，可以支持十多种不同的分割任务，并显著减少跨各种任务和数据集的计算和参数开销。本文在共同训练期间严格评估了任务间的影响和相关性。代码和模型可在 https://github.com/lxtGH/OMG-Seg 获取。

主要贡献

引入OMG-Seg，一个模型应对所有分割任务。结构上：

• 先前的模型通常使用共享的视觉骨干，有几个特定任务的分支，但OMG-Seg采用了共享的编码器-解码器架构。
• 将所有任务的输出统一为一个统一的查询表示。一个查询可以表示一个掩模标签/图像或掩模/唯一的ID/视觉提示。然后，采用一个共享的解码器处理所有类型的查询及其特征。这个设置促进了通用分割训练和推理的统一，并实现了跨任务的广泛参数共享。

效果上：

• 通过在合并的图像和视频数据集上进行共同训练，OMG-Seg能够处理多达十种不同的分割任务，涵盖不同的数据集。OMG-Seg是首个在图像、视频、开放词汇和交互分割四种设置上在八个不同的数据集上取得了不错的结果，包括COCO、ADE20k、VIPSeg、Youtube-VIS-2019、Youtube-VIS-2021和DAVIS-1。

相关工作：

通用图像/视频分割 视觉Transformer的出现引发了通用分割领域的创新浪潮。最近的研究开发了基于端到端集合预测方法的掩模分类架构，在图像和视频分割任务中表现出色，超过了专业模型。尽管取得了这些进展，大多数现有方法仍然依赖于不同的模型来处理不同的分割任务和数据集。最近，有一种趋势是在各种数据集和任务之间训练单一模型，以获得参数共享的好处。例如，OneFormer在一个单一模型中集成了三个图像分割任务，而UNINEXT专注于统一实例级任务。类似地，TarVIS使用目标提示组合各种视频分割任务。然而，这些现有工作中没有一个彻底研究过图像、视频和由提示驱动的数据的联合训练，以构建一个全面的分割模型。本文的工作是朝着这个方向的首次尝试，拓展了跨这些领域的共同训练的潜力。

视觉基础模型 视觉基础模型领域的最新研究展示了优化技术的多样化，涵盖了各种学习范式。这些包括仅视觉预训练策略、联合视觉-语言预训练方法以及包含视觉提示的多模态框架。一个值得注意的例子是SAM，它展示了广泛培训在实现通用分割方面的泛化性和可扩展性。在此基础上，Semantic-SAM通过添加语义标签和增加细粒度级别的方式增强了SAM模型。然而，尽管它们的印象深刻，这些视觉基础模型通常在视频分割任务上表现不佳，需要在更动态的情境中进行进一步的优化以实现最佳性能。

开放词汇分割 这一系列的视觉分割研究旨在识别和分割超出有限封闭集视觉概念的新目标。许多研究利用视觉语言模型（VLMs）提供的可转移表示，在训练过程中探索区域和文本表示之间的对齐。在推理阶段，检测器可以使用从VLMs派生的文本嵌入来识别新的类别。本文的模型遵循这一理念，以实现开放词汇分割。具体来说，本文使用冻结的VLMs同时作为特征提取器和分类器。这种策略允许在开放词汇设置中实现无缝过渡。

统一建模 Transformer架构的适应性使得在不同模态之间共享基本模块成为可能。这种多功能性启发了几个研究倡议，使用一个通用的Transformer框架处理不同领域。值得注意的是，在视觉通用领域方面的努力主要集中在统一视觉领域内不同任务。例如，Pix2Seq系列通过自回归标记预测的方式实现任务统一。类似地，Unified-IO实现了一个序列到序列的流水线，将各种输入和输出转换为离散的标记序列。此外，最近的进展探索了在上下文中学习作为结合各种视觉任务的手段。这些方法主要针对跨领域的任务统一。然而，统一分割模型与为特定目的构建的分割模型之间的性能差距仍然是一个待解决的问题。

方法设计：

OMG-Seg是一个单一而多才多艺的模型，具有减少任务特定定制和最大参数共享的特点，可以支持各种分割任务，使其成为一个适用于所有分割任务的模型。目标不是追求每个任务的最先进结果，而是增加一个可广泛推广的分割模型的建模能力，同时允许任务之间的广泛知识共享。

统一任务表示（看看都支持哪些任务～～）

图像分割 给定输入图像 ，图像分割的目标是输出一组掩模 ，其中 表示二进制掩模 的类别标签， 是掩模数量， 是空间大小。根据类别标签和掩模的范围，本文报告了三种不同的分割任务的结果，包括语义分割 (SS)、实例分割 (IS) 和全景分割 (PS)。PS是SS和IS的统一，其中包含可计数的thing类别和不可计数的stuff类别。对于这三个任务，本文采用掩模分类架构，其中每个掩模对应一个语义标签。

视频分割 给定视频剪辑输入 ，其中 表示帧数，视频分割的目标是获取一个掩模tube ，其中 是tube掩模数量， 。 表示tube掩模 的类别标签， 表示每个tube掩模的实例ID。每个tube掩模可以分类到可计数的thing或不可计数的stuff类别，其中thing类别还分配了一个唯一的 。对于stuff掩模，默认情况下跟踪为零。当 且任务只包含stuff类别，且所有thing类别没有ID时，VPS变成了视频语义分割 (VSS)。如果 重蛋且 只包含 thing类别，且所有stuff类别被忽略，VPS变成了视频实例分割 (VIS)。视频目标分割 (VOS) 的目标是跟踪第一个框架的掩模，而不进行分类。受图像分割的启发，本文还采用了tube掩模分类架构来训练和沿时间维度链接短管。对于VOS，本文采用类别不可知的 tube掩模训练，这类似于VPS和VIS。目标分割 (VOS) 的目标是跟踪第一个框架的掩模，而不进行分类。受到图像分割的启发，本文还采用了tube掩模分类架构来训练和沿时间维度链接短管。对于VOS，本文采用了类别不可知的tube掩模训练，这与VPS和VIS类似。

交互式分割 SAM框架中的交互式分割同时使用图像 和视觉提示 (例如点和框) 作为输入，输出相应的二进制图像掩模 ， 其中 是视觉提示的数量。每个视觉提示被编码成一个目标查询，这自然地可以成为解码器的输入。在本文的实验中，本文使用共享的解码器处理所有不同任务查询。

所有的事物都在查询中 如上所述，通过结合所有不同的设置，本文可以使用相同基于查询的掩模分类框架表示所有输出分割实体。特别是，一个目标查询对应一个掩模 ，标签 和 ID 根据不同的任务设置， 和 的格式和范围是不同的。然而， 和 的格式和范围是相似的。因此，将所有这些任务放入一个共享的编码器和解码器框架中，并共同训练一个模型用于所有分割任务是很自然的。

OMG-Seg架构（具体是怎么做的～～）

OMG-Seg遵循Mask2Former的架构设计。如图所示，它包含一个骨干网络、一个像素解码器和一个掩模解码器。区别在于以下几个方面，包括冻结骨干的设计、合并的目标查询 (包括目标查询和视觉提示)，以及共享的多任务解码器。根据不同的任务设置，解码器输出相应的掩模和标签。

使用VLM编码器作为冻结骨干 为了实现开放词汇的识别，在骨干部分，本文采用冻结的 CLIP视觉模型作为特征提取器。本文使用ConvNeXt架构来自OpenCLIP。给定图像/视频输入，VLM编码器提取多尺度的冻结特征 ，以供进一步处理。

像素解码器作为特征适配器 像素解码器与Mask2Former相同，包含多阶段的可变形注意力层。它将冻结特征 转换为具有相同通道维度的融合特征 ，其中 是特征的层索引， 是最高分辨率的特征。

合并的目标查询 如上所分析，每个目标查询代表一种掩模输出。然而，从功能的角度来看，图像、视频和交互模式代表不同的属性。对于图像，目标查询侧重于目标级别的定位和识别。对于视频，目标查询可能涉及时间一致性，例如同一目标跨越不同帧。对于交互式分割，目标查询被迫定位特定区域。对于图像和视频输入，本文采用目标查询来表示图像掩模或被跟踪的tube掩模。因为两者都需要语义标签。本文将它们称为语义查询， 。对于交互模式，按照SAM 的做法，本文采用提示编码器将各种视觉提示编码成与目标查询相同形状的查询。本文将其称为位置查询， 。因此，本文可以共享变压器解码器的相同接口。

共享的多任务解码器 其主要操作是交叉注意力，接收合并的目标查询 和 和图像/视频特征 ，并输出精炼的目标查询。最终的掩模是通过精炼查询和高分辨率特征 的点积获得的。对于图像语义级别的任务，本文采用了与Mask2Former相同的过程。具体而言， 使用多尺度特征 进行掩码交叉注意力 [18]。 是查询，而 是键和值。然后，对精炼的查询应用多头自注意力 (MHSA) 层。

对于视频任务，本文采用相同的交叉注意力设计。唯一的区别是金字塔特征 沿着时间维度与3D位置嵌入连接，这是先前的工作中的默认设置。合并的视频特征和精炼的查询用于预测tube掩模。

对于交互式分割，本文执行相同的交叉注意力设计。然而，本文跳过自注意力，以避免在 MHSA层中的掩模查询之间发生交互，因为交互式分割只关心输入视觉提示区域。在获取精炼的目标查询后，它经过一个预测FFN，通常包括一个带有ReLU激活层和线性投影层的3层感知器。所有的查询都受到掩模分类损失和掩模预测损失的监督。解码过程以级联的方式进行，每个特征金字塔有三个阶段。共享的多任务解码器 其主要操作是交叉注意力，接收合并的目标查询 和