CVPR 2023 | ODISE: Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models

CVPR 2023 | ODISE: Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models

• 论文：https://arxiv.org/pdf/2303.04803.pdf
• 代码：https://github.com/NVlabs/ODISE

架构设计

• 冻结的CLIP编码器：
  • 图像编码器和可学习的MLP，用于编码图像到隐式的文本嵌入。
  • 文本编码器编码数据集类别标签或者当前图像的caption中的名词。
• 在LAION子集上预训练的文本对图像扩散模型Stable Diffusion，使用冻结权重的版本，对输入的加噪图像和文本嵌入处理获得多层级金字塔形式的图像特征。实际扩散使用的时间步为0。
• 可学习的基于Mask2Former的mask generator，从输入的特征中预测N=100个mask。

解码设计

从扩散模型中可以提取从图像输入和隐式文本嵌入获取的图像表征，通过mask generator获得N个mask prediction和对应的mask embedding。

• 先对mask预测使用匈牙利匹配于真值配对，之后对配对的预测进行损失计算。
• mask prediction使用binary mask loss监督。
• mask embedding与CLIP文本编码器输出的文本嵌入之间计算相似性矩阵从而实现mask的类别分配。由于有两种不同的设定：
  • 基于数据集类别标签的形式中，使用交叉熵损失。属于各个类别的概率由mask embedding与每个真值类别的嵌入计算向量内积并通过softmax处理后获得。
  • 基于对应的图像caption的形式中，使用grounding loss。

实际推理中，假定输入图像没有给定的caption和真值标签，只有给定的测试集类别集合。这里对于分类，在原有的mask generator的mask embedding的基础上，额外引入了一组新的mask embedding。新引入的是通过使用mask prediction对implicit caption中的CLIP视觉编码器的图像特征执行masked average pooling来获得对应的额外嵌入。二者利用一个固定的超参数 $\lambda \in [0,1]$ （实际使用0.65）通过几何平均的形式组合。获得最终的分类概率预测。

其他设计

• 使用的基于CLIP视觉编码特征和masked pooling的处理，与ZegFormer等方法中使用的方案（从原始图像中剪裁出N个预测的掩码区域对应的边界框区域，再通过图像编码器编码）一样有效。单独来看，基于扩散特征的预测在两个数据集上的表现都比仅基于判别特征的更好。然而，将两者融合在一起会导致更高的性能。这里实际上涉及到了两种不同的利用CLIP图像编码器的方式：
  • Bounding box cropping：基于mask对原图前景区域crop后提取目标视觉嵌入。对于多个mask而言，需要独立多次masking和cropping，并执行多次前向传播。由于CLIP的输入尺寸是固定的，所以多次前向传播较为费时。实验中，ADE20K获得了23.7 PQ，速度为0.38FPS。
  • Mask pooling：基于mask对原图对应的视觉嵌入masked average pooling。对于多个mask而言，只需要执行单次前向传播，多次masked pooling即可。速度相较于第一种方案更有效率。实验中，ADE20K获得了23.4 PQ，速度是cropping策略的3倍。
• 一次训练，多种任务推理：
  • 训练时使用COCO训练集，利用全景分割注释作为监督信号。如果使用图像caption信息，则随机从对应的caption注释中选择一个。
  • 推理时模型使用统一的一个checkpoint结果，对于语义分割，则是通过预测的结果进行合并后得到的。例如对于语义分割，合并所有属于thing的类别到单一一个类别输出。这种方案可以实现语义分割，但实际上效果其实是次优的。

• 方法中，训练在大规模互联网数据的基于扩散的生成模型在视觉表征上优势明显。在基于K-means的聚类效果上，本文的设计更具有语义判别性。而在定量性能上，所提的基于扩散的方案超过了:
  • 预训练在相同规模数据的文本匹配模型CLIP。
  • 预训练在较小规模更多类别的ImageNet的类别条件生成模型LDM、判别模型、自监督模型。

• 实验对比发现，对于文本图像扩散模型中使用的文本嵌入来说，使用一个显式的或者是隐式的caption要比直接使用空字符串更好。但是显式的caption模型由于对训练数据分布的偏好，导致跨数据集的应用效果稍弱于本文使用的方法。因为本文是基于预训练在大规模互联网数据的CLIP模型的驱动下学习的，所以具有更好的泛化能力。

• 实验对比发现，对于采样步数，本文发现t=0的时候整体模型性能是最好的。即使将t在训练过程中设置为一个可学习的量，其最终也是非常接近0的。
• 实际使用Open-vocabulary image segmentation中相同的prompt engineering strategy来创建类别文本prompt的集成。