Semantic Segmentation with Generative Models: Semi-Supervised Learning and Strong Out-of-Domain Generalization

Abstract

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 [45], 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 testtime 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

EditGAN: High-Precision Semantic Image Editing

“…GANs and Latent Space Image Embedding. EditGAN builds on top of DatasetGAN [1] and SemanticGAN [2], which proposed to jointly model images and their semantic segmentations using shared latent codes. However, these works leveraged this model design only for semi-supervised learning, not for editing.…”

“…3 High-Precision Semantic Image Editing with EditGAN Deep generative models such as StyleGAN2, which are trained to synthesize highly realistic images, acquire a semantic understanding of the modeled images in their high-dimensional feature space. Recently, DatasetGAN [1] and SemanticGAN [2] built on this insight to learn a joint distribution p(x, y) over images x and pixel-wise semantic segmentation labels y, while requiring only a handful of labeled examples. EditGAN utilizes this joint distribution p(x, y) to perform high-precision semantic image editing of real and synthesized images.…”

“…Both methods [1,2] model p(x, y) by adding an additional segmentation branch to the image generator, which is a pre-trained StyleGAN [1]. We follow DatasetGAN [1], which applies a simple three-layer multi-layer perceptron classifier on the layer-wise concatenated and appropriately upsampled feature maps.…”

“…To both train the segmentation branch and perform segmentation on a new image, we embed an image into the GAN's latent space using an encoder and optimization. To this end, we build on previous works [2,62,66] and train an encoder that embeds images into W + space, which is defined as W but where the w's are modeled independently [62,66]. Our objectives to train this encoder consist of standard pixel-wise L2 and perceptual LPIPS reconstruction losses using both the real training data as well as samples from the GAN itself.…”

“…In this work, we are addressing these limitations and propose EditGAN, a novel GAN-based image editing framework that enables high-precision semantic image editing by allowing users to modify detailed object part segmentations. EditGAN builds on a recently proposed GAN that jointly models both images and their semantic segmentations based on the same underlying latent code [1,2], and requires as few as 16 labeled examples -allowing it to scale to many object classes and choices of part labels. We achieve editing by modifying the segmentation mask according to a desired edit and optimizing the latent code to be consistent with the new segmentation, thus effectively changing the RGB image.…”