Mask2Former 与 OneFormer：🤗 Transformers 中的通用图像分割模型

TL;DR

Mask2Former 和 OneFormer 现已在 🤗 Transformers 中提供，提供统一的架构，能够在无需任务特定模型的情况下执行实例、语义和全景分割。

图像分割任务

实例分割 标识每个对象实例（例如，每个人），并为每个实例输出二值掩码。

语义分割 为每个像素分配单一的类别标签，而不区分同一类别的不同实例。

全景分割 将两者结合：它生成一组不重叠的分割，每个分割都有二值掩码和类别标签，覆盖“事物”（实例）和“材料”（背景类别）。

这三个子任务历来需要不同的模型系列，但近期研究已汇聚到一种统一的“掩码分类”范式，对所有任务进行统一处理。

通用图像分割

自 2020 年起，诸如 DETR 的模型引入了基于 Transformer 的解码器，能够并行预测一组二值掩码和类别标签，从而以统一的方式实现全景分割。MaskFormer 证明了相同的范式同样适用于语义分割。

Mask2Former 将此思路扩展到实例分割，通过改进骨干网络、像素解码器和 Transformer 解码器。该架构包括：

  1. 一个骨干网络（ResNet 或 Swin Transformer），生成低分辨率特征图。
  2. 一个像素解码器，将这些特征图上采样至高分辨率特征。
  3. 一个 Transformer 解码器，接收固定数量的查询并输出二值掩码提案和类别 logits。

Mask2Former 仍然需要针对每个任务进行单独训练，以实现最先进的性能。

OneFormer 在 Mask2Former 的基础上添加了文本编码器，使模型能够根据任务描述（“instance”、 “semantic” 或 “panoptic”）进行条件化。仅在全景风格的数据集上训练，OneFormer 在所有三项任务上均达到最先进的结果，但由于额外的文本编码器，推理延迟更高。它支持 Swin Transformer 或 DiNAT 骨干网络。

使用 Transformers 库进行推理

两种模型都可以通过一行代码使用 AutoImageProcessor（或 OneFormerProcessor）和相应的模型类加载：

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"
)

该库提供了超过 30 个预训练检查点，涵盖各种数据集和骨干网络。

典型的推理流程：

from PIL import Image, ImageDraw
import requests, torch

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

inputs = processor(image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# Convert raw mask proposals to panoptic output
prediction = processor.post_process_panoptic_segmentation(
    outputs, target_sizes=[image.size[::-1]]
)[0]
print(prediction.keys())  # dict_keys(['segmentation', 'segments_info'])

prediction['segmentation'] 是一个 (H, W) 的映射，其中每个像素值编码实例 ID；segments_info 包含类别 ID、分数以及其他元数据。

可使用 Matplotlib 进行可视化，将每个分割 ID 映射到不同的颜色，并添加显示类别名称和实例计数的图例。

OneFormer 推理 使用相同的 API，但需要额外的文本提示，例如全景分割使用 "segment everything"，实例分割使用 "segment instances"，语义分割使用 "segment semantics"。完整的演示 notebook 可在 Hugging Face Transformers‑Tutorials 仓库中获取。

在自定义数据上微调

微调使用与 MaskFormer 相同的高级 API。将 MaskFormerForInstanceSegmentation 替换为 Mask2FormerForUniversalSegmentation 或 OneFormerForUniversalSegmentation。处理器类也相应更改：

  • Mask2FormerImageProcessor（或 AutoImageProcessor）用于 Mask2Former。
  • OneFormerProcessor 用于 OneFormer，能够处理图像和文本输入。 演示 notebook 逐步演示数据集准备、训练循环以及针对所有三种分割任务的评估。

含义及其重要性

  • 统一工作流 – 研究人员和实践者不再需要为每个分割任务维护独立的代码库。
  • 降低工程开销 – 单一模型检查点即可部署到多个下游应用（例如自动驾驶、医学影像、内容审核）。
  • 最先进的性能 – OneFormer 在单一全景数据集上训练即可匹配或超越专用模型，简化了数据收集。
  • 开源可及性 – 将这些模型集成到 🤗 Transformers 中，Hugging Face 降低了高质量分割的入门门槛，促进快速原型和可复现的研究。

资源

Sources