[论文解读] 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 detection被引用 4

一句话总结

本文提出了一种基于 GAN 反演的结肠镜病变分类数据增强方法，利用单一预训练的 StyleGAN-ADA 模型，通过操纵潜在向量实现风格迁移（白光成像至窄带成像）和病变大小插值，生成合成图像。该方法在不微调生成模型的前提下，使多巴胺分类准确率最高提升 6%，实现了来自公开数据集的多样化、标签保持的增强。

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%.

研究动机与目标

• 解决在训练深度学习模型进行多巴胺分类时，标注的结肠镜数据有限的挑战。
• 克服传统数据增强技术在结肠镜成像中无法覆盖测试分布的局限性。
• 通过单一 GAN 模型实现高效、多任务的数据增强，无需为每项任务重新训练模型。
• 利用公开的结肠镜数据集，通过多样化成像模式和病变特征丰富内部训练数据。
• 通过解耦潜在空间的操纵，提升多巴胺分类模型的鲁棒性和性能。

提出的方法

• 在多样化的公开结肠镜数据集上训练 StyleGAN-ADA 生成器，以学习一个丰富且解耦的潜在空间，编码成像模式、病变形态和患者特异性特征。
• 训练一个编码器（e4e）将真实结肠镜图像反演到 StyleGAN 潜在空间中，以实现图像重建和潜在变量操纵。
• 通过编辑潜在码在白光成像（WLI）和窄带成像（NBI）之间实现域迁移，同时保持原始图像的类别标签。
• 通过在两个同类别同模态的真实图像之间进行潜在空间线性插值，生成具有渐进病变大小的合成图像。
• 使用增强后的训练数据（包括原始图像、反演图像、转换图像和插值图像）重新训练多巴胺分类模型。
• 将同一预训练 GAN 模型应用于多种增强任务（如风格迁移、插值），避免为每项任务训练独立模型。

实验结果

研究问题

• RQ1 能否利用 GAN 反演生成逼真且标签保持的合成结肠镜图像以实现数据增强？
• RQ2 与标准数据增强和 CycleGAN 相比，基于 GAN 反演的增强是否能提升多巴胺分类性能？
• RQ3 单一预训练 GAN 模型能否执行多种结肠镜特异性数据增强任务，如模态转换和病变大小插值？
• RQ4 在潜在空间中插值在多大程度上改善了模型的泛化能力和决策边界的平滑性？
• RQ5 能否利用公开的结肠镜数据集增强内部数据集，以支持多巴胺分类等下游任务？

主要发现

• 该方法在完整训练集上使多巴胺分类准确率相比基线最高提升 6%，在仅含 NBI 的子集上提升 4.6%。
• 该方法在准确率上比 CycleGAN 提升 3.3%，表明其在域迁移方面表现更优，原因在于可为单个输入生成多张图像。
• 将输入图像对应的域迁移结果扩展为五张 WLI 和五张 NBI 图像，使准确率相比基线提升 2.9%，比 CycleGAN 提升 2%。
• 基于插值的增强在 NBI 域中使准确率提升 2.8%，在 WLI 域中提升 1.1%，在较小的 NBI 数据集中增益更显著。
• 插值生成的图像使分类器接触到更密集、更连续的病变大小流形，从而实现更平滑的决策边界并提升鲁棒性。
• 该方法仅使用一个预训练 GAN 模型即可实现多样化、标签保持的增强，显著减少了为每项任务训练多个生成模型的需求。