OneFormer: One Transformer to Rule Universal Image Segmentation

Idea

Propose a universal model for semantic, instance and panoptic segmentation (See differences), which just requires once-and-for-all training.

• Introduce a task token to condition the model on the task at hand
• Use a query-text contrastive loss

Architecture

Multi-Scale Feature Modeling

Extract multi-scale features for an input image using a backbone, followed by a pixel decoder (translating the high-level features into pixel-level predictions)

Unified Task-Conditional Query Formulation

Task Input Token

Uniformly sample {task} from {panopitc, instance, semantic} to form the task input token

At the same time, uniformly sample the corresponding ground truth to generate text queries ${\text{Q}}_{\text{text}}$

Text Query (${\text{Q}}_{\text{text}}$)

As shown below, first iterate over the set of masks to create a list of text (${\text{T}}_{\text{text}}$) with a template a photo with a {CLS}. Then pad ${\text{T}}_{\text{text}}$ with a/an {task} entries to obtain a padded list (${\text{T}}_{\text{pad}}$) of constant length $N_{\text{text}}$, with padded entries representation no-object masks. To obtain the text queries ${\text{Q}}_{\text{text}}$, we first tokenize the text entries ${\text{T}}_{\text{pad}}$ and pass the tokenized representation through a text-encoder, which is a $6$-layer transformer. The encoded $N_{\text{text}}$ text embeddings represent the number of binary masks and their corresponding classes in the input image. We further concatenate a set of $N_{\text{ctx}}$ learnable text context embeddings ${\text{Q}}_{\text{ctx}}$ to the encoded text embeddings to obtain the final $N$ text quires ${\text{Q}}_{\text{text}}$  The goal of ${\text{Q}}_{\text{ctx}}$ is to provide a unified textual context that captures the relevant information necessary for various tasks, such as image segmentation.

Object Query ($\text{Q}$)

First initialize the object queries (${\text{Q}}^{\prime }$) as $N-1$ times repetitions of the task-token ${\text{Q}}_{\text{task}}$. Then update ${\text{Q}}^{\prime }$ with guidance from the flattened $\frac{1}{4}$-scale features inside a $2$-layer transformer. The updated ${\text{Q}}^{\prime }$ from the transformer (rich with image-contextual information) is concatenated with ${\text{Q}}_{\text{task}}$ to obtain a task-conditioned representation of $N$ queries, $\text{Q}$

Query-Task Contrastive Loss

Considering that we have a batch of $B$ object-text query pairs ${\left\{(q_i^{\text{obj}},q_i^{\text{txt}})\right\}}_{i=1}^B$, where $q_i^{\text{obj}}$ and $q_i^{\text{txt}}$ are the corresponding object and text queried, respectively, of the $i$-th pair, we measure the similarity between the queries by calculating a dot product.

$$\begin{array}{rl}{\mathcal{L}}_{\text{Q}\to {\text{Q}}_{\text{text}}}& =-\frac{1}{B}\sum _{i=1}^B\log \frac{\exp \left(q_i^{\text{obj}}\odot q_i^{\text{txt}}/\tau \right)}{{\sum }_{j=1}\exp \left(q_i^{\text{obj}}\odot q_j^{\text{txt}}/\tau \right)}\\ {\mathcal{L}}_{{\text{Q}}_{\text{text}}\to \text{Q}}& =-\frac{1}{B}\sum _{i=1}^B\log \frac{\exp \left(q_i^{\text{txt}}\odot q_i^{\text{obj}}/\tau \right)}{{\sum }_{j=1}\exp \left(q_i^{\text{txt}}\odot q_j^{\text{obj}}/\tau \right)}\\ {\mathcal{L}}_{\text{Q}\leftrightarrow {\text{Q}}_{\text{text}}}& ={\mathcal{L}}_{\text{Q}\to {\text{Q}}_{\text{text}}}+{\mathcal{L}}_{{\text{Q}}_{text}\to \text{ Q}}\end{array}$$

Here $\tau$ is a learnable temperature parameter to scale the contrastive logits.