[논문 리뷰] Performance of GAN-based augmentation for deep learning COVID-19 image classification

Oleksandr Fedoruk, K. Klimaszewski외 2명|arXiv (Cornell University)|2023. 04. 18.

COVID-19 diagnosis using AI참고 문헌 51인용 수 4

한 줄 요약

이 연구는 스타일 제너레이티브 어드버서리 모델(StyleGAN2-ADA)을 사용한 GAN 기반 데이터 증강 기법을 통해 코로나19 흉부 X-ray 이미지의 딥러닝 분류 성능을 향상시키는지를 평가한다. 네 가지 클래스(코로나19, 바이러스성 폐렴, 박테리아성 폐렴, 정상)로 구성된 데이터셋을 대상으로 한다. 고품질의 합성 이미지를 생성했음에도 불구하고, GAN 기반 증강은 전통적인 데이터 증강 기법과 단순한 클래스 균형 조정 방식보다 성능이 열 劣했으며, 효율넷-B0 모델은 증강 없이 균형 잡힌 데이터에서 90.2%의 정확도를 기록했다. 이는 이 데이터셋에 대해 추가적인 GAN 기반 합성 데이터가 유익하지 않으며, 오히려 성능 저하를 초래할 수 있음을 시사한다.

ABSTRACT

The biggest challenge in the application of deep learning to the medical domain is the availability of training data. Data augmentation is a typical methodology used in machine learning when confronted with a limited data set. In a classical approach image transformations i.e. rotations, cropping and brightness changes are used. In this work, a StyleGAN2-ADA model of Generative Adversarial Networks is trained on the limited COVID-19 chest X-ray image set. After assessing the quality of generated images they are used to increase the training data set improving its balance between classes. We consider the multi-class classification problem of chest X-ray images including the COVID-19 positive class that hasn't been yet thoroughly explored in the literature. Results of transfer learning-based classification of COVID-19 chest X-ray images are presented. The performance of several deep convolutional neural network models is compared. The impact on the detection performance of classical image augmentations i.e. rotations, cropping, and brightness changes are studied. Furthermore, classical image augmentation is compared with GAN-based augmentation. The most accurate model is an EfficientNet-B0 with an accuracy of 90.2 percent, trained on a dataset with a simple class balancing. The GAN augmentation approach is found to be subpar to classical methods for the considered dataset.

연구 동기 및 목표

• 딥러닝을 통한 코로나19 분류에 있어 제한적이고 불균형한 의료 영상 데이터 문제를 해결하기 위한 목적이 있다.
• 전통적 증강 기법을 초월해 GAN 기반 증강 기법이 다중 클래스 코로나19 X-ray 분류 성능 향상에 기여하는지 조사하기 위한 목적이 있다.
• 다양한 이미지 유사도 측정 지표를 활용해 GAN으로 생성된 합성 X-ray 이미지의 품질과 대표성을 검증하기 위한 목적이 있다.
• 다양한 데이터 증강 전략 하에서 여러 딥러닝 모델(효율넷-B0, 인셉션-v3)의 성능을 비교하기 위한 목적이 있다.
• 시각화 지도를 활용한 모델의 해석 가능성 평가를 통해 모델이 생물학적으로 타당한 특징을 학습하는지, 또는 '단서'에 의존하는지 평가하기 위한 목적이 있다.

제안 방법

• 4개 클래스의 흉부 X-ray 이미지(코로나19, 바이러스성 폐렴, 박테리아성 폐렴, 정상)로 구성된 제한적이고 불균형한 데이터셋을 기반으로 스타일 제너레이티브 어드버서리 모델(StyleGAN2-ADA)을 훈련하여 합성 이미지를 생성한다.
• FID, RMSE, SRE, SSIM 등의 지표를 활용해 GAN으로 생성된 이미지의 품질을 평가하며, 내부 클래스 및 외부 클래스 간 유사도 분포를 비교한다.
• 전통적 데이터 증강 기법(회전, 자르기, 밝기 변화 등)을 적용하고, 이를 GAN 기반 증강과 여러 훈련 시나리오에서 비교한다.
• 균형 잡힌 데이터셋과 증강된 데이터셋을 대상으로 전이 학습 기법을 활용해 두 모델(효율넷-B0, 인셉션-v3)을 미세조정하고 평가한다.
• 다양한 증강 전략 간 성능 차이의 통계적 유의성을 평가하기 위해 스타트-맥스웰 검정을 사용한다.
• 시각화 지도를 활용해 모델의 주의 집중 영역을 분석하여, 단서 학습의 가능성을 탐지하고 폐 부위에 주의가 집중되는지 평가한다.

실험 결과

연구 질문

• RQ1 전통적 데이터 증강 기법과 단순한 클래스 균형 조정 방식에 비해 GAN 기반 증강 기법이 다중 클래스 코로나19 X-ray 분류 성능 향상에 기여하는가?
• RQ2 FID 및 유사도 측정 지표로 측정했을 때, GAN으로 생성된 합성 X-ray 이미지가 실제 이미지의 분포와 품질 면에서 얼마나 유사한가?
• RQ3 GAN 증강을 통해 관찰된 성능 저하가 낮은 품질의 합성 이미지 때문인지, 또는 데이터 분포 자체의 본질적 한계 때문인가?
• RQ4 GAN 증강 데이터로 훈련된 딥러닝 모델이 폐 부위에 의미 있는 주의 패턴을 학습하는가, 아니면 비의미적인 상관관계에 의존하는가?
• RQ5 시각화 지도 분석을 통해 모델이 폐와 같은 생물학적으로 관련 있는 영역에 집중하는지, 아니면 피부나 횡격막과 같은 구조물이나 잡음에 집중하는지 확인할 수 있는가?

주요 결과

• 가장 높은 정확도를 기록한 모델은 증강 없이 균형 잡힌 데이터셋으로 훈련된 효율넷-B0 모델로, 정확도는 90.2%였다.
• 모든 평가 모델과 지표에서 GAN 기반 증강은 전통적 증강 기법과 단순한 클래스 균형 조정 방식보다 일관되게 성능이 열 劣했다.
• 전통적 증강 기법은 인셉션-v3의 성능 향상에 기여했지만, 이미 높은 정확도를 기록한 효율넷-B0에는 추가적인 이득이 없었다.
• 이미지 품질 측정 지표(FID, RMSE, SRE, SSIM)는 GAN으로 생성된 이미지가 실제 이미지와 매우 유사함을 확인했으며, 이는 고품질의 합성 결과임을 시사한다.
• 시각화 지도 분석 결과, 모델은 주로 폐 부위에 주의를 기울였지만, 횡격막과 같은 진단에 무관한 영역에도 일부 주의가 집중되는 경향을 보였으며, 이는 최소한의 단서 학습이 존재함을 시사한다.
• 이 연구는 이 특정 데이터셋에 대해 GAN 증강의 복잡성이 성능 향상으로 이어지지 않음을 발견했으며, 균형 잡힌 데이터셋이 충분히 제공된다면 추가적인 합성 데이터가 유익하지 않을 수 있음을 시사한다.