Self-Supervised One-Shot Learning for Automatic Segmentation of StyleGAN Images

TL;DR

This paper introduces a self-supervised one-shot segmentation method for StyleGAN-generated images, leveraging multi-scale features and contrastive learning to improve speed and accuracy over existing semi-supervised approaches.

Contribution

It presents a novel contrastive clustering framework utilizing multi-scale GAN features and a unique data augmentation strategy for efficient one-shot segmentation.

Findings

Outperforms semi-supervised baselines with 1.02% higher wIoU

Increases inference speed by 4.5 times

Effective in generating annotated synthetic baggage X-ray scans

Abstract

We propose a framework for the automatic one-shot segmentation of synthetic images generated by a StyleGAN. Our framework is based on the observation that the multi-scale hidden features in the GAN generator hold useful semantic information that can be utilized for automatic on-the-fly segmentation of the generated images. Using these features, our framework learns to segment synthetic images using a self-supervised contrastive clustering algorithm that projects the hidden features into a compact space for per-pixel classification. This contrastive learner is based on using a novel data augmentation strategy and a pixel-wise swapped prediction loss that leads to faster learning of the feature vectors for one-shot segmentation. We have tested our implementation on five standard benchmarks to yield a segmentation performance that not only outperforms the semi-supervised baselines by an average wIoU margin of 1.02 % but also improves the inference speeds by a factor of 4.5. Finally, we also show the results of using the proposed one-shot learner in implementing BagGAN, a framework for producing annotated synthetic baggage X-ray scans for threat detection. This framework was trained and tested on the PIDRay baggage benchmark to yield a performance comparable to its baseline segmenter based on manual annotations.