Autoencoder-based conditional optimal transport generative adversarial network for medical image generation
Jun Wang, Bohan Lei, Liya Ding, Xiaoyin Xu, Xianfeng Gu, Min Zhang
Abstract
Recently, there has been a significant surge of interest in medical image generation. In this study, we developed a model known as AE-COT-GAN (autoencoder-based conditional optimal transport generative adversarial network) to generate medical images that belong to specific categories. The primary objective of our research is to address the prevalent challenges often encountered during the training of generative adversarial networks (GANs), including issues such as mode collapse and mode mixing. The training process of our model encompasses three fundamental components. First, we employ an autoencoder model to obtain a low-dimensional manifold representation of real images. Second, we apply extended semi-discrete optimal transport to map Gaussian noise distribution to the latent space distribution and obtain corresponding labels effectively. This procedure leads to the generation of new latent codes with known labels. Finally, we integrate a GAN to train the decoder further to generate medical images. To evaluate the performance of the AE-COT-GAN model, we conducted experiments on two medical image datasets, namely DermaMNIST and BloodMNIST. The model’s performance was compared with state-of-the-art generative models. Results show that the AE-COT-GAN model had excellent performance in generating medical images. Moreover, it effectively addressed the common issues associated with traditional GANs.