[논문 리뷰] GAN Inversion for Data Augmentation to Improve Colonoscopy Lesion Classification

Mayank Golhar, Taylor L. Bobrow외 2명|arXiv (Cornell University)|2022. 05. 04.

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 모델을 사용해 효율적이고 다중 작업 데이터 증강을 가능하게 한다.
• 공개된 대장내시경 데이터셋을 활용해 내부 훈련 데이터에 다양한 영상 모odalities와 병변 특성을 통합한다.
• 잠재 공간의 분리된 조작을 통해 폴립 분류 모델의 강건성과 성능을 향상시킨다.

제안 방법

• 다양한 공개 대장내시경 데이터셋을 기반으로 StyleGAN-ADA 생성자 훈련을 통해 영상 모달리티, 병변 형태, 환자 고유 특징을 인코딩한 풍부하고 분리된 잠재 공간을 학습한다.
• 실제 대장내시경 영상을 StyleGAN 잠재 공간으로 역전환할 수 있도록 인코더(e4e)를 훈련시켜 영상 복원 및 잠재 벡터 조작을 가능하게 한다.
• 원본 영상 클래스 레이블을 유지하면서 잠재 코드 편집을 통해 백색광 영상(WLI)과 좁역대 영상(NBI) 간 도메인 번역을 수행한다.
• 동일한 클래스와 모달리티의 두 실제 영상 간 잠재 공간에서 선형 보간을 통해 점진적인 병변 크기를 갖는 합성 영상을 생성한다.
• 원본 영상에 더해 역전환된, 번역된, 보간된 영상으로 구성된 증강 훈련 데이터를 사용해 폴립 분류 모델을 재훈련한다.
• 동일한 사전 훈련된 GAN 모델을 다양한 증강 작업(예: 스타일 전이, 보간)에 적용함으로써 각 작업마다 별도의 모델을 훈련할 필요 없이 해결한다.

실험 결과

연구 질문

• RQ1 식별자 기반 GAN을 활용해 현실적이고 레이블 유지 가능한 합성 대장내시경 영상을 생성하여 데이터 증강에 활용할 수 있는가?
• RQ2 기존 데이터 증강 기법과 CycleGAN 대비, GAN 기반 역전환 증강이 폴립 분류 성능을 향상시키는가?
• RQ3 단일 사전 훈련된 GAN 모델이 모달리티 번역 및 병변 크기 보간과 같은 대장내시경 전용 다중 증강 기법을 수행할 수 있는가?
• RQ4 잠재 공간 내 보간이 모델 일반화 능력과 결정 경계의 매끄러움에 얼마나 기여하는가?
• RQ5 공개된 대장내시경 데이터셋을 활용해 후속 작업(예: 폴립 분류)을 위한 내부 데이터셋을 향상시킬 수 있는가?

주요 결과

• 제안된 방법은 전체 훈련 세트에서 기준 모델 대비 6% 향상된 폴립 분류 정확도를 달성했으며, NBI 전용 세트에서는 4.6% 향상되었다.
• CycleGAN 대비 3.3% 높은 정확도를 기록하여, 다중 영상 생성 능력 덕분에 도메인 번역에서 뛰어난 성능을 입증했다.
• 입력당 5장의 WLI 및 5장의 NBI 영상으로 도메인 번역을 수행한 결과, 기준 모델 대비 2.9% 향상되었으며, CycleGAN 대비 2% 향상되었다.
• 보간 기반 증강은 NBI 영역에서 2.8% 향상되었고, WLI 영역에서는 1.1% 향상되었으며, 더 작은 NBI 데이터셋에서 더 큰 성과를 보였다.
• 보간된 영상들은 병변 크기의 더 조밀하고 연속적인 다양성을 분류기에게 노출시켜 결정 경계를 더욱 매끄럽게 하고 강건성을 향상시켰다.
• 이 방법은 단일 사전 훈련된 GAN 모델만으로도 다양한 레이블 유지 증강을 가능하게 하여, 각 작업마다 별도의 생성 모델을 훈련할 필요를 줄였다.