迈向分割的大一统！OMG-Seg：一个模型搞定所有分割任务-51CTO.COM

本文经自动驾驶之心公众号授权转载，转载请联系出处。

写在前面&笔者个人思考

图像分割已经从单任务分割走到了语义分割、实例分割、全景分割三种分割任务的统一；大模型以及多模态的发展又带来了文本和图像统一，使得跨模态端到端成为可能；追求更高级、更全面任务似乎成了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。

方法设计：

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

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

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

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

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

开放词汇和多数据集分割 任务的制定与先前的图像和视频分割相同。然而，这个设置超越了固定的标签空间。特别是，它需要在各种数据集上进行开放集识别。同时，多数据集分割要求一个模型在不同数据集下分割更多的概念。作为一种常见做法，本文采用CLIP文本嵌入作为掩模分类器，避免了分类学冲突，并同时实现了开放集识别。因此，本文测量了视觉查询特征与类别嵌入之间的距离，而不是学到的分类器。

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

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

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

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

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

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

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

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

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

Training and Inference（多数据集&多任务放一起怎么训～～）

联合图像视频数据集联合训练 与首先在图像数据集上进行预训练不同，本文的目标是仅联合训练所有分割任务一次。所有训练目标都是三种不同情况下的一个实体标签和掩模。实体可以是物体、杂项、类别无关的掩模及其相应的标签。注意，具有相同ID 的实例掩模形成tube掩模。在训练过程中，本文在预测的实体掩模和地面实体掩模之间应用匈牙利匹配，以将目标查询分配给视频/图像实体，然后监督它们的预测掩模和分类。为了避免跨数据集的分类冲突，分类器被替换为CLIP文本嵌入。最终的损失函数表示为。这里，是掩模分类的交叉熵（CE）损失，和分别是分割的掩模交叉熵（CE）损失和Dice损失。

通用推断 对于图像分割，本文遵循Mask2Former的相同推断过程。例如，对于 PS，本文根据排序后的分数合并物体和杂项。分数由 CLIP 文本嵌入生成。对于视频分割任务，对于 VIS 和 VPS，为了生成实例ID，本文遵循先前的工作，使用查询匹配而不是引入额外的跟踪组件。对于 VOS 任务，本文在第一帧和其余帧之间采用掩模匹配。对于交互式分割任务，本文遵循原始的 SAM，通过提供框和点提示来获取二进制掩模。对于开放词汇分割任务，由于本文有一个冻结的 CLIP 编码器，本文将掩模池化分数和学习分数与开放词汇嵌入合并。

结合任务以获取更多应用 由于本文的模型可以执行各种分割任务，结合交互、开放词汇和图像/视频分割任务可以导致一些新的应用。例如，本文可以结合交互和视频分割，实现灵活的基于提示驱动的视频目标分割。或者本文可以将交互分割与开放词汇设置结合起来，从而实现开放词汇交互分割。

实验结果：

数据集和评估指标 与常规设置不同，本文的目标是尽可能地探索在多个数据集上进行联合训练。本文使用COCO panoptic、COCO-SAM、VIPSeg 和Youtube-VIS-2019（YT-VIS19）作为训练数据集。除了封闭集测试，本文还包括使用Youtube-VIS-2021、ADE-20k和DAVIS-2017数据集进行开放词汇（OV）推理，它们的注释在训练过程中未被使用。COCO-SAM 是通过使用真实边界框，掩模中心点作为视觉提示来创建的。注释是通过COCO全景掩模获得的。此外，本文还在表3中包括了多数据集设置，以验证本文的OMG-Seg的多数据集联合训练的有效性。除了表2之外，本文添加了更多数据集，包括ADE-20k 和YT-VIS21 用于联合训练。本文为每个数据集使用相应的评估指标，包括 PQ、掩模 mAP、VPQ、tube mAP、J&F和 mIoU。

实验细节 本文在MMDetection中实现本文的模型和所有其他基线。本文使用32个A100 GPU的分布式训练框架。每个小批次在每个GPU上有一张图像。对于数据增强，本文采用大规模的抖动，与先前的工作一样，以构建强大的基线。对于每个表中的所有模型，本文采用相同的训练步骤。本文使用OpenCLIP来初始化骨干网络，并用其相应的文本嵌入替换学到的分类器。对于图像输入，将它们视为伪视频，通过将两个图像及其掩模连接成一个。本文采用不同的采样率来平衡每个数据集的训练样本。本文报告冻结和训练后的骨干网络的结果供参考。

Main Results

虽然在COCO图像分割基准上表现略有下降，但在VIPSeg数据集上取得了接近最先进水平的结果；开放词汇视频分割领域效果也不错。

基于K-Net和 Mask2Former进行了多数据集联合训练。联合共同训练通常可以提高大多数视频分割数据集的性能，从而实现了大幅度的模型参数减少（从1326M减少到221M）。这种改进在三个VPS和VIS数据集中是一致的，无论是否冻结了骨干网络。

结论：

本文介绍了第一个用于图像、视频、开放词汇和交互式分割的联合共训练框架。本文的解决方案OMG-Seg是一个新颖而简单的框架，它使用统一的查询表示和共享的解码器来处理多样化的任务。首次有可能训练一个单一的分割模型，能够在十个不同的任务上表现出与特定任务模型相媲美的性能。这种方法显著减少了参数大小和在各种应用程序的模型设计中需要专门工程的需求。预计本文高效且多功能的框架将成为多任务和多数据集分割的稳健基线。

原文链接：https://mp.weixin.qq.com/s/Hz2gDOy5POxU_bFyJCxUDw

迈向分割的大一统！OMG-Seg：一个模型搞定所有分割任务