Optimal transport theory-based GAN for medical image augmentation and classification

Abstract

Generative Adversarial Networks (GANs) can create realistic images without relying on original data, which is particularly useful in addressing class imbalance in medical imaging due to disparities in illness occurrence. GAN-based methods have been proposed to tackle this issue, but they often struggle with small-scale disorders due to the challenge of extracting significant information from few pixels. To combat data scarcity and imbalance, data augmentation techniques, like image modifications, are employed to maintain performance. Optimal Transport (OT) theory aids in defining distance in function spaces, interpolating functions, and establishing barycenters of weighted functions. However, using OT in generative models poses challenges such as lack of smoothness, computational burden, and difficulty in estimating high-dimensional gradients. To overcome these, we propose using Sinkhorn divergence to create a loss family that combines Wasserstein (OT) and Maximum Mean Discrepancy (MMD) losses, leveraging OT's geometry for generating highdimensional spaces from low-dimensional manifolds. This approach is suitable for image augmentation and classification, ensuring the generated images are physically and aesthetically convincing. Our Optimal Transport Theory-based GAN model excels in detecting and differentiating chest X-ray images, achieving 96.94% accuracy in classifying normal, COVID-19-affected, lung opacity, and viral pneumonia images. This performance surpasses CNN models used in our research, demonstrating superior accuracy, precision, recall, and F1 score on the available datasets. The integration of OT and GANs represents a significant advancement in medical imaging, addressing critical challenges and improving diagnostic capabilities.