[Paper Review] GAN Inversion for Data Augmentation to Improve Colonoscopy Lesion Classification

Mayank Golhar, Taylor L. Bobrow+2 more|arXiv (Cornell University)|May 4, 2022

AI in cancer detection4 citations

TL;DR

This paper proposes a GAN inversion-based data augmentation method for colonoscopy lesion classification, leveraging a single pre-trained StyleGAN-ADA model to generate synthetic images by manipulating latent vectors for style transfer (white light to narrowband imaging) and lesion size interpolation. The approach improves polyp classification accuracy by up to 6% over baselines without retraining generative models, enabling diverse, label-preserving augmentations from public datasets.

ABSTRACT

A major challenge in applying deep learning to medical imaging is the paucity of annotated data. This study demonstrates that synthetic colonoscopy images generated by Generative Adversarial Network (GAN) inversion can be used as training data to improve the lesion classification performance of deep learning models. This approach inverts pairs of images with the same label to a semantically rich & disentangled latent space and manipulates latent representations to produce new synthetic images with the same label. We perform image modality translation (style transfer) between white light and narrowband imaging (NBI). We also generate realistic-looking synthetic lesion images by interpolating between original training images to increase the variety of lesion shapes in the training dataset. We show that these approaches outperform comparative colonoscopy data augmentation techniques without the need to re-train multiple generative models. This approach also leverages information from datasets that may not have been designed for the specific colonoscopy downstream task. E.g. using a bowel prep grading dataset for a polyp classification task. Our experiments show this approach can perform multiple colonoscopy data augmentations, which improve the downstream polyp classification performance over baseline and comparison methods by up to 6%.

Motivation & Objective

• Address the challenge of limited annotated colonoscopy data in training deep learning models for polyp classification.
• Overcome limitations of traditional data augmentation techniques that fail to span the test distribution in colonoscopy imaging.
• Enable efficient, multi-task data augmentation using a single GAN model without retraining for each task.
• Leverage publicly available colonoscopy datasets to enrich internal training data with diverse imaging modalities and lesion attributes.
• Improve robustness and performance of polyp classification models through disentangled latent space manipulation.

Proposed method

• Train a StyleGAN-ADA generator on diverse public colonoscopy datasets to learn a rich, disentangled latent space encoding imaging modality, lesion morphology, and patient-specific features.
• Train an encoder (e4e) to invert real colonoscopy images into the StyleGAN latent space, enabling image reconstruction and latent manipulation.
• Perform domain translation between white light imaging (WLI) and narrowband imaging (NBI) by editing the latent code while preserving the original image class label.
• Generate synthetic images with progressive lesion sizes by linear interpolation in the latent space between two real images of the same class and modality.
• Use the augmented training data—comprising original images plus inverted, translated, and interpolated images—to retrain a polyp classification model.
• Apply the same pre-trained GAN model to multiple augmentation tasks (e.g., style transfer, interpolation), eliminating the need for separate models per task.

Experimental results

Research questions

• RQ1 Can GAN inversion be used to generate realistic, label-preserving synthetic colonoscopy images for data augmentation?
• RQ2 Does GAN inversion-based augmentation improve polyp classification performance compared to standard data augmentation and CycleGAN?
• RQ3 Can a single pre-trained GAN model perform multiple types of colonoscopy-specific data augmentation, such as modality translation and lesion size interpolation?
• RQ4 To what extent does interpolating in the latent space improve model generalization and decision boundary smoothness?
• RQ5 Can public colonoscopy datasets be leveraged to enhance internal datasets for downstream tasks like polyp classification?

Key findings

• The proposed method improved polyp classification accuracy by 6% over the baseline on the full training set and by 4.6% on the NBI-only subset.
• The method outperformed CycleGAN by 3.3% in accuracy, demonstrating superior performance in domain translation due to multi-image generation per input.
• Domain translation to five WLI and five NBI images per input improved accuracy by 2.9% over the baseline, with a 2% improvement over CycleGAN.
• Interpolation-based augmentation increased accuracy by 2.8% in the NBI domain and 1.1% in the WLI domain, with greater gains in the smaller NBI dataset.
• The interpolated images exposed the classifier to a denser, more continuous manifold of lesion sizes, leading to smoother decision boundaries and improved robustness.
• The approach enables diverse, label-preserving augmentations using only one pre-trained GAN model, reducing the need for training multiple generative models per task.