A Swin transformer encoder-based StyleGAN for unbalanced endoscopic image enhancement

Abstract

With the rapid development of artificial intelligence, automated endoscopy-assisted diagnostic systems have become an effective tool for reducing the diagnostic costs and shortening the treatment cycle of patients. Typically, the performance of these systems depends on deep learning models which are pre-trained with large-scale labeled data, for example, early gastric cancer based on endoscopic images. However, the expensive annotation and the subjectivity of the annotators lead to an insufficient and class-imbalanced endoscopic image dataset, and these datasets are detrimental to the training of deep learning models. Therefore, we proposed a Swin Transformer encoder-based StyleGAN (STE-StyleGAN) for unbalanced endoscopic image enhancement, which is composed of an adversarial learning encoder and generator. Firstly, a pre-trained Swin Transformer is introduced into the encoder to extract multi-scale features layer by layer from endoscopic images. The features are subsequently fed into a mapping block for aggregation and recombination. Secondly, a self-attention mechanism is applied to the generator, which adds detailed information of the image layer by layer through recoded features, enabling the generator to autonomously learn the coupling between different image regions. Finally, we conducted extensive experiments on a private intestinal metaplasia grading dataset from a Grade-A tertiary hospital. The experimental results show that the images generated by STE-StyleGAN are closer to the initial image distribution, achieving a Fréchet Inception Distance (FID) value of 100.4. Then, these generated images are used to enhance the initial dataset to improve the robustness of the classification model, and achieved a top accuracy of 86 %.