【论文速览】图像分割领域的通用大模型SegGPT - Segmenting Everything in Context

文章目录

代码地址：https://github.com/baaivision/Painter
Demo地址：https://huggingface.co/spaces/BAAI/SegGPT

研究背景

图像分割一直是计算机视觉领域的一项基础研究问题，在深度学习时代得到了迅猛发展，对于不同的分割子问题，涌现了很多又快又好的算法。但这些方法都是针对具体的子任务进行设计的，比如前景分割(foreground segmentation)、交互式分割(interactive segmentation)、语义分割(semantic segmentation)、实例分割(instance segmentation)以及全景分割(panoptic segmentation)。

交互式分割是一种图像处理技术，它允许用户通过进行正/负点击来提取目标掩模

不同子任务的分割模型有着其特殊的任务需求、分割类别、分割细粒度、数据类型等，这意味着要针对具体的问题进行模型的设计与训练。而该论文的目标就是提出一个模型，它能够解决所有的分割任务。如下图所示，该论文提出了SegGPT模型，即只用一个模型就能处理任何的上下文分割问题(segmenting everything in context)。

SegGPT全称是segment everything with a generalist Painter, 其中Painter是该团队的另一项工作，用上下文训练框架进行各项视觉任务（下节会对Painter进行简单介绍）

在这里插入图片描述

橘色的box表示对应的prompt图像与掩码，蓝色框为模型针对输入图像得到的结果。可以看到，SegGPT能够处理任意的物体分割(红色球、所有球体、阴影、上表面和物体轮廓等)，多部分分割(自由女神像的特殊部分)，彩虹分割、视频分割以及用可学习的prompt tuning进行闭集语义分割(close set semantic segmentation)。

解决思路

Painter

Images Speak in Images: A Generalist Painter for In-Context Visual Learning

这里先介绍SegGPT基于的前置工作"Painter"，Painter是该团队在CVPR2023年发布的一个通用的视觉大模型。它希望能够像NLP领域一种，基于预训练的大模型，通过设计不同的prompt来完成不同的任务。但与NLP不同，视觉任务的输出空间差异非常大，哪怕都是稠密预测任务【Painter主要基于稠密预测领域】，语义分割和深度估计的输出形式也是不一样的（一个是分类/另一个是回归）。其次，哪怕能够统一输出形式，如何设计prompt也是一个面临的问题。

为了解决这两个问题，Painter首先将输出空间统一成连续的输出图像表示（输出RGB图像），比如深度估计将0-10映射到0-255之间向下取整，语义分割的类别映射到具体的RGB值。

在这里插入图片描述

论文中，Painter统一了语义分割、深度估计、实例分割、关键点检测、图像去噪、图像去雨、图像增强七项任务（在in-domain和out-of-domain数据集上都有着较好效果）

为了让模型能够根据不同的prompt执行不同的任务，Painter在训练时会拼接两组图像，然后进行"masked image modeling"。具体来说：在输入时选择同域同任务下的两组图像进行拼接（每组图像包括输入图像和它对应的任务输出），然后进行如下所示的掩码重构学习（类似simMiM）。

在这里插入图片描述
在推理阶段，给定prompt( $prompt_{in},prompt_{out}$ )、测试图像( $y_{in}$ )，将它们和mask一起组合送入模型即可，如第一个大图所示，Painter能够根据不同的prompt来执行不同的任务。

为了得到推理阶段的prompt，有两种方式：一种是基于训练集搜索的方式；另一种即为生成式的方式，即为不同任务学习一个对应的prompt。

SegGPT

SegGPT与Painter不同，其聚焦于图像分割任务，旨在采用一个模型来解决所有分割问题。这又面临着两方面的挑战：① 如何融合不同任务中的数据集类型；② 如何设计有效的训练策略使得模型能够灵活地处理各项任务。

在Painter框架中，不同任务的颜色空间都是预定义的，这样会使得模型沦为多任务学习。比如语义分割就是每种语义类别对应一种颜色，而实例分割会根据实例所在空间位置得到对应颜色（类似SOLO）。

为了解决这个限制，SegGPT采用了一种“随机着色”的策略：对于每幅图像随机选择另一幅图像（两幅图像都有着相似的语义类别或者对象实例，也可以通过数据增广获得），从目标图像中随机采样一组颜色，进行随机映射，使得相应像素的重新着色。

在这里插入图片描述
通过这种方式，不同的分割数据类型就可以转换成同一种图像格式统一起来，对于不同的任务，采用不同的采样策略就好。比如语义分割就采样类别，实例分割就采样实例数，然后进行随机颜色映射+着色。

经过这一步，不同分割任务下的数据得到了统一。模型要做的事情就是：根据上下文信息来进行区域着色问题（颜色是随机的）

在训练方式上，SegGPT和Painter一样，都是建模成MIM的形式。具体地说，SegGPT进行上下文着色(in-context coloring)学习，即根据上下文信息对图像的mask区域(类别、对象实例、物体区域等)进行着色。

模型采用普通的ViT + smooth-l1损失进行训练（过程与Painter基本一致）

测试过程与Painter一致，将Prompt图像与测试图像一起组合送入SegGPT即可。不过为了利用多幅prompt来进一步提升分割准确性，SegGPT采用了两种ensemble策略：① 空间融合；② 特征融合。

在这里插入图片描述

空间融合即在空间层面上进行图像组合（拼接+缩放）；特征融合在每个attention层之后对query进行了一次平均。

为了针对具体的任务设计更好的prompt格式，SegGPT也采用了可学习的prompt，如下所示，冻结整个SegGPT，在对应任务的数据集上训练一个可学习的图像向量即可。

在这里插入图片描述

实验效果（部分）

SegGPT在各项分割任务的基准数据集上都进行测试，无论是域内(in-domain)还是域外(out-of-domain)的图像都能够有效地进行分割。

在这里插入图片描述

SegGPT在COCO-20i和PASCAL-5i的few-show语义分割效果（in-domain）以及在FSS-1000上的few-show分割效果（out-of-domain）

在这里插入图片描述
表3展示了在视频物体分割上的表现，其与先进的方法非常接近（SegGPT并没有用到视频数据训练）。表5表6展示了在ADE20k语义分割和COCO全景分割上的效果，距离最优方法还有一段距离，并且相比于通用的Painter性能还下降了。这是因为SegGPT并没有针对语义分割任务进行训练，而是进行的随机颜色着色(利用上下文进行着色)，这种训练难度更高。

思考

该论文主要针对图像分割问题提出了一种通用大模型SegGPT，其不局限于具体的分割类型，可以根据提供的prompt来分割任意区域。

SegGPT本质上是Painter的一种特殊变体，只不过在训练策略上有了一些变化，不根据预定义的颜色进行上下文着色，而是随机着色，强迫模型去根据当前的上下文信息来进行分割。这种策略带来的好处就是SegGPT具有更强的泛化能力与零样本/少样本迁移能力，但坏处在于对于具体的in-domain任务上，距离专业模型还有一定性能差距。

此外，目前发布的分割领域大模型（比如SegGPT、SAM），其都是基于监督式学习方式【需要提供分割的GT】，这就意味着在数据的收集上比较耗时与费力。如果能找到一种更好的方式，像NLP领域中的GPT那样进行自监督式的预训练学习就更完美了。

关于SAM，segment anything model我也会写相应的blog进行介绍，欢迎关注

参考资料

[1] CVPR/AAAI 2023论文分享