InvSeg: Test-Time Prompt Inversion for Semantic Segmentation

Jiayi Lin, Jiabo Huang, Jian Hu, Shaogang Gong
Queen Mary University of London
arXiv Supplementary Code

Abstract

The precise visual-textual correlations embedded in the attention maps derived from text-to-image generative diffusion models have been shown beneficial to open-vocabulary dense visual prediction tasks,e.g., semantic segmentation. However, a significant challenge arises due to the distributional discrepancy between the context-rich sentences used for image generation training and the isolated class names typically available for visual discrimination. This discrepancy in the richness of textual context limits the effectiveness of diffusion models in capturing accurate visual-textual correlations. To tackle this challenge, we propose a novel approach called InvSeg, a test-time prompt inversion method that leverages per-image visual context to optimize the context-insufficient text prompts composed of isolated class names, so as to associate every pixel and class for open-vocabulary semantic segmentation. Specifically, we introduces a Contrastive Soft Clustering (CSC) method to derive the underlying structure of images according to the assumption that different objects usually occupy distinct while continuous areas within visual scenes. Such structural information is then used to constrain the image-text cross-attention for calibrating the input class embeddings without requiring any manual label or additional training data. By incorporating sample-specific context at test time, InvSeg learns context-rich text prompts in embedding space to achieve accurate semantic alignment across modalities. Experiments show that InvSeg achieves state-of-the-art performance on the PASCAL VOC and Context datasets.

Our Framework

MY ALT TEXT


The framework of InvSeg. Our proposed Contrastive Soft Clustering method can achieve region-level prompt inversion. The text tokens are first initialized with the pretrained text encoder from the diffusion model (dashed box on left) and then are used as the only learnable parameters during the test time training. After the adaption process, the learned text tokens can be used to derive a more accurate and complete refined attention maps $$\{M\}$$ for segmentation.

Comparison with SOTA

MY ALT TEXT


Comparison with existing methods. Models in the first three rows are finetuned on target datasets while the rest approaches do not require mask annotations. Bold fonts refer to the best results among the models and underline fonts refer to the second best.

Visualization

MY ALT TEXT

Examples of Segmentation on VOC (top), Context (middle) and COCO (bottom). For each sample (image group of four), from left to right is input, GT, InvSeg, Diffusion baseline.

MY ALT TEXT

Visualization of refined cross-attention maps derived from text prompts before (top) and after (bottom) prompt inversion. Before prompt inversion, the segmentation of background elements such as "grass" or "trees" is influenced by foreground objects like "cow" or "horse", resulting in mistakenly ignoring background classes or segmenting foreground (and background) classes. After applying prompt inversion, this phenomenon is suppressed by improving the distinction between foreground and background through proposed Contrastive Soft Clustering.

BibTeX

@misc{lin2024invsegtesttimepromptinversion,
      title={InvSeg: Test-Time Prompt Inversion for Semantic Segmentation}, 
      author={Jiayi Lin and Jiabo Huang and Jian Hu and Shaogang Gong},
      year={2024},
      eprint={2410.11473},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2410.11473}, 
}