使用 Mask2Former 和 OneFormer 进行通用图像分割

发布日期：2023 年 1 月 19 日

在 GitHub 上更新

本指南介绍了 Mask2Former 和 OneFormer，这是两种最先进的图像分割神经网络模型。这些模型现已在 🤗 transformers 中提供，这是一个开源库，提供易于使用的最先进模型实现。在此过程中，您将了解各种图像分割形式之间的区别。

图像分割

图像分割是将图像中不同的“区域”进行识别的任务，例如人或汽车。更专业地说，图像分割是将具有不同语义的像素分组的任务。有关简要介绍，请参阅 Hugging Face 任务页面。

图像分割大致可以分为 3 个子任务——实例分割、语义分割和全景分割——每个子任务都有许多方法和模型架构。

实例分割 是识别图像中不同“实例”（如个体人物）的任务。实例分割与目标检测非常相似，只是我们希望输出一组二值分割掩码，而不是边界框，并带有相应的类标签。实例通常也称为“对象”或“事物”。请注意，单个实例可能会重叠。
语义分割 是识别图像中每个像素的不同“语义类别”（如“人”或“天空”）的任务。与实例分割相反，语义分割不区分给定语义类别的单个实例；例如，只希望为“人”类别而不是为单个个体创建掩码。没有单个实例的语义类别，如“天空”或“草地”，通常被称为“背景物”，以区别于“事物”（很棒的名字，是吧？）。请注意，语义类别之间不可能有重叠，因为每个像素都属于一个类别。
全景分割，由 Kirillov 等人于 2018 年引入，旨在通过让模型简单地识别一组“段”，每个段都有相应的二值掩码和类标签，从而统一实例分割和语义分割。段可以是“事物”或“背景物”。与实例分割不同，不同段之间不可能重叠。

下图说明了 3 个子任务之间的区别（取自此博客文章）。

drawing

在过去几年中，研究人员提出了几种架构，这些架构通常非常适合实例分割、语义分割或全景分割。实例分割和全景分割通常通过为每个对象实例输出一组二值掩码 + 相应的标签来解决（与目标检测非常相似，只是为每个实例输出二值掩码而不是边界框）。这通常被称为“二值掩码分类”。另一方面，语义分割通常通过让模型输出一个带有每个像素一个标签的“分割图”来解决。因此，语义分割被视为“逐像素分类”问题。采用此范式的流行语义分割模型包括 SegFormer（我们为此撰写了一篇详细的博客文章）和 UPerNet。

通用图像分割

幸运的是，自 2020 年左右以来，人们开始提出能够使用统一架构和相同范式解决所有 3 个任务（实例、语义和全景分割）的模型。这始于 DETR，它是第一个使用“二值掩码分类”范式解决全景分割的模型，通过统一处理“事物”和“背景物”类别。关键创新是让 Transformer 解码器以并行方式生成一组二值掩码 + 类别。这在 MaskFormer 论文中得到了改进，该论文表明“二值掩码分类”范式也适用于语义分割。

Mask2Former 通过进一步改进神经网络架构将其扩展到实例分割。因此，我们已经从独立的架构发展到研究人员现在所称的“通用图像分割”架构，能够解决任何图像分割任务。有趣的是，这些通用模型都采用了“掩码分类”范式，完全放弃了“逐像素分类”范式。下图描绘了 Mask2Former 的架构（取自原始论文）。

drawing

简而言之，图像首先通过骨干网络（在论文中可以是 ResNet 或 Swin Transformer）以获得低分辨率特征图。接下来，使用像素解码器模块增强这些特征图以获得高分辨率特征。最后，Transformer 解码器接收一组查询，并根据像素解码器的特征将其转换为一组二值掩码和类别预测。

请注意，Mask2Former 仍然需要针对每个任务单独训练才能获得最先进的结果。这已通过 OneFormer 模型得到改进，该模型仅通过训练数据集的全景版本（！）即可在所有 3 个任务上获得最先进的性能，通过添加文本编码器来根据“实例”、“语义”或“全景”输入对模型进行条件化。该模型目前也在 🤗 transformers 中提供。它比 Mask2Former 更准确，但由于额外的文本编码器，延迟更高。有关 OneFormer 的概述，请参阅下图。它利用 Swin Transformer 或新的 DiNAT 模型作为骨干网络。

drawing

使用 Transformers 中的 Mask2Former 和 OneFormer 进行推理

Mask2Former 和 OneFormer 的用法相当简单，与它们的先行者 MaskFormer 非常相似。让我们从 Hub 实例化一个在 COCO 全景数据集上训练的 Mask2Former 模型，以及它的处理器。请注意，作者发布了不下 30 个检查点，这些检查点在各种数据集上进行了训练。

from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation

processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")

接下来，让我们加载 COCO 数据集中熟悉的猫图像，我们将对其执行推理。

from PIL import Image

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image

我们使用图像处理器为模型准备图像，然后将其通过模型。

inputs = processor(image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

模型输出一组二值掩码和相应的类 logits。Mask2Former 的原始输出可以很容易地使用图像处理器进行后处理，以获得最终的实例、语义或全景分割预测。

prediction = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
print(prediction.keys())

Output:
----------------------------------------------------------------------------------------------------
dict_keys(['segmentation', 'segments_info'])

在全景分割中，最终的 prediction 包含两部分：一个形状为 (高, 宽) 的 segmentation 图，其中每个值编码给定像素的实例 ID，以及相应的 segments_info。segments_info 包含有关图中单个段的更多信息（例如它们的类别/类别 ID）。请注意，Mask2Former 为了效率输出形状为 (96, 96) 的二值掩码提案，并且 target_sizes 参数用于将最终掩码调整到原始图像大小。

让我们将结果可视化

from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import cm

def draw_panoptic_segmentation(segmentation, segments_info):
    # get the used color map
    viridis = cm.get_cmap('viridis', torch.max(segmentation))
    fig, ax = plt.subplots()
    ax.imshow(segmentation)
    instances_counter = defaultdict(int)
    handles = []
    # for each segment, draw its legend
    for segment in segments_info:
        segment_id = segment['id']
        segment_label_id = segment['label_id']
        segment_label = model.config.id2label[segment_label_id]
        label = f"{segment_label}-{instances_counter[segment_label_id]}"
        instances_counter[segment_label_id] += 1
        color = viridis(segment_id)
        handles.append(mpatches.Patch(color=color, label=label))
        
    ax.legend(handles=handles)

draw_panoptic_segmentation(**panoptic_segmentation)

在这里，我们可以看到模型能够检测图像中的猫和遥控器个体。另一方面，语义分割只会为“猫”类别创建一个单一的掩码。

要使用 OneFormer 执行推理，其 API 除了还接受额外的文本提示作为输入外，其余与 Mask2Former 相同，我们参考演示笔记本。

在 Transformers 中微调 Mask2Former 和 OneFormer

要将 Mask2Former/OneFormer 微调到自定义数据集上进行实例、语义和全景分割，请查看我们的演示笔记本。MaskFormer、Mask2Former 和 OneFormer 共享相似的 API，因此从 MaskFormer 升级很容易，只需要最少的更改。

演示笔记本使用 MaskFormerForInstanceSegmentation 加载模型，而您需要切换到使用 Mask2FormerForUniversalSegmentation 或 OneFormerForUniversalSegmentation。在 Mask2Former 的图像处理情况下，您还需要切换到使用 Mask2FormerImageProcessor。您也可以使用 AutoImageProcessor 类加载图像处理器，该类会自动加载与您的模型对应的正确处理器。另一方面，OneFormer 需要 OneFormerProcessor，它准备图像以及文本输入，以便模型使用。