Mask2Former and OneFormer: Universal Image Segmentation Models in šŸ¤— Transformers

TL;DR

Mask2Former and OneFormer are now available in šŸ¤— Transformers, offering a unified architecture that can perform instance, semantic, and panoptic segmentation without taskā€‘specific models.

Image segmentation tasks

Instance segmentation identifies each object instance (e.g., each person) and outputs a binary mask per instance.

Semantic segmentation assigns a single class label to every pixel, without distinguishing between separate instances of the same class.

Panoptic segmentation combines the two: it produces a set of nonā€‘overlapping segments, each with a binary mask and a class label, covering both "things" (instances) and "stuff" (background categories).

These three subtasks have historically required distinct model families, but recent research has converged on a single "mask classification" paradigm that treats all tasks uniformly.

Universal image segmentation

Since 2020, models such as DETR introduced a transformerā€‘based decoder that predicts a set of binary masks and class labels in parallel, enabling panoptic segmentation with a unified approach. MaskFormer demonstrated that the same paradigm works for semantic segmentation.

Mask2Former extends this idea to instance segmentation by improving the backbone, pixel decoder, and transformer decoder. The architecture consists of:

  1. A backbone (ResNet or Swin Transformer) producing lowā€‘resolution feature maps.
  2. A pixel decoder that upsamples these maps to highā€‘resolution features.
  3. A transformer decoder that receives a fixed set of queries and outputs binary mask proposals and class logits.

Mask2Former still requires separate training per task to achieve stateā€‘ofā€‘theā€‘art performance.

OneFormer builds on Mask2Former by adding a text encoder that conditions the model on a task description ("instance", "semantic", or "panoptic"). Trained only on a panopticā€‘style dataset, OneFormer attains stateā€‘ofā€‘theā€‘art results on all three tasks, at the cost of higher inference latency due to the extra text encoder. It supports Swin Transformer or DiNAT backbones.

Inference with the Transformers library

Both models can be loaded with a single line of code using AutoImageProcessor (or OneFormerProcessor) and the appropriate model class:

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

The library provides over 30 preā€‘trained checkpoints covering various datasets and backbones.

A typical inference pipeline:

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'] is a (H,ā€ÆW) map where each pixel value encodes the instance ID; segments_info holds class IDs, scores, and other metadata.

Visualization can be done with Matplotlib by mapping each segment ID to a distinct color and adding a legend that shows the class name and instance count.

OneFormer inference follows the same API but requires an additional text prompt, e.g., "segment everything" for panoptic, "segment instances" for instance, or "segment semantics" for semantic segmentation. A full demo notebook is available in the Hugging Face Transformersā€‘Tutorials repository.

Fineā€‘tuning on custom data

Fineā€‘tuning uses the same highā€‘level API as MaskFormer. Replace MaskFormerForInstanceSegmentation with Mask2FormerForUniversalSegmentation or OneFormerForUniversalSegmentation. The processor class also changes:

  • Mask2FormerImageProcessor (or AutoImageProcessor) for Mask2Former.
  • OneFormerProcessor for OneFormer, which handles both image and text inputs. Demo notebooks walk through dataset preparation, training loops, and evaluation for all three segmentation tasks.

Implications and why it matters

  • Unified workflow ā€“ Researchers and practitioners no longer need to maintain separate codebases for each segmentation task.
  • Reduced engineering overhead ā€“ A single model checkpoint can be deployed for multiple downstream applications (e.g., autonomous driving, medical imaging, content moderation).
  • Stateā€‘ofā€‘theā€‘art performance ā€“ OneFormer matches or exceeds specialized models while being trained on a single panoptic dataset, simplifying data collection.
  • Openā€‘source accessibility ā€“ By integrating these models into šŸ¤— Transformers, Hugging Face lowers the barrier to entry for highā€‘quality segmentation, enabling rapid prototyping and reproducible research.

Resources

Sources

Related