摘要
我们提出了 SegGPT，一个在上下文中执行“分割一切”的通用模型。我们将各种分割任务统一为一个通用的上下文学习框架，通过将不同类型的分割数据转换为相同的图像格式来兼容这些差异。SegGPT 的训练被建模为一个带有随机颜色映射的上下文着色问题。训练目标是根据上下文完成多样的任务，而非依赖特定的颜色。训练完成后，SegGPT 可通过上下文推理（in-context inference）执行图像或视频中的任意分割任务，如实例、区域、部件、轮廓和文本等。SegGPT 在多种任务上进行了评估，包括小样本语义分割、视频目标分割、语义分割和全景分割。实验结果显示其在域内和跨域目标分割中都展现出强大的能力，无论是定性还是定量评估。
1. 引言
分割是计算机视觉中最基本的问题之一，其目标是在像素级上定位并重组有意义的概念，如前景、类别、目标实例等。近年来，针对各类分割任务（如前景分割 [41]、交互式分割 [51, 34]、语义分割 [32, 28, 54, 39]、实例分割 [18, 11, 2, 48] 和全景分割 [23, 5, 8]）的算法取得了显著进展，表现越来越精确且快速。
然而，这些专用的分割模型通常局限于特定的任务、类别、粒度和数据类型。当面临不同场景（如分割新概念，或从图像迁移到视频）时，往往需要重新训练一个新模型。这不仅代价高昂，还难以适应大量不同分割任务的需求，缺乏可持续性。
在本研究中，我们的目标是训练一个单一模型，能够处理多样且无限的分割任务。这一目标面临两个主要挑战：(1) 如何在训练中融合截然不同的数据类型，如部件、语义、实例、全景、人物、医学图像、航空图像等；(2) 如何设计一种可泛化的训练机制，区别于传统的多任务学习方式，具备任务定义的灵活性，并能处理跨领域任务（out-of-domain tasks）。
为了解决这些挑战，我们提出了 SegGPT——一个在上下文中执行“分割一切”的通用模型。我们将分割视为一种通用的视觉感知格式，并将不同的分割任务统一到一个通用的上下文学习框架中 [46]。该框架通过将各种分割数据转换为相同的图像格式，使模型可以处理不同类型的数据。SegGPT 的训练被建模为一个上下文着色问题，对每个样本使用随机颜色映射。模型的目标是根据上下文对类别、实例、部件等区域进行着色。通过使用随机颜色，模型被迫依赖上下文信息来完成任务，而非依赖预定义颜色，这使训练方式更具灵活性与泛化能力。其余训练流程与 [46] 保持一致，使用基础 ViT [42] 和简单的 smooth-ℓ1 损失函数 [17]。
训练完成后，SegGPT 能够在图像或视频中通过少量示例进行上下文推理，完成多种分割任务，如目标实例、区域、部件、轮廓、文本等。为了更有效地整合多个上下文示例，我们提出了一种简单但有效的上下文集成策略——特征集成（feature ensemble），该策略在多示例提示设置下能提升模型性能。此外，SegGPT 在不更新模型参数的前提下，也可以作为专用模型使用，通过调整特定的提示即可适配具体场景，例如在 ADE20K 数据集上进行语义分割。
我们的主要贡献如下：
- 首次展示了一个通用模型，可以自动执行一系列多样的分割任务；
- 无需微调，我们直接评估了预训练的 SegGPT 在多个任务上的表现，包括小样本语义分割、视频目标分割、语义分割和全景分割；
- 实验结果表明，SegGPT 在域内和跨域目标分割方面均具有出色的泛化能力，无论是定性还是定量。
需要指出的是，本研究的目标并不是在所有基准测试中超越已有的专用方法或达到新的 SOTA（state-of-the-art）性能，因为我们认为这并不是一个通用模型的责任所在。
2. 相关工作
2.1. 图像分割
分割是计算机视觉中的一个基础问题，其目标是在像素级别定位和组织有意义的概念。分割任务的类型根据概念的定义而异，例如foreground、category或object instance。例如，semantic segmentation [55] 涉及对图像进行像素级的语义分类，而instance segmentation [30] 则旨在识别不同的object instance及其类别。video object segmentation [52, 37, 12] 的任务是：在给定视频第一帧中目标mask的情况下，在整个视频序列中对特定目标进行分割。
以往的分割方法 [32, 28, 54, 39, 18, 11, 2, 48, 23, 5, 8] 通常是针对特定任务设计的，不能灵活地切换任务或更换类别。本文提出了一个通用接口，该接口与所有分割任务兼容，通过恰当的训练方案，一个单一的通用模型就能在in-domain和out-of-domain的分割任务上都获得良好的性能，无论是从定性还是定量的角度看。
2.2. 通用视觉模型
近年来，人们尝试使用基于Transformer的模型来统一视觉领域中的不同任务，从而诞生了一些通用视觉模型（vision generalists）[6, 7, 56, 33, 24]。DETR [5] 是最早使用Transformer [42] 作为目标检测任务特定head的模型之一。Pix2Seq系列 [6, 7] 将视觉任务的输出空间定义为离散的token序列，并以自回归方式执行目标检测、instance segmentation、关键点估计以及图像描述等任务。Unified-IO [33] 和 OFA [45] 采用sequence-to-sequence方式对视觉、视觉-语言和NLP任务进行联合建模，其中输入和输出都被定义为离散token的序列。UViM [24] 将panoptic segmentation、深度估计、图像着色等像素标注任务统一到一起，但为每种任务单独训练模型。
尽管这些工作看起来是在将不同任务统一到相似的空间中，实际上它们仍通过某种形式的硬指示符（如特殊token）来完成每个任务，从而难以推广到新任务。与之相对，本文提出的方法使用in-context框架，保持了任务定义的灵活性，并利用随机上色方案来避免模型退化为多任务学习解决方案，而是通过引用上下文信息来完成指定任务。另一个区别在于任务范围：本工作主要聚焦于视觉感知中的一个核心类别，即图像分割。
2.3. In-Context视觉学习
GPT-3 [3] 将in-context learning引入深度学习，使得一系列NLP任务可以在给定prompt和示例的条件下被表述为文本补全问题。在计算机视觉领域，[1] 首次提出一种使用离散token进行图像修复的in-context训练框架，训练数据来自视觉论文中的图表和信息图，展示了其在前景分割、单目标检测和图像着色方面的能力。Painter [46] 采用基于连续像素的masked image modeling方法，在七个多样且具有挑战性的视觉任务上进行in-context训练，使用有监督数据集，并在这些任务上取得了竞争性的结果。
我们的工作构建在Painter框架之上，但专注于分割任务，因为其在视觉感知中扮演着核心角色。因此，本工作统一了多样的分割数据，包括semantic segmentation、instance segmentation、part segmentation，甚至是特殊场景（如航空图像）下的分割数据。此外，我们设计了一种随机上色方案，迫使模型必须引用上下文信息才能完成任务，从而避免退化为多任务解法。由于分割任务和数据集的变化性小于深度/姿态估计，因此更容易共享内部结构以实现对in-domain任务的有效训练，同时保持对out-of-domain分割任务的泛化能力。
3. 方法
SegGPT 是 Painter [46] 框架的一个特殊版本，它能够通过一个通用的Painter对一切进行分割，因此命名为SegGPT。该训练框架将视觉任务的输出空间重新定义为“图像”，并将不同任务统一为同一种图像修复问题，即：随机遮盖任务输出图像并重建缺失像素。为了保持简洁性与通用性，我们对架构与损失函数不作修改，即使用原始的ViT [13] 和简单的smooth-ℓ1 [17] 损失，但在in-context训练中设计了一种新的随机上色方案，以提升模型的泛化能力。
3.1. In-Context Coloring
在传统的Painter框架中，每个任务的颜色空间是预定义的，这会导致模型的解法退化为多任务学习。例如，对于semantic segmentation，会预先定义一组颜色，并为每个语义类别分配固定的颜色。同样地，在instance segmentation中，实例目标的颜色是根据其空间位置类别分配的，即颜色的数量等于空间位置的数量，这导致模型只依赖颜色本身来判断任务，而不是利用不同segment之间的关系。
为了解决上述限制，我们提出了一种用于in-context coloring的随机着色方案。我们首先随机采样另一张与输入图像具有相似上下文的图像，例如具有相同的语义类别或object instance。接着，从目标图像中随机采样一组颜色，并将每种颜色映射为一个随机颜色，从而对对应像素进行重新着色。由此我们获得两对图像，定义为一个in-context pair。
此外，我们引入了mix-context training方法，通过使用混合样例对模型进行训练。这一方法将多张具有相同颜色映射的图像拼接在一起，然后对拼接后的图像进行随机裁剪和缩放，从而生成一个mixed-context训练样本。通过这种方式，模型学会关注图像的上下文信息，而不是仅依赖特定的颜色信息来判断任务。
这种统一方式使我们能够以一致的方式利用所有segmentation数据集，仅根据具体任务调整数据采样策略。我们根据不同的数据类型定义不同的上下文。例如，对于semantic segmentation，我们随机采样类别；对于instance segmentation，我们以随机数量采样object instances。同一图像的不同视角，例如通过一组图像增强方式变换而得的视角，也被视为上下文中的图像。在实现中，采样完全基于颜色，例如相同的颜色可表示相同的类别或相同的instance。
3.2. 上下文集成（Context Ensemble）
一旦训练完成，SegGPT在推理过程中能够释放出其全部潜力。SegGPT支持在上下文中进行任意的分割任务，例如，使用单张图像及其目标图像的例子。目标图像可以是单一颜色（不包括背景），也可以是多种颜色，例如在一次分割中分割多个类别或感兴趣的对象。具体而言，给定一张待测试的输入图像，我们将其与示例图像拼接，并将其输入到SegGPT中，以获得相应的in-context预测。
为了提供更准确和具体的上下文，可以使用多个示例。例如，可以使用相同语义类别的多个示例，或视频中的前几帧。为了有效地利用多个示例，我们提出了两种上下文集成方法。一种是叫做空间集成（Spatial Ensemble），将多个示例按n × n网格拼接，然后重新采样至与单个示例相同的大小。这种方法与in-context着色的直觉一致，多个示例的语义信息几乎不增加额外成本就可以被提取到上下文中。另一种方法是特征集成（Feature Ensemble）。多个示例在批次维度上合并并独立计算，唯一不同的是，在每个注意力层之后，查询图像的特征会被平均。通过这种方式，在推理过程中，查询图像可以从多个示例中收集信息。
3.3. 上下文调优（In-Context Tuning）
SegGPT能够在不更新模型参数的情况下，适应特定的使用案例。我们冻结整个模型，并初始化一个可学习的图像张量作为输入上下文。在训练过程中，只有这个可学习的图像张量会被更新，其余部分的训练保持不变，例如使用相同的损失函数。调优完成后，我们可以将学习到的图像张量提取出来，并将其作为特定应用的即插即用键。例如，给定一个固定对象类别集的数据集，如ADE20K，我们可以为该数据集训练一个定制的提示，而不会损害模型的通用性。或者，我们可以为特定场景（如你的公寓）或特定人物（如Bert的面部）优化一个提示图像。这为广泛的应用场景开辟了机会。
温馨提示：
阅读全文请访问"AI深语解构" SegGPT: 在上下文中分割一切