AttCL-GAN: Attentional contrastive learning-based generative adversarial network for modality completion of medical images

Abstract

Most multi-modal medical image analysis models require datasets to be modal-balanced. However, existing multi-modal medical datasets often suffer from modality missing problems due to practical factors. Therefore, in this study, we propose a novel generative model, AttCL-GAN, to remedy three issues encountered when applying the multi-domain model StarGAN v2 to medical images generation. Further by utilizing the proposed modality completion strategy, we achieve the data imputation under arbitrary modality missing scenarios. The advancements of AttCL-GAN are mainly threefold: Firstly, we integrate the idea of contrastive learning and attention mechanism to guide the network in learning domain-independent content during image translations using attention information from spatial and channel dimensions. Then, the attentional AdaIN generator is introduced into the network. By adding the refined features and original features together and feeding them into AdaIN, the generator's ability to generate tissue details is enhanced. Finally, to alleviate the problem of high inter-modal similarity of generated images, we propose a style code diversity loss, which increases the diversity of images by enlarging the Euclidean distance between codes of different modalities in the latent space. Extensive experimental results on a real-world multi-modal brain MRI dataset show that (i) The proposed AttCL-GAN significantly outperforms the state-of-the-art GAN-based data augmentation methods in both the generation and segmentation tasks in terms of all metrics; (ii) The proposed three advancements are all effective and essential for AttCL-GAN to achieve the superior performances in both tasks.