CycleGAN Clinical Image Augmentation Based on Mask Self-Attention Mechanism

Abstract

With the development of society and the advancement of science and technology, artificial intelligence has also emerged as the times require. In computer vision, deep learning based on convolutional neural networks(CNN) achieves state-of-the-art performance. However, the massive data requirements of deep learning have long been a pain point in the field, especially in the medical field, where it is often difficult (and sometimes impossible) to obtain enough training data for some specific tasks. To overcome insufficient and unbalanced data, in this paper, we focus on the generation and balance of data on radiation-induced pneumonia, an extremely rare disease with a low incidence. As a result, datasets on this disease are extremely sparse and unevenly distributed. To address the above problems, the predecessors’ method is often to use generative models to generate data as a complement of the fewer samples to achieve a balanced distribution of data samples. Among various generative models, CycleGAN is widely used in medical image generation due to its cycle consistency to achieve style migration without changing the basic content. However, the original CycleGAN method has many shortcomings, especially in Few-shot and the data unevenly distributed, its performance will be greatly reduced. To make the generated data samples retain the original structure and have finer and clearer details, this paper proposes a mask-based self-attention CycleGAN data augmentation method. A self-attention branch is added to the generator and two different loss functions named Self-Attention Loss and Mask Loss are designed. To stabilize the training process, spectral normalization is introduced to improve the discriminator and MS-SSIM and L1 joint loss are used to improve the original identity loss. The ResNet18 is used to complete classification experiments on the radiation-induced pneumonia dataset and the COVID-19 dataset respectively. Four classification performance indicators: the area under the ROC curve (AUC), Accuracy (ACC), Sensitivity (SEN), and Specificity (SPE) are calculated to verify the effectiveness and generalization of our method. Compared with the original CycleGAN and traditional data augmentation, the classifier trained by data augmentation using our method has outstanding performance in multiple classification indicators and has better classification performance. Experimental results show that our method solves the problem of insufficient samples and data imbalance in the pneumonia dataset by generating high-quality pneumonia images. Code is available at https://github.com/ngfufdrdh/CycleGAN-lung.