Source-linked AI summary
Semantic Segmentation with Generative Models: Semi-Supervised Learning and Strong Out-of-Domain Generalization
Daiqing Li, Junlin Yang, Karsten Kreis, Antonio Torralba, Sanja Fidler
TL;DR
Limited labeled data and costly pixel-level annotation make semantic segmentation difficult, especially when models must generalize across domains. The paper models the joint image-label distribution with a StyleGAN2-based GAN trained on mostly unlabeled images, then labels inputs through latent inversion. It reports competitive or better in-domain performance and strong generalization from CT to MRI and from face photographs to paintings, sculptures, cartoons, and animal faces.
Problem
Semantic segmentation needs expensive pixel-level labels, while discriminative semi-supervised methods may overfit because they do not explicitly model the input distribution.
Method
A StyleGAN2-based GAN models the joint image-label distribution using many unlabeled images and few labeled ones, with test-time latent inversion for labeling.
Results
The method achieves competitive or better in-domain performance and significantly higher out-of-domain generalization across medical-image and face-segmentation tasks.
Takeaways & Limitations
The model generalizes from CT to MRI and from natural face photographs to paintings, sculptures, cartoons, and even animal faces.
Takeaways & Limitations
The method is limited by generative-model expressivity and is not applicable to highly complex data such as vivid outdoor scenes.
Abstract
from arXiv · showhide
Training deep networks with limited labeled data while achieving a strong generalization ability is key in the quest to reduce human annotation efforts. This is the goal of semi-supervised learning, which exploits more widely available unlabeled data to complement small labeled data sets. In this paper, we propose a novel framework for discriminative pixel-level tasks using a generative model of both images and labels. Concretely, we learn a generative adversarial network that captures the joint image-label distribution and is trained efficiently using a large set of unlabeled images supplemented with only few labeled ones. We build our architecture on top of StyleGAN2, augmented with a label synthesis branch. Image labeling at test time is achieved by first embedding the target image into the joint latent space via an encoder network and test-time optimization, and then generating the label from the inferred embedding. We evaluate our approach in two important domains: medical image segmentation and part-based face segmentation. We demonstrate strong in-domain performance compared to several baselines, and are the first to showcase extreme out-of-domain generalization, such as transferring from CT to MRI in medical imaging, and photographs of real faces to paintings, sculptures, and even cartoons and animal faces. Project Page: \url{https://nv-tlabs.github.io/semanticGAN/}
1. Introduction
The paper addresses the need to reduce expensive pixel-level annotation while improving semantic-segmentation generalization beyond the training distribution. It proposes a generative image-label model trained semi-supervisedly and reports strong in-domain and out-of-domain performance.
- Motivation: Deep networks require large datasets, while semantic-segmentation annotations are especially time-consuming and expensive to collect.The problem is particularly acute in medical imaging, where annotations require highly skilled experts and sensors vary across sites.
- Results: The model generalizes across CT and MRI volumes and from natural face photographs to sculptures, paintings, cartoons, and animal faces.These experiments demonstrate the reported extreme out-of-domain behavior across substantially different visual domains.
- Motivation: Existing semi-supervised methods use unlabeled data but often do not explicitly model the input distribution, leaving them vulnerable to overfitting.This limitation is especially important when only a small amount of labeled segmentation data is available.
- Approach: The proposed GAN models the joint image-label distribution and synthesizes images with semantic masks using a large unlabeled collection and a small labeled subset.The approach builds on StyleGAN2 with an additional label-generation branch and uses adversarial objectives.
- Results: The method achieves competitive or better in-domain performance and significantly higher generalization on out-of-domain tests in medical imaging and face segmentation.The evaluation covers both medical images and human faces, with comparisons against competitive baselines.
2. Related Work
The related work contrasts discriminative and generative approaches to semi-supervised segmentation, positioning this paper as a fully generative model of the joint image-label distribution. It also connects the method to GAN inversion and generative image understanding.
- Semi-Supervised Learning and Semantic Segmentation: Prior semi-supervised segmentation methods use pseudo-labeling, adversarial training, transformation consistency, weak supervision, and related techniques.The cited literature includes medical and computer-vision segmentation methods, including mean-teacher and discriminator-based approaches.
- Generative Modeling: Many SSL methods train discriminative models of p(y|x), whereas this paper models the joint distribution p(x, y) and directly synthesizes images with pixel-wise labels.The paper distinguishes its approach from methods that use generative models only to augment training data.
- Generative Modeling: The proposed method is presented as the first fully generative semantic-segmentation approach using only adversarial objectives and no cross-entropy terms.Its generator models the joint image-label distribution and directly produces images together with labels.
- Generator Inversion: GAN inversion methods infer latent embeddings through optimization, encoders, or hybrid procedures; this paper uses inferred embeddings for pixel-wise labeling.The cited prior work primarily focused on image reconstruction and editing rather than semantic labeling.
- Generator Architecture: The generator architecture modifies StyleGAN2 to produce masks through additional branches that share style layers with image synthesis.The tImage and tSeg blocks output intermediate images and segmentation masks at different resolutions.
- Generative Models for Image Understanding: Generative modeling has also been used for image recognition, inpainting, colorization, super-resolution, calibration, and robustness, motivating its use for semantic segmentation.The paper interprets its approach as inpainting missing labels with a joint image-label generative model.
3. Method
The method models the joint distribution of images and segmentation labels with a StyleGAN2-based generator, enabling semi-supervised training and latent-space inference. It generates aligned image-label pairs and labels target images by reconstructing them through encoder initialization and test-time optimization.
- 3.2. Motivation: The model replaces conditional segmentation with a GAN that jointly models images and pixel-wise labels.A shared latent vector generates both outputs, while image-label alignment is enforced by a discriminator operating on concatenated pairs.
- 3.5. Inference: At inference, the optimized embedding generates a reconstructed image and aligned mask, yielding an approximate segmentation for the target image.The optimization begins from the encoder output and minimizes reconstruction error; the generator's learned alignment is then used to obtain the mask.
- 3.4. Training: Semi-supervised training combines a large unlabeled image set with a small labeled subset using adversarial objectives.The framework treats image synthesis as learning semantic information that can support label synthesis, with most generator parameters shared across branches.
- 3.3. Model: The StyleGAN2-based generator adds a segmentation branch at every style layer to synthesize an image and its mask from the same latent input.Two discriminators separately encourage realistic images and aligned image-label pairs.
- 3.5. Inference: The encoder maps images to W+-space, where independently modeled per-layer embeddings support test-time inversion.The encoder uses a feature-pyramid backbone and a fully convolutional mapping to provide a strong initialization for optimization.
- 3.4. Training: Encoder training combines supervised segmentation losses with unsupervised image reconstruction so embeddings regenerate inputs and, when available, their masks.The supervised objective uses pixel-wise cross-entropy and Dice loss, while the reconstruction objective uses LPIPS and pixel-wise L2 terms.
4. Experiments
The experiments evaluate the generative segmentation model across medical and face tasks, using limited labels and larger unlabeled sets. It performs competitively in-domain and shows strong generalization across datasets and substantially different image domains.
- Experimental setup: The evaluation covers chest X-ray, skin lesion, CT-MRI liver, and face part segmentation with small labeled and relatively large unlabeled datasets.Models are tested on both in-domain and out-of-domain datasets, using DICE, JC index, or mIoU according to task.
- Experimental setup: The model is compared with fully supervised U-Net-style approaches and semi-supervised methods including mean teacher, adversarial training, and Guided Collaborative Training.The semi-supervised baselines use the same ResNet-50 backbone and additional unlabeled data.
- Chest X-ray Segmentation: On chest X-ray segmentation, the model is on-par or better in-domain and outperforms supervised and semi-supervised baselines on out-of-domain datasets, often by a large margin.Table 1 reports DICE scores for in-domain CXR14 and JSRT and out-of-domain NLM, NIH, and SZ, with 108k additional unlabeled samples for the model and semi-supervised methods.
- Skin Lesion Segmentation: On skin lesion segmentation, the model consistently outperforms supervised and semi-supervised baselines both in-domain and out-of-domain.Table 2 evaluates JC index on ISIC and three additional out-of-domain datasets, using approximately 33k unlabeled samples.
- Face Part Segmentation: For face part segmentation, the model outperforms baselines and, with 1.5k labels, exceeds a DeepLab model trained with all 28k CelebA masks on out-of-domain MetFaces data.The paper also reports reasonable qualitative segmentations for cartoons, animal faces, and other extreme out-of-domain images.
- CT-MRI Transfer: On CT-MRI liver segmentation, the model outperforms chosen supervised baselines by a large margin when trained on CT and evaluated on unseen MRI data.The experiment uses 70 additional unlabeled LITS2017 volumes.
- Ablations and efficiency: Embedding optimization generally improves results over using the encoder prediction directly, while direct generative segmentation requires test-time optimization and is not suitable for real-time applications.The authors identify speed-ups as future work.
- Ablations and efficiency: Performance with 1500 labeled and 3K unlabeled examples is almost equivalent to performance with 150 labeled and 28K unlabeled examples.This ablation measures data value using segmentation mIoU.
5. Conclusion
The paper presents a StyleGAN2-based fully generative approach to semantic segmentation that supports semi-supervised training and strong in- and out-of-domain generalization.
- The method models the joint image-label distribution with a fully generative StyleGAN2-based approach.
- The approach naturally supports semi-supervised training while targeting strong generalization capabilities.
- Medical-domain experiments address settings where annotations are expensive and models must transfer between imaging sensors.
- The method significantly outperforms strong baselines both in-domain and out-of-domain.
- Face-part experiments show plausible segmentations for paintings, sculptures, cartoons, and animal faces.