Per-Pixel Classification is Not All You Need for Semantic Segmentation

Idea

Instead for pre-pixel classification, we could first predict a set of binary masks and then assign a single class to each mask.

Each prediction is supervised with a per-pixel binary mask loss and a classification loss

Architecture

1. Using a backbone to extract images features $\mathcal{F}$
2. Upsample the features to obtain per-pixel embeddings ${\mathcal{E}}_{\mathrm{pixel}}$
3. A transformer decoder attends to image features and produces $N$ per-segment embeddings $\mathcal{Q}$, then embeddings then independently generate
   • $N$ class predictions (shape $N\times (K+1)$), where $\varnothing$ is introduced to represent "no-object" class
   • $N$ corresponding mask embeddings ${\mathcal{E}}_{\mathrm{mask}}$
4. Predict $N$ possibly overlapping binary mask predictions via a dot product between pixel embeddings ${\mathcal{E}}_{\mathrm{pixel}}$ and mask embeddings ${\mathcal{E}}_{\mathrm{mask}}$ followed by a sigmoid activation.
5. Finally we can get the prediction by combining $N$ binary masks with their class predictions using a simple matrix multiplication.