DiffuMask: Synthesizing Images with Pixel-level Annotations for Semantic Segmentation Using Diffusion Models

In plain words

Labelling every pixel of an image by hand is slow and expensive. Text-to-image generators can already draw a "horse on the grass" from a sentence. While they do this, they keep an internal map of which pixels respond to which word. The idea of DiffuMask is to read that map as a rough outline of the object named by the word, clean it up, and use the generated image plus its outline as a training example for an ordinary segmentation model.

The hard part is that the raw maps are blurry and the right cut-off between "object" and "background" changes from image to image. The authors therefore pick the cut-off per image, drop the worst generated examples, and make the prompts and images more varied so that the synthetic data looks less like a set of copies. The result is a segmentation model trained without any hand-drawn masks, which the authors report comes reasonably close to a model trained on real labelled photos for some classes.

TL;DR

• Training data is produced from a frozen, text-only-trained Stable Diffusion model: images plus pixel masks taken from its cross-attention maps §3.1 §4.1.
• The mask pipeline is: average attention over layers and steps eq. 2; threshold, with a per-image threshold chosen by matching against an AffinityNet estimate eq. 3 eq. 4; DenseCRF refinement; noise pruning §3.3; a second segmenter trained to smooth the masks §4.1.
• Author-reported VOC 2012 val: 70.6 mIoU with synthetic data only (Mask2Former, Swin-B) vs 84.3 with the full real set tab. 1.
• The paper is a conference paper (ICCV 2023) with the released code's airplane-only data path and several gaps listed below. Read the Paper vs code and Pitfalls sections before implementing.

Problem

Dense segmentation labels are costly. Cityscapes labelling is quoted at up to 60 minutes per image §1. Earlier generative-data methods (DatasetGAN and BigDatasetGAN) need some pixel-level labels to train a decoder on GAN features §1 §2. The paper asks whether a text-supervised diffusion model can supply masks with no pixel labels at all §1.

Two obstacles are named: masks from attention are coarse and not class-discriminative enough at one resolution, and the synthetic-to-real domain gap §1 §3.2 §3.4 §3.5.

Method

3.1 Cross-attention in text-conditioned diffusion

In the U-Net of Stable Diffusion, image latents produce queries Q. The prompt is encoded by the text encoder into keys K and values V. The cross-attention map is A = Softmax(QKᵀ / √d), reshaped to H × W × N for N text tokens §3.1 eq. 1. For token j, the slice A_j is a spatial map of how strongly the image responds to that word §3.1.

3.2 Mask generation and refinement

Multi-scale, multi-step average eq. 2. The paper takes attention maps from four U-Net resolutions, 8×8, 16×16, 32×32 and 64×64, and from several denoising steps §3.2 fig. 2. Each map is divided by its own maximum, and the maps are averaged over the S layers and T steps eq. 2. The text says S is the total number of steps and T the number of layers, but the subscripts say the opposite (see Pitfalls).

Standard binarisation eq. 3. A fixed threshold γ is applied to the averaged map, and the result is refined with DenseCRF using colour and pixel distance §3.2.1 eq. 3. The paper shows that the best γ differs between classes and images, and that a fixed value is not robust §3.2.1 fig. 3 tab. 4a.

Adaptive threshold. The authors train an AffinityNet per category §4.1. Pixels whose averaged-map score lies in a middle band are treated as "neutral" and ignored; the other pixels give positive or negative affinity labels for neighbouring pixel pairs §3.2.2. At inference the AffinityNet gives a coarse affinity map B̂. For each image and class, the threshold γ̂ is the candidate in a search set Ω that maximises a matching cost between B̂ and the binary mask at γ §3.2.2 eq. 4.

Mask smoothing. The CRF output has jagged edges and small spurious regions. The authors train a Mask2Former on those masks and use its predictions on the synthetic training images as the final annotation §4.1.

3.3 Noise learning (Prune by Class)

Masks with IoU under 80% against the model's prediction make up a noticeable share of the data §3.3 fig. 5. The method is:

1. Get out-of-sample predictions with a segmenter (Mask2Former or Mask R-CNN) using cross-validation §3.3.
2. For each class, score each sample by its self-confidence (an IoU-based quality estimate) §3.3.
3. Prune the α fraction of samples with the lowest score and train on the rest §3.3.

The text sets α = 50% in the method section §3.3; the experiments use α = 0.7 §4.1 tab. 4c.

3.4 Prompt engineering

• Sub-classes. For each category, K sub-classes are taken from Wikipedia and inserted into the template "Photo of a [sub-class] " §3.4 fig. 6.
• Retrieval. The prompt is used to query a CLIP retrieval index built on LAION-5B. The top N captions are kept, giving K × N prompts. One prompt is drawn at random for each image §3.4 fig. 4. N is not stated.

3.5 Data augmentation

Four augmentations are used to close the gap to real photos §3.5 fig. 7:

• Splicing: grids of 1×2, 2×1, 2×2, 3×3, 5×5 or 8×8 images, sampled at random from the training set §3.5.
• Gaussian blur: kernel length sampled from 6 to 22 §3.5.
• Occlusion: CutMix-style patch paste with labels mixed by patch area §3.5.
• Perspective transform: described only in words; no parameters §3.5.

Procedure used in the experiments

# Build the synthetic training set per class  [§4.1] [§3.2] [§3.3]
def build_dataset(cls, n_images, alpha):
    # Generate images and attention maps with the frozen diffusion model [§3.1]
    images, attn_maps = generate(cls, n_images)
    # Average the maps over layers and steps  [eq. 2]
    avg_maps = [average_attention(m) for m in attn_maps]
    # Adaptive threshold per image, then DenseCRF  [eq. 4] [§3.2.2]
    masks = [refine(avg, affinity_for(avg)) for avg in avg_maps]
    # Prune the alpha fraction with the lowest quality  [§3.3]
    kept = prune_by_class(images, masks, alpha)
    # Train a segmenter on the kept masks, then smooth labels  [§4.1]
    return smooth_labels(kept)

# Average one token's attention over layers and steps  [eq. 2]
def average_attention(maps, layers, steps):
    # maps[s][t]: (H, W) map for layer s at step t  [§3.2]
    total = zeros_like(maps[0][0])
    for s in layers:
        for t in steps:
            # divide each map by its own maximum  [eq. 2]
            total += maps[s][t] / maps[s][t].max()
    # divide by the number of (layer, step) pairs  [eq. 2]
    return total / (len(layers) * len(steps))

# Pick a per-image threshold by matching an affinity map  [eq. 4] [§3.2.2]
def adaptive_threshold(affinity, avg_map, search_space):
    # affinity: coarse map from AffinityNet, shape (H, W)  [§3.2.2]
    best_gamma = None
    best_score = float("-inf")
    for gamma in search_space:
        # binary mask from the averaged map at this threshold  [§3.2.2]
        binary = avg_map > gamma
        # matching cost between affinity map and binary mask  [eq. 4]
        score = matching_cost(affinity, binary)
        if score > best_score:
            best_score = score
            best_gamma = gamma
    return best_gamma

# Standard binarisation with DenseCRF  [eq. 3] [§3.2.1]
def standard_binarise(avg_map, gamma):
    # stack the fixed threshold with the map, take the argmax label  [eq. 3]
    labels = argmax(stack([gamma, avg_map]))
    # refine with colour and distance pairwise terms  [§3.2.1]
    return dense_crf(labels)

# Prune by class: keep the samples with the highest self-confidence  [§3.3]
def prune_by_class(samples, scores, alpha):
    # alpha = 0.7 in the experiments  [§4.1] [tab. 4c]
    ranked = sort_by_score(samples, scores)
    # the alpha fraction with the lowest scores is removed as noisy  [§3.3]
    cut = int(alpha * len(ranked))
    return ranked[cut:]

# Splice a grid of synthetic images  [§3.5]
def splice(images, rows, cols):
    # grid shapes used: 1x2, 2x1, 2x2, 3x3, 5x5, 8x8  [§3.5]
    tiles = sample(images, rows * cols)
    return grid_concat(tiles, rows, cols)

Hyperparameters

Code-only values (from the repository, not the paper):

Results

Author-reported. Setups below come from the paper's tables.

Semantic segmentation on VOC 2012 val, mIoU over 20 classes tab. 1

Cityscapes val, Human and Vehicle only tab. 2 (Mask2Former, mIoU is the mean of the two)

Zero-shot (open-vocabulary) segmentation on VOC, seen / unseen / harmonic tab. 3

The table's "Categories" column lists 15 for the real-mask methods and "15+5" for the synthetic rows; the paper does not explain the difference tab. 3.

ADE20K val, three classes tab. 5

Domain generalisation (train → test, VOC/Cityscapes/ADE) tab. 6 (Mask2Former, ResNet-50; Person and Rider merged in Cityscapes)

Ablations on VOC val tab. 4. The mIoU columns in these four sub-tables are the mean over Bird and Dog only (checked by arithmetic against the rows).

• (a) Annotation: affinity map 81.6; fixed γ = 0.4 85.3; γ = 0.5 78.9; γ = 0.6 44.4; DiffuMask adaptive threshold 89.5.
• (b) Prompts: no prompt 76.9; retrieval only 77.7; retrieval + 10 sub-classes 87.6; + 50 sub-classes 89.0; + 100 sub-classes 89.5.
• (c) Noise learning α: 0.3 → 83.2; 0.4 → 84.7; 0.5 → 88.2; 0.6 → 89.1; 0.7 → 89.5.
• (d) Augmentation: none 84.3; + splicing 87.0; + blur 87.9; + occlusion 88.2; + perspective 89.5.

Backbone (five classes: bird, dog, sheep, horse, person) tab. 7: ResNet-50 70.3; ResNet-101 70.5; Swin-B 87.3; Swin-L 87.4 (mIoU column is the mean of these five classes).

Mask precision vs domain gap, four classes tab. 8: real images with manual labels 87.2; synthetic images with pseudo labels from a VOC-trained Mask2Former 82.7; synthetic images with DiffuMask labels 76.3. The paper attributes about 6.4 points to mask precision and 4.5 to the image gap §4.5.

Paper vs code

Official code: weijiawu/DiffuMask README.md, HEAD 913d2327deffb0e927ef2cc6c343700a2e2096a6. The first author is the owner. The repository has no LICENSE file, so no code is quoted here; everything below is a reading of the logic.

• Threshold selection is not the paper's method. The paper matches an AffinityNet affinity map against candidate thresholds §3.2.2 eq. 4. curve_threshold.py sweeps thresholds 0.35 to 0.54 in steps of 0.01 and keeps the one with the highest IoU against a file read from a refine_gt directory curve_threshold.py. The source of refine_gt is not in this repository. If it is ground-truth masks, the released threshold choice uses labels, which the paper's method does not. Not determined from the code alone.
• CRF settings are stated only in code. The paper gives no DenseCRF parameters §3.2.1. Code: 10 iterations, unary confidence 0.5, a Gaussian term (sxy 3, compat 3) and a bilateral term (sxy 80, srgb 13, compat 10), then a 0.5 cut curve_threshold.py curve_threshold.py.
• Attention generation. Code uses 50 diffusion steps and guidance scale 7 parallel_generate_VOC_Attention_AnyClass.py parallel_generate_VOC_Attention_AnyClass.py. Each attention map is resized to 512 × 512 with cubic interpolation and divided by its maximum parallel_generate_VOC_Attention_AnyClass.py. Averaging across layers is by sum divided by count parallel_generate_VOC_Attention_AnyClass.py. The per-step averaging of eq. 2 was not traced through the aggregation function.
• AffinityNet training. The paper says only that settings follow the original §4.1. Code: crop 448, batch 8, 8 epochs, base LR 0.1 with group multipliers 1, 2, 10, 20, weight decay 5e-4, a polynomial schedule, random crop and horizontal flip, colour jitter, and loss weights 1/4, 1/4, 1/2 for background, foreground and neutral-negative terms train_aff.py train_aff.py train_aff.py. Labels are average-pooled by 8 train_aff.py.
• Augmentation. The paper lists six splicing grids §3.5. Data_Augmentation.py concatenates size images with a default of 1, 2 or 3 Data_Augmentation.py Data_Augmentation.py. No code for Gaussian blur, occlusion or perspective transform was found in the repository (searched for blur, perspective, cutmix, kernel).
• Synthetic counts. The generation example asks for 8000 images per class VOC_data_generation.sh; the paper uses 10k §4.1. The README says only airplane data is generated in this repository and suggests adding other classes as negatives README.md.
• Noise-learning α. Not in the released code path inspected; the value follows the paper's experiments (0.7).

Pitfalls

• Ablation numbers are two-class means. Every mIoU in the ablation table is the average of Bird and Dog only, not the 20-class VOC mIoU tab. 4. Do not compare them with Table 1.
• Contribution claim conflicts with the table. The contributions list names "dog" as within 3% of real data §1. Table 1 gives dog 86.0 synthetic vs 96.8 real (10.8 points) and bird 92.9 vs 93.7 (0.8 points) tab. 1. Section 4.2 says "within 5%" for bird, cat, cow, horse and sheep, but horse is 89.0 vs 94.4 (5.4 points) tab. 1.
• Pairing error in the finetune claim. The text pairs "84.9% finetuned on 5.0k real" with "83.4% on full real (10.6k)" §4.2. Table 1 gives 83.4 for 5.0k real only and 84.3 for full real tab. 1.
• "19.2% mIoU improvement" is not reproduced. Swin-B over ResNet-50 on the 20-class synthetic run is 13.2 points (70.6 vs 57.4) tab. 1. The five-class backbone table gives 17.0 (87.3 vs 70.3) tab. 7. Neither gives 19.2 §4.5.
• Motorbike number. The text says Cityscapes-trained motorbike mIoU is 28.9; the table gives 28.3 §4.4 tab. 6.
• ADE20K is not "competitive". Synthetic Swin-B reaches 69.6 vs 87.9 real tab. 5, a gap of about 18 points, despite the text calling it competitive §4.2.
• Noise-learning α is inconsistent. Method text: α = 50% §3.3. Experiments and Figure 5: α = 0.7 §4.1 tab. 4c. Figure 5 caption says 30k images per class; the ablation says 10k fig. 5 tab. 4c.
• Cross-validation folds. The text says "five-fold" and then "k = 3" §4.1.
• Attention-map averaging notation. Eq. 2 is said to define S as total steps and T as layers, but s indexes layers and t steps §3.2 eq. 2.
• Adaptive threshold is underspecified. The search set Ω has L unstated candidates, and the sum in eq. 4 has no stated index. "Neutral" band limits are not given §3.2.2 eq. 4. The released code is different (see Paper vs code).
• Cityscapes counts do not add up. The text says 30k generated and 10k kept per sub-category across six sub-categories §4.1, which gives 60k kept. The table says 100k training images tab. 2. Not reconciled.
• Mask smoothing hides a second model. The final labels come from a Mask2Former trained on the CRF masks §4.1. Reproductions must budget a second segmentation training run.
• Compute is the hidden cost. Generation is done per class with a 50-step sampler [code] and 8 V100s are used for all experiments §4.1. The paper gives no per-image time. The paper's LaTeX also contains commented-out limitation text; it is not part of the compiled paper and is not used here.
• Single object per image. Only one object is considered per image §4.1.
• Backbone insight is one-sided. The claim that a stronger backbone reduces the domain gap §4.5 rests on one backbone table with five classes tab. 7 and a visual comparison fig. 2; it is not tested across datasets.
• Domain generalisation is mixed. On Cityscapes the DiffuMask-trained model (59.4) is below the ADE20K-trained one (60.1) tab. 6.

Lineage

• Builds on: Stable Diffusion, used frozen §3.1 §4.1; AffinityNet for affinity-based refinement §3.2.2; DenseCRF §3.2.1; CLIP retrieval over LAION-5B §3.4; Mask2Former as the evaluation and noise-learning segmenter §4.1.
• Compared against: DatasetGAN and BigDatasetGAN §1 §2; Li et al. (guided synthetic data with a COCO-trained model) §4.3 tab. 3; real-mask zero-shot methods (ZegFormer, SIGN, ZS3, CaGNet, Joint, STRICT) tab. 3.
• Superseded by: not stated in this card.