Source-linked AI summary

Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models

Jiarui Xu, Sifei Liu, Arash Vahdat, Wonmin Byeon, Xiaolong Wang, Shalini De Mello

arXiv:2303.04803v4cs.CV

TL;DR

Open-vocabulary panoptic segmentation requires recognizing arbitrary categories while jointly parsing instances and scene semantics, and discriminative models may lack sufficient spatial understanding. ODISE combines frozen text-to-image diffusion and discriminative representations, using diffusion features for mask generation and text embeddings for mask classification. It achieves state-of-the-art performance across open-vocabulary panoptic and semantic segmentation tasks and demonstrates the potential of diffusion representations for recognition.

  • Problem

    Open-vocabulary recognition must identify arbitrary categories while unified panoptic segmentation remains uncommon, and diffusion visual representations depend on paired captions that are unavailable for many application images.

  • Method

    ODISE combines frozen text-to-image diffusion and discriminative models, using diffusion features to generate panoptic masks and text embeddings to classify them with label or caption supervision.

  • Results

    ODISE achieves state-of-the-art accuracy across multiple open-vocabulary recognition tasks and outperforms CLIP(H) by a large margin on all reported metrics despite equal-sized training data.

  • Takeaways & Limitations

    The results demonstrate the potential of text-to-image diffusion models’ internal representations for open-vocabulary segmentation and other downstream recognition tasks.

  • Takeaways & Limitations

    The evaluation includes category definitions that can be ambiguous and non-exclusive, such as ADE20K’s overlap between “tower” and “building.”

Abstract

from arXiv · show

We present ODISE: Open-vocabulary DIffusion-based panoptic SEgmentation, which unifies pre-trained text-image diffusion and discriminative models to perform open-vocabulary panoptic segmentation. Text-to-image diffusion models have the remarkable ability to generate high-quality images with diverse open-vocabulary language descriptions. This demonstrates that their internal representation space is highly correlated with open concepts in the real world. Text-image discriminative models like CLIP, on the other hand, are good at classifying images into open-vocabulary labels. We leverage the frozen internal representations of both these models to perform panoptic segmentation of any category in the wild. Our approach outperforms the previous state of the art by significant margins on both open-vocabulary panoptic and semantic segmentation tasks. In particular, with COCO training only, our method achieves 23.4 PQ and 30.0 mIoU on the ADE20K dataset, with 8.3 PQ and 7.9 mIoU absolute improvement over the previous state of the art. We open-source our code and models at https://github.com/NVlabs/ODISE .

1. Introduction

Open-vocabulary recognition aims to identify limitless categories, but unified panoptic segmentation remains difficult because discriminative models may lack spatial and relational understanding. ODISE addresses this by combining frozen text-to-image diffusion and discriminative models for open-vocabulary panoptic segmentation.

  • Open-vocabulary recognition seeks to parse all object instances and scene semantics across limitless categories in a unified panoptic segmentation framework.
  • Discriminative text-image models generalize well for open-vocabulary classification but may confuse spatial relations, creating a bottleneck for scene-level panoptic understanding.
  • Diffusion models’ internal visual representations are plausibly correlated with language-described high- and mid-level semantic concepts through text-conditioned cross-attention.
  • ODISE extracts frozen diffusion features from an image and its caption, generates panoptic masks, and classifies them with open-vocabulary text embeddings.
  • ODISE is presented as the first work to explore large-scale text-to-image diffusion models for open-vocabulary segmentation and reports state-of-the-art results across multiple tasks.

2. Related Work

Prior work largely addressed closed-vocabulary panoptic segmentation or only one open-vocabulary segmentation setting, typically using discriminative models. ODISE instead provides a unified framework that uses diffusion-model representations for both instance and semantic segmentation.

  • Panoptic Segmentation: Closed-vocabulary panoptic methods recognize only training categories, limiting segmentation to finite vocabularies smaller than real-world language vocabularies.
  • Open-Vocabulary Segmentation: Prior open-vocabulary methods commonly perform either instance segmentation or semantic segmentation alone rather than unifying both.
  • Open-Vocabulary Segmentation: ODISE differs by using large-scale text-to-image diffusion representations alongside discriminative models for unified open-vocabulary instance and semantic segmentation.
  • Generative Models for Segmentation: Earlier generative segmentation approaches commonly trained generative models on small-vocabulary datasets before learning category classification from few-shot annotations.

3. Method

ODISE combines frozen text-to-image diffusion and text-image discriminative representations with a mask generator for open-vocabulary panoptic segmentation. Its diffusion features provide dense, spatially differentiated representations, while text-based classification assigns open-vocabulary labels to predicted masks.

  • Method overview: ODISE trains a mask generator on frozen text-to-image diffusion features to produce class-agnostic panoptic masks and mask embeddings.The generator can be instantiated with bounding-box-based or direct segmentation mask-based architectures.
  • Diffusion features: Diffusion UNet features are used as dense visual representations because text enters the denoising network through cross-attention and encourages semantic correlations.ODISE extracts these features with a single forward pass using a noisy image at a selected diffusion timestep.
  • Diffusion features: An implicit captioner generates a text embedding from the input image itself, addressing the dependence of diffusion representations on paired captions.This avoids requiring explicitly captioned images during feature extraction.
  • Mask classification: Mask classification uses text-image discriminative models to compare mask embeddings with encoded training-category names under category-label supervision.The classification probability is computed from mask–text similarity with a learnable temperature parameter.
  • Mask classification: ODISE also supports image-caption supervision by treating nouns in captions as grounding labels, and trains with either label or caption supervision alongside mask losses.The two supervision signals are selected according to the comparison setting.
  • Open-vocabulary inference: At inference, diffusion and discriminative predictions are fused with a geometric mean to classify each predicted mask into test categories.The diffusion representation produces many plausible instance masks, while the discriminative image encoder further classifies masked regions.

4. Experiments

ODISE is evaluated across open-vocabulary panoptic, semantic, and related segmentation settings, with ablations examining representations, captioning, and diffusion time steps. It achieves strong performance while using mostly frozen parameters.

  • ODISE is evaluated on ADE20K for open-vocabulary panoptic, instance, and semantic segmentation, and on Pascal datasets for semantic segmentation.
  • ODISE has 28.1M trainable parameters and 1,493.8M frozen parameters, running at 1.26 FPS with 11.9 GB memory.
  • 8.3 PQ and 8.4 mAP gains over MaskCLIP are reported for ADE20K open-vocabulary panoptic and instance segmentation.
  • 7.6 mIoU on A-150, 4.7 mIoU on A-847, and 4.8 mIoU on PC-459 are reported with caption supervision versus the next-best method.
  • Diffusion-based internal representations outperform CLIP(H) by a large margin on all metrics despite equal-sized LAION training data.
  • Explicit or implicit captions outperform empty text, while implicit captioning generalizes best among the compared variants.
  • All metrics decrease as diffusion time increases, with t=0 producing the best results; concatenating three time steps is 3× slower with similar accuracy.
  • Fusing diffusion and discriminative class predictions improves over either model individually on ADE20K and COCO.

5. Conclusion

The conclusion presents ODISE as an initial demonstration that frozen internal representations from large-scale text-to-image diffusion models can support open-vocabulary recognition. It identifies open-vocabulary segmentation as a promising direction for using such representations in future tasks.

  • ODISE is presented as the first step in leveraging frozen internal representations of large-scale text-to-image diffusion models for downstream recognition.
  • The work reports strong potential for text-to-image generation models in open-vocabulary segmentation and establishes a new state of the art.
  • The authors conclude that text-to-image diffusion models learn rich semantic representations beyond generating plausible images.
  • The paper opens a direction for leveraging text-to-image model representations for other tasks in the future.

A. Implementation Details

ODISE combines frozen Stable Diffusion and CLIP components with a segmentation architecture, prompt ensembles, and COCO-based training details. The implementation emphasizes extensive frozen capacity and practical inference trade-offs.

  • Training uses 90k iterations, 1024^2 images, large-scale jittering, batch size 64, AdamW, learning rate 0.0001, and weight decay 0.05.
  • Prompt ensembles are created for each test category, and the category with maximum probability is predicted.
  • ODISE contains 28.1M trainable and 1,493.8M frozen parameters, runs at 1.26 FPS, and uses 11.9 GB memory on an NVIDIA V100.
  • Mask feature pooling gives a 3x speedup over bounding-box cropping while maintaining similar ADE20K PQ.

B.1. Comparison with State of the Art

Across ADE20K, Cityscapes, Mapillary Vistas, COCO, LVIS, UVO, and ADE20K, ODISE consistently outperforms MaskCLIP, CLIP(H), GGN, or existing state of the art on open-vocabulary tasks.

  • ADE20K panoptic segmentation: ODISE outperforms MaskCLIP on all detailed ADE20K panoptic metrics when both models are trained on COCO.The metrics include segmentation quality, recognition quality, and panoptic quality for thing and stuff categories.
  • Cross-dataset evaluation: ODISE’s diffusion-feature model outperforms its CLIP(H) variants by large margins on Cityscapes and Mapillary Vistas.The evaluation uses original test-dataset text labels without carefully selecting category names.
  • Swapped training and evaluation: ODISE outperforms its CLIP(H) variant by a large margin when trained on ADE20K and evaluated on COCO.This experiment swaps the training and evaluation datasets.
  • Open-vocabulary object detection: ODISE surpasses MaskCLIP by a large margin on LVIS object detection using either category-label or caption supervision.LVIS evaluation reports mAP and mAPr, with mAPr measuring performance on 337 rare categories.
  • Open-world instance segmentation: ODISE outperforms the existing state of the art by 14.3 points on UVO and 9.3 points on ADE20K for open-world instance segmentation.The metric is average recall of 100 mask proposals, AR@100.

B.2. Ablation Study

The ablation evidence indicates that diffusion-model internal features are more semantically differentiated than CLIP features and support qualitative open-vocabulary segmentation across several datasets.

  • Visual representations: Diffusion features are much more semantically differentiated than CLIP’s frozen internal features in K-means visualizations.Quantitative comparisons of ODISE and its CLIP(H) variant further substantiate diffusion features’ superiority.
  • Qualitative evaluation: Qualitative open-vocabulary panoptic segmentation results are visualized on COCO, ADE20K, and Ego4D.The supplied passages identify the datasets and figures but do not report a quantitative comparison.

C. Qualitative Results

ODISE performs open-vocabulary inference over approximately 1.5k categories, including categories absent from COCO training annotations, and produces plausible results across multiple datasets.

  • Open-vocabulary inference: ODISE performs open-vocabulary inference with approximately 1.5k merged LVIS, COCO, and ADE20K test classes.It is trained only on COCO and evaluated on ADE20K and Ego4D, with qualitative results also shown on COCO validation data.
  • Novel categories: ODISE segments novel categories such as police cruiser, chandelier, grocery bag, and power shovel that are not annotated in COCO.The results include qualitative visualizations on COCO, ADE20K, and Ego4D.
  • Cross-domain qualitative results: ODISE outputs good-quality plausible panoptic segmentation on Ego4D despite the large domain gap from COCO training.Ego4D consists of diverse ego-centric videos.

D. Limitations and Future Work

The paper identifies ambiguity in current category definitions and potential bias in diffusion representations as limitations and directions for future analysis.

  • Category definitions: Ambiguous and non-exclusive dataset categories can cause confusions such as ADE20K’s “tower” being mis-classified as “building”.The authors suggest studying how category definitions affect evaluation accuracy and note prompt engineering may mitigate the issue.
  • Representation bias: Web-crawled image-text pretraining data may introduce potential bias into the diffusion model’s internal representation despite filtering.This limitation concerns the pretrained text-to-image diffusion model used by ODISE.
Loading 2303.04803v4…