EfficientNet and DenseNet Ensemble with GAN Augmentation for Imbalanced COVID-19 and Viral Pneumonia Classification in Chest X-Rays

Abstract

Chest X-ray  imaging is a critical tool for diagnosing various diseases and plays a significant role in computer-aided diagnosis (CAD) systems. Recent research has leveraged CXR datasets and deep learning models to detect multiple diseases; however, these studies often suffer from dataset imbalance, where one class is overrepresented. This imbalance hampers deep learning model training, leading to models that perform well on the majority class but struggle with underrepresented classes. To address this challenge in CXR datasets comprising three classes Covid-19, Viral pneumonia, and Normal images, we employ a generative adversarial network to generate synthetic CXR images for the minority classes, thereby enhancing model robustness. Additionally, we fine-tune DenseNet201 and EfficientNetV2-B3 and integrate their predictions using a soft voting ensemble. Our method reaches an accuracy of 99.33%, an F1-score of 99.01%, and a Matthews correlation coefficient (MCC) of 98.59%, which helps reduce the problems caused by uneven class distribution. These results show that using GAN-based augmentation with an ensemble of deep learning methods can enhance the classification of different types of CXR, leading to more trustworthy AI-assisted diagnoses.