See our collection for all versions of OneFormer.

Run OneFormer with Keras 3: JAX, PyTorch, or TensorFlow

zeromodels/oneformer_coco_swin_large

Paper: OneFormer: One Transformer to Rule Universal Image Segmentation (arXiv:2211.06220) · HF Papers

OneFormer trains one model on semantic, instance, and panoptic segmentation jointly, conditioned on a text token that names the task. At inference you pass task="semantic", "instance", or "panoptic" to the same weights. The processor combines an image processor and a tokenizer for that task token.

For more details on the model, please go to SHI Labs' original model card.

Pure-Keras 3 conversion of shi-labs/oneformer_coco_swin_large for zeromodels. One implementation runs unmodified on TensorFlow / Torch / JAX.

This is a universal segmentation checkpoint (OneFormerUniversalSegment) trained on COCO (133 classes, Swin-Large). All three tasks share these weights.

✨ Quick start

import os
os.environ["KERAS_BACKEND"] = "torch"  # or "jax" / "tensorflow"

from PIL import Image
from zeromodels.models.oneformer import (
    OneFormerUniversalSegment,
    OneFormerProcessor,
)

model = OneFormerUniversalSegment.from_weights("zeromodels/oneformer_coco_swin_large")
processor = OneFormerProcessor.from_weights("zeromodels/oneformer_coco_swin_large")

image = Image.open("your_image.jpg").convert("RGB")
# task is an argument: "semantic" | "instance" | "panoptic"
inputs = processor(images=image, task="panoptic")
output = model(inputs)
result = processor.post_process_panoptic_segmentation(
    output, target_size=(image.height, image.width)
)
print(result["segmentation"].shape)

Load any OneFormer variant the same way with from_weights("zeromodels/<variant>"):

Tips

• Set KERAS_BACKEND before importing Keras / zeromodels.
• Prefer OneFormerProcessor.from_weights(...) so image size and tokenizer match the variant.
• Pick the task at call time via processor(..., task=...); the checkpoint vocabulary comes from the training set, not the task.
• See OneFormer docs and Loading Weights.
• Community / upstream safetensors still work via the hf: prefix, e.g. OneFormerUniversalSegment.from_weights("hf:shi-labs/oneformer_coco_swin_large").

Special Thanks

A huge thank you to the SHI Labs OneFormer authors for creating and releasing these models.

License: MIT.