A method for classifying medical images of covid-19 infection based on self-supervised learning

Abstract

A kind of COVID-19 infection medical image classification method based on self-supervised learning belongs to image classification technical field, based on MAE self-supervised method, and is supplemented with self-distillation learning, constructs SSL-Covid model, SSL-Covid model learns domain knowledge from data by self-supervised learning, and is added with self-distillation learning to improve model feature extraction capability, provides auxiliary information for COVID-19 infection classification;MAE is introduced into this field, and corresponding improvement is carried out, a self-distillation module is added, and a self-supervised learning COVID-19 medical image classification model SSL-Covid is constructed, which aims to shorten the diagnosis time of COVID-19 patients by using computed tomography CT, reduce the duration and number of disease spread and provide decision support system for radiologists to diagnose COVID-19.The SSL-Covid model of the application has more advantages than MAE in the recognition of COVID-19 infection in CT image, and well relieves the problem of large-scale data demand, improves the performance of ViT in medical image classification.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of image classification, and particularly relates to a COVID-19 infection medical image classification method based on self-supervised learning. BACKGROUND

[0002] Corona Virus Disease 2019 (COVID-19) is a highly contagious disease. Efficient and rapid diagnosis of COVID-19 is an effective way to block its spread and control its transmission. Current diagnostic techniques (such as reverse transcription polymerase chain reaction (RT-PCR), serum testing) are expensive, time-consuming, and require a strict experimental environment, making them difficult to implement in other more remote areas. At present, COVID-19 is a highly contagious disease, and can be clinically diagnosed based on radiodiagnosis. Deep learning can mine the rich information hidden in patient image data and complete classification of different stages of the disease course. However, a large amount of training data is one of the important conditions for training an excellent deep learning model. Unfortunately, the scarcity of medical field data due to privacy, labeling difficulties, and other issues is a common phenomenon, especially the labeled data of COVID-19 infection is extremely scarce. Therefore, we introduce the Masked Autoencoder (MAE) to directly pre-train and fine-tune on small-scale target datasets. On this basis, we propose a new model (Self-Supervised Learning on COVID-19 classification, SSL-Covid), which, in addition to the reconstruction loss calculation of the masked image blocks, further applies a self-distillation loss to the latent representation output by the encoder and decoder. The additional loss calculation can transfer the knowledge brought by the global attention of the decoder to the encoder which only gets local attention. Our model achieves an identification accuracy of 97.78% on the SARS-COV-2-CT dataset containing 2481 images, and further verifies the identification accuracy of 81.76% on the COVID-CT dataset containing 746 images. The experimental results show that the deep learning model and classification method studied in the application can effectively mine the COVID-19 infection features from rare data, and can be used as an auxiliary decision support system to improve the detection efficiency of COVID-19 disease.

[0003] Rapid diagnosis of COVID-19 infection is one of the effective ways to block the spread of COVID-19 infection. Currently, reverse transcription polymerase chain reaction (RT-PCR) is considered the gold standard for diagnosing COVID-19 infection. The rapid increase in patients and the delay in PCR test results and its requirements for testing environment make it extremely challenging to rapidly diagnose COVID-19 infection. In addition, RT-PCR has a high false negative rate, leading to many COVID-19 positive cases being diagnosed as negative, posing greater risks. In particular, studies have shown that the sensitivity of PCR testing is 71%, far lower than the target sensitivity value for effectively containing the spread of the virus.

[0004] Radiological imaging by computed tomography (CT) has become a promising alternative diagnostic form due to its ability to visualize lung structures. Related studies have shown that CT scan images are considered a better method for detecting COVID-19 infection, with a sensitivity of up to 98%, compared to the 71% sensitivity of RT-PCR. In COVID-19 cases, CT scan images show some specific manifestations, including bilateral, peripheral or posterior distribution of multiple lobular ground-glass opacities (GGO), mainly in the lower lobe, and less in the middle lobe. Diffuse distribution, blood vessel thickening and fine reticular opacity are other common features reported for COVID-19. Some CT images of COVID-19 and non-COVID-19 infections are shown in Figures 1-2 However, studying and judging scan results by artificial methods is very time-consuming and prone to misjudgment. In contrast, deep learning can automatically extract features from data, reducing the dependence on artificial judgment and better reducing the negative impact on diagnosis results due to human factors, which provides a good way to improve diagnosis efficiency and reduce misdiagnosis rate.

[0005] In recent years, the successful application of deep learning in computer vision and biomedical fields has greatly stimulated researchers to explore the application of deep learning in CT images. Some researchers have conducted related research on CT images based on deep learning methods. In the biomedical field, medical images are usually classified using CNN-based network structures (e.g., ResNet). The common practice is: first, encode the image into a high-level feature representation, then extract the relevant features of the image through the network, and finally classify the image by fusing its features or decisions. With the automatic feature learning ability of CNN, COVID-19 classification based on deep neural networks has been widely applied.

[0006] In recent years, with the development of deep learning, self-supervised learning (SSL) has also received more and more attention. SSL aims to train a model on unlabeled data through the design of a pretext task. In fact, it is to train an unsupervised representation learning model to obtain a general representation of data or a representation for a specific task, and then fine-tune it for downstream tasks. A large number of researchers have proved that this training method can provide more effective weight representation for downstream tasks. In addition, SSL is also very suitable for the field of medical images. So far, supervised learning is still the main technology for many medical image analysis. However, it relies on time-consuming and laborious manual data labeling, which is a relatively difficult engineering in the medical field, and SSL can well alleviate this problem.

[0007] Based on the above analysis, we studied a new coronavirus infection medical image classification task based on a masked autoencoder (MAE) self-supervised pre-training model with self-distillation. Unlike the usual self-supervised method, we added a self-distillation optimization method to it, and additionally introduced label information, so that the model can not only know what concept it wants to reconstruct, but also guide the self-distillation to provide more effective information gain to the student model. By comparing the results obtained by the model with those based on VGG19, ResNet101, DenseNet169, MAE and BEiT, the performance of the proposed technology was evaluated. The experimental results show that the model has better feature learning ability. SUMMARY

[0008] In order to solve the above problems, the present application proposes: a new coronavirus infection medical image classification method based on self-supervised learning, based on the masked autoencoder MAE self-supervised method, supplemented by self-distillation learning, and constructs an SSL-Covid model. The SSL-Covid model learns domain knowledge from data through self-supervised learning, and improves the feature extraction ability of the model through self-distillation learning, and provides auxiliary information for new coronavirus infection classification;

[0009] The overall model framework is divided into two stages: a pre-training stage and a fine-tuning stage, and both stages are performed on the target data set without introducing external data. The pre-training stage consists of two auxiliary tasks: image reconstruction and self-distillation modeling. Image reconstruction completes the reconstruction task of the mask image, promotes the encoder to better learn the domain knowledge, and self-distillation modeling further applies distillation constraints to the former to optimize the feature extraction ability of the encoder.

[0010] After pre-training, the learned weights of the encoder are migrated to the fine-tuning stage, and then the classification of new coronavirus infection is performed.

[0011] Further, a visual self-attention model ViT is composed of a patch embedding layer, a position embedding layer and a Transformer Block;

[0012] Patch embedding layer: ViT takes a sequence as input, and patch converts any high-dimensional data into a sequence form. For a multi-modal image where X, H, W, C are the vector representation of the image, real matrix, image height, image width and image channel number, after a series of image preprocessing, the image is processed into a unified resolution size of 224x224, which is passed through a convolution layer and straightened into where P is the resolution size of each image block, N=HW / P 2 is the number of image blocks, x p is the reshaped vector representation, then a trainable linear projection is used to map them to D-dimensional vectors, i.e. the length D=P 2 C.

[0013] Position embedding layer: position embedding information is added at the patch embedding. The standard ViT uses 1D learnable position embedding, i.e. learnable position embedding

[0014] Self-attention block Transformer Block: a Transformer Block is composed of multi-head attention MSA, multi-layer perception MLP and layer normalization LN alternately;

[0015] First, the input image is passed through the Patch embedding layer, then it is spliced with a class token for classification, and the obtained patches plus position embedding E pos are sent into the Transformer Encoder module for operation, and finally the class token in the result is extracted and sent into the MLP module for classification. The calculation process is as follows:

[0016]

[0017] S′ l = MSA(LN(S l-1 ))+S l-1 , l=1, 2, …, L (2)

[0018] S l = MLP(LN(S′ l ))+S′ l , l=1, 2, …, L (3)

[0019]

[0020] where S0denotes the vector representation of the image after patch embedding layer and positional embedding processing, S' l denotes the vector representation after processing by the Transformer Blocks, S l denotes the vector representation after processing by the layer normalization and multi-layer perception, y denotes the final prediction output.

[0021] Further, the masked image modeling process is as follows: the encoder in MAE uses the Transformer Block in ViT, which first divides the input image into non-overlapping image blocks Then after linear mapping to get patch embedding, then random mask is performed at a certain proportion, define visible patches and invisible patches as Z vis , Z mask , retain visible patches Z vis plus position information as the input of the encoder, process, and then get the corresponding latent representation for the subsequent image reconstruction;

[0022] The decoder in MAE uses a complete token set as input, which is composed of the encoding vector output by the encoder and the mask token. For the mask token, it is a shared, learnable vector, which represents the missing patch to be predicted, and for all tokens, plus the corresponding position embedding is recovered to the patch at each specific masking position through the decoder model.

[0023] MAE uses the mean square error MSE to calculate the reconstruction loss between the predicted pixel value y pred and the original pixel value y of the unmasked area:

[0024] L MSE = MSE(y pred , y) (5).

[0025] where L MSE denotes the mean square error loss between the true pixel and the predicted pixel

[0026] Further, two supervised branches are introduced in the pre-training stage, respectively located after the encoder and the decoder, branch Student and branch Teacher, and relative to the local processing of the encoder part patches, the decoder processes all patches more globally, the result after the decoder is used to map the vector through the Teacher branch to the result after the encoder through the Student branch to apply distillation constraints, and the global information is added to the encoder to further improve the feature extraction capability of the encoder;

[0027] For the two branches, only one layer of MLP is used for the Student branch at the encoder end and the Teacher branch at the decoder end, and the output vector Z and the output vector Z all after the decoder are respectively processed through the branches at the respective ends to obtain two high-dimensional vectors V1 and V2, respectively, and finally the two vector distributions are matched by minimizing the cross-entropy loss, and the process is expressed as:

[0028]

[0029] V2=MLP(Z all ) (7)

[0030]

[0031] where L CE represents the cross-entropy loss function, y is the true label, and L distill represents the distillation loss.

[0032] The calculation formula of the total loss is as follows:

[0033] L=L MSE +L distill (9)。

[0034] The beneficial effects of this invention are as follows: It introduces MAE (Medical Image Classification) into this field and improves it accordingly, incorporating a self-distillation module to construct a self-supervised learning COVID-19 infection medical image classification model (SSL-Covid). This aims to shorten the diagnosis time for COVID-19 patients by using computed tomography (CT) scans, reduce the duration and number of cases spreading, and provide a decision support system for radiologists diagnosing COVID-19. In this invention, global information is introduced into self-supervised learning by adding labels. Further self-distillation transfers the global information learned after decoding to the encoder, optimizing and improving the encoder's feature extraction capabilities. Related experiments demonstrate that the SSL-Covid model has advantages over MAE in identifying COVID-19 infection in CT images and effectively alleviates the problem of large-scale data requirements in the ViT architecture. In summary, these experimental results show that SSL-Covid can further improve the performance of ViT in medical image classification. Attached Figure Description

[0035] Figure 1 The COVID-19 CT image is from this invention;

[0036] Figure 2 The images shown are of non-COVID-19 infected individuals according to the present invention.

[0037] Figure 3 This is a schematic diagram of the overall model framework of the present invention;

[0038] Figure 4 This is a diagram of the ViT model framework of the present invention;

[0039] Figure 5 This is a schematic diagram of the MAE model framework of the present invention;

[0040] Figure 6 This is a schematic diagram of the self-distillation structure of the present invention;

[0041] Figure 7 This is the confusion matrix 1 for the SSL-Covid model of this invention;

[0042] Figure 8 This is the confusion matrix 2 of the BEiT model of this invention;

[0043] Figure 9 This is the confusion matrix 3 of the MAE model of this invention;

[0044] Figure 10 This is the confusion matrix 4 of the DenseNet169 model of this invention;

[0045] Figure 11 The confusion matrix of the VGG19 model of this invention is shown in Figure 5.

[0046] Figure 12 Original picture of the present invention;

[0047] Figure 13 Grad-CAM heat map of the original picture of the present invention;

[0048] Figure 14 SSL-Covid original picture of the present invention;

[0049] Figure 15 Grad-CAM heat map of the SSL-Covid original picture of the present invention;

[0050] Figure 16 MAE and SSL-Covid fine-tuning performance comparison curve of the present invention. DETAILED DESCRIPTION

[0051] A new coronavirus infection medical image classification method based on self-supervised learning is proposed, and the relevant work is as follows:

[0052] In this part, the research fields related to our research work--self-supervised learning, knowledge distillation and coronavirus infection detection are discussed in detail.

[0053] Self-supervised learning

[0054] In deep learning, most of the training of the model is based on supervised learning, that is, there are a large number of labeled data. However, in some fields, it may not be like the general field to have a relatively large labeled data set (such as ImageNet). Therefore, the pre-training model can well alleviate this problem. The model is usually pre-trained on a large data set, and then fine-tuned on other visual tasks. The reasons for doing so are mainly the following two points: after training on a large data set, the network parameters can be used as the starting point for training other tasks to converge faster; after training, the network can effectively extract hierarchical semantic information, thereby preventing overfitting on small data sets. Thus, the performance of the network depends largely on the size of the labeled data set. However, at the same time, a contradiction needs to be faced, which is that in the medical field, there is no large-scale data set like the general field ImageNet for pre-training, and due to patient privacy and labeling quality and other problems, it is very time-consuming and expensive to collect and label a large amount of data. In order to solve this problem, self-supervised learning is proposed, which can be trained on a data set without human labeling.

[0055] Self-supervised learning alleviates the problem of insufficient labeled data by using innovative auxiliary tasks in the way of active learning, anomaly detection, data augmentation and registration. The self-supervised method can be regarded as a special form of unsupervised learning method with a supervised manner, and the supervision here is induced by the self-supervised task rather than the preset prior knowledge. Compared with the completely unsupervised setting, self-supervised learning uses the information of the dataset itself to construct pseudo-labels. In terms of representation learning, self-supervised learning has great potential to replace fully supervised learning. Its general research idea is: first, propose an auxiliary task, then use the auxiliary task to automatically generate pseudo-labels for images, these pseudo-labels are obtained from the attributes of the image, do not involve any manual annotation, and are used as the supervision information of the image to train the computer vision model. The trained computer vision model can be migrated to the target data domain or target task. In terms of representation learning, self-supervised learning has great potential to replace fully supervised learning. The essence of human learning tells us that large annotation datasets may not be necessary, and we can learn from unlabeled datasets spontaneously.

[0056] Jing et al. investigated the self-supervised general visual feature learning methods based on deep learning, summarized the development of self-supervised learning in recent years, and proposed some future directions for the development of self-supervised learning. In the field of medical images, it is usually difficult to obtain enough labeled images for training. In some cases, the data set used in the study may contain more unlabeled data than labeled data. In this regard, Chen et al. proposed a new self-supervised learning strategy based on context restoration, which makes better use of unlabeled data to bring more meaningful semantic image features for subsequent image analysis tasks. This method takes background restoration as a self-supervised task. In detail, given an image, randomly select and exchange two small blocks. Repeat this operation several times, and the intensity distribution of the new image remains unchanged, but its spatial information is changed. Then train the CNN structure model to restore the changed image to the original version. The experimental results are good on classification, positioning and segmentation tasks. Most existing few-shot image classification methods are more dependent on efficient embedding networks, but due to the limited labeled data, the size of the embedding network is limited, making the few-shot learning reach a bottleneck. Chen et al. studied a more general embedding network with self-supervised learning (SSL) for the problem of limited data in few-shot image classification tasks. This network model can learn from the data itself to provide robust representations for downstream tasks, and the model performs well through experiments on few-shot classification data sets. By taking advantage of the complementarity of few-shot learning and self-supervised learning, the former model can effectively learn to recognize patterns in low-data states, and the latter can learn more general representations using unlabeled data. Gidaris et al. proposed a method that combines supervised loss and self-supervised loss to pretrain the embedding network. Masood et al. proposed a semi-supervised, self-suggested learning model for automatically identifying melanoma using dermoscopy images. Data augmentation is a common training technique that usually uses linear changes, which do not change the high-level semantic information in the image, are not sufficient to effectively summarize the input distribution, and cannot fully exploit the performance potential of neural networks. Based on this, Xue et al. proposed a new data augmentation method by using MAE to reconstruct the distorted view of the input image to achieve the purpose of expanding the data volume. Using this model-based nonlinear transformation as data augmentation, it performs well in supervised, semi-supervised and few-shot classification tasks.

[0057] Knowledge distillation

[0058] With the rise and evolution of deep neural networks, knowledge distillation (KD) as a new and general model compression and transfer learning architecture, due to its simplicity and effectiveness, plays an increasingly important role in various fields of artificial intelligence, and is a general means to solve many practical problems. KD was first defined by Ba and Caruana, and then popularized by Hinton et al. Its training method based on the "teacher-student network idea" distills the knowledge contained in the already trained model (teacher network) into another model (student network), and enhances the learning ability of the student model by transferring knowledge. Generally speaking, the teacher network is a large-capacity model, while the student network is a smaller-capacity model. In other words, KD refers to a method that helps a smaller student network to train under the supervision of a larger teacher network. Unlike other compression methods, KD can ignore the structural differences between the two networks and compress the student network. The structural differences between the teacher and student networks can be ignored to reduce the size of the network. Knowledge distillation can be roughly divided into two branches: logits distill and intermediate representation distillation.

[0059] Deep learning has made remarkable achievements in many fields, such as computer vision (CV) and natural language processing (NLP). One of the keys to its success is that deep learning models are relatively large and can learn complex patterns and features in data. However, there is a big problem, that is, the leading large model is difficult to deploy on terminal devices. The goal of knowledge distillation is to provide smaller models to solve the same tasks as larger models. Based on this, more and more researchers have focused on knowledge distillation.

[0060] For knowledge distillation within the same model, Zhang et al. proposed a self-distillation framework that extracts knowledge within the network itself. It extracts knowledge from deep layers to shallow layers, enhancing the feature representation of shallow layers. Heo et al. proposed a novel feature distillation method to make parts including teacher transformation, student transformation, distillation feature position, and distance function work together. The proposed distillation loss includes feature transformation with newly designed margin ReLU, new distillation feature position, and partial L2 distance function to achieve the purpose of skipping redundant information that adversely affects student compression. Touvron et al. introduced a teacher-student strategy for transformers, which mainly relies on a distillation token to ensure that students can learn knowledge from teachers through attention. Hinton et al. transferred knowledge from the teacher model to the student model by minimizing the difference between the logits produced by the teacher model (the input to the final softmax) and the logits produced by the student model. Wen et al. observed that training with teacher logits can be affected by incorrect and uncertain supervision, so they proposed two new methods (Logits Adjustment and Dynamic Temperature Distillation) to handle incorrect logits and uncertain logits, respectively. The main idea is to correct the incorrect predictions (knowledge) of the teacher through smooth regularization, and avoid using dynamic temperature for excessive uncertain supervision. Liu et al. focused on making the student independent of the teacher, so that it can receive “knowledge” from it without being constrained by the teacher’s network structure, output space, and task. Mishra et al. combined the use of low-precision numerical and model compression techniques for knowledge distillation, and showed that the use of knowledge distillation techniques can significantly improve the performance of low-precision networks.

[0061] New coronavirus infection image recognition model

[0062] Deep learning-based methods are more convenient and fast, and researchers have developed many effective COVID-19 detection models based on this.

[0063] Wang et al. modified the InceptionV3 structure and fine-tuned the modified structure using pre-trained weights. The model achieved a significant result, with an accuracy of 89.5% in the experimental test. Heidarian et al. developed a two-stage fully automated CT framework (COVID-FACT), which is mainly composed of a capsule network. COVID-FACT does not require extensive data augmentation and large datasets to capture spatial information, and it operates in two stages: the first stage detects infected slices, and the second stage classifies patient CT scan images. Compared with similar models, it has less dependence on data supervision and annotation. Meng et al. used the transfer learning method and used the Vision Transformer as the backbone model to build a two-stage transfer learning COVID-19 recognition framework (TL-Med), which greatly alleviated the problem of data insufficiency. Shui-Hua et al. proposed a rank-based average pooling module (NRAPM) and inspired by the VGG network, they proposed an average pooling network (DRAPNet) model, i.e., a deep rank-based average pooling network, for COVID-19 recognition. The F1 value of the proposed model reached 95.49% on 1164 CT images. With the continuous growth of COVID-19 cases, the importance of deep learning research in this field has become increasingly apparent. Garg et al. conducted experiments on COVID-19 detection by initializing 20 pre-trained weight CNN structures in neural networks to study the differences between different models in order to determine the most accurate model for COVID-19 recognition. In CNN, convolution operations are good at extracting local features, but they have certain limitations in capturing global feature representations. In Vision Transformer, cascading self-attention modules can capture long-range feature dependencies, but they ignore the details of local features. Fan et al. combined the local feature extraction capabilities of convolutional neural networks and the global feature extraction advantages of Transformers to propose a parallel dual-branch model (Trans-CNNNet) based on Transformer modules and convolutional neural network modules. Through experiments, it was verified that the classification accuracy of a large-scale COVID-19 dataset (COVIDx-CT) reached 96.7%. Based on the model's generalization performance, Ravi et al. proposed a large-scale learning model for COVID-19 classification. First, they extracted the features of the second-to-last layer of the pre-trained EfficientNet model, supplemented by using kernel principal component analysis (PCA) to reduce the dimensionality of the extracted features, and then used a feature fusion method to further fuse various extracted features. Finally, a method based on a stacked ensemble meta-classifier was used for classification.Through the use of pre-trained EfficientNet, kernel principal component analysis (PCA) and feature fusion methods, the features are continuously optimized, and finally, the optimized features are sent to the method based on the stacked ensemble meta-classifier for classification. Through experiments, it is verified that the classification performance of the method is good.

[0064] The above COVID-19 classification method is either tested on a large-scale labeled COVID-19 dataset or uses a method of introducing external knowledge through transfer learning. However, we start from another angle, fully excavate the characteristics of the data itself by constructing an image reconstruction auxiliary task and self-distilling between the encoder and decoder output representations, to improve the feature extraction capability of the model, and greatly alleviate the problem of insufficient labeled data.

[0065] The research is based on the MAE self-supervised method, and is assisted by self-distillation learning. Through self-supervised learning, the model can effectively learn domain knowledge from data, and further improve the feature extraction capability of the model through self-distillation learning, to provide good auxiliary information for COVID-19 infection classification.

[0066] The overall model framework researched by the application is as shown in Figure 3 It mainly includes two stages: a pre-training stage and a fine-tuning stage, and both stages are performed on the target dataset without introducing external data. The pre-training stage is composed of two auxiliary tasks: image reconstruction and self-distillation modeling. The former completes the reconstruction task of the mask image, and promotes the encoder to better learn the domain knowledge. The latter further applies distillation constraints to the former to optimize the feature extraction capability of the encoder. After pre-training, the learned weight of the encoder is migrated to the fine-tuning stage (the encoder and the fine-tuning stage use the same model), and then the COVID-19 infection classification is performed. The detailed experimental process is introduced as follows.

[0067] Vision Transformer

[0068] A ViT is mainly composed of a patch embedding layer, a position embedding layer and a Transformer Block. The structure is as shown in Figure 4 .

[0069] Patch embedding layer: since ViT takes sequence as input, patch needs to convert any high-dimensional data into sequence form. For a multi-modal image After a series of image preprocessing, the image is processed into a uniform resolution size of 224*224, which is passed through a convolutional layer and flattened into where P is the resolution size of each image block, and N=HWP 2For the number of image patches, they are then mapped to a D-dimensional vector using a trainable linear projection, i.e., length D = P 2 C.

[0070] Position Embedding Layer: Add position embedding information at patch embedding. The standard ViT uses 1D learnable position embedding, i.e., learnable position embedding

[0071] Transformer Blocks: A Transformer Block consists of multi-head self-attention (MSA), multi-layer perceptron (MLP) and layer normalization (LN) alternately.

[0072] Compared with CNN structure model, the attention mechanism centered on the Transformer can capture long-range relationships, which is particularly important for the positioning and identification of lesions in the research of the present application. Similar to the Transformer in the field of natural language processing (NLP), in the field of computer vision (CV), we also process a picture into a sequence mode. In order to process the 2D shape of the picture, we first convert the input image Through the Patch Embedding Layer, then splice it with a class token for classification, and obtain the patches plus the position embedding E pos together into the Transformer Encoder module for operation, and finally extract the class token in the result to send it to the MLP module for classification. The calculation process is as follows:

[0073]

[0074] S′ l = MSA(LN(S l-1 ))+ S l-1 , l = 1, 2, …, L (2)

[0076] S l = MLP(LN(S′ l ))+ S′ l , l = 1, 2, …, L (3)

[0078]

[0079] Masked Image Modeling

[0080] In this part, we will introduce the encoder, decoder and loss function in MAE. The structure is as shown in Figure 5 ​

[0081] The encoder in Masked Autoencoder (MAE) uses the same Transformer Block as in ViT. Similar to the ViT processing, the input image is first divided into non-overlapping image blocks Then the patch embedding is obtained by linear mapping, and then random mask is performed according to a certain proportion. We define the visible patches and invisible patches as Z vis , Z mask . The visible patches Z vis are retained and added with position information as the input of the encoder for processing, and then the corresponding latent representation is obtained for image reconstruction.

[0082] The decoder in MAE uses a complete token set as input, which is composed of the encoding vector output by the encoder and the mask token. For the mask token, it is a shared and learnable vector representing the missing patch to be predicted. And for all tokens, the corresponding position embedding is added. If the position embedding is not added, the mask token will not have the position information of the corresponding patch, resulting in poor image reconstruction effect. Through the decoder, the model can restore the patch at each specific masking position. In order to reduce the pre-training budget, the decoder is designed to be more lightweight.

[0083] MAE uses mean square error (MSE) to calculate the reconstruction loss by predicting the pixel value y pred of the masked image block and the original pixel value y

[0084] L MSE = MSE (y pred , y) (5)

[0086] Self-distillation modeling

[0087] The self-distillation model structure used in the present application is shown in Figure 6 Compared with the high semanticity of language, images have a large amount of spatial redundant signals. Therefore, we believe that a part of the masked image area, the remaining visible part can still well express this image. Just like we humans watching a defective image, we can still well distinguish the main object. And in MAE, the visible part of the image is used for training, but we believe that this is the lack of global features, that is: not due to training on the whole image. And for global feature learning, a direct way is to add supervised labels, allowing the network to know what concept to reconstruct. Based on this, we introduce two supervised branches in the pre-training stage, which are located after the encoder and the decoder, respectively, asFigure 2 The Student and Teacher are shown. And relative to the encoder part patches (i.e. 40%) processing has locality, the decoder processing all patches is more global. Therefore, we use the result after the decoder to map the vector of the Teacher branch to the result after the encoder to apply the distillation constraint to the vector of the Student branch, using its encoder to add global information, thereby improving the feature extraction ability of the encoder.

[0088] For both branches, the Student branch on the encoder side and the Teacher branch on the decoder side only use one layer of MLP (see Tables 6b and 6c). Specifically, we map the output vector Z and the output vector Z all after the encoder and decoder, respectively, to get two high-dimensional vectors V1 and V2, respectively, and finally match the distribution of the two vectors by minimizing the cross-entropy loss, which can be expressed as:

[0089]

[0090] V2 = MLP(Z all ) (7)

[0091]

[0092] The total loss is calculated as follows:

[0093] L = L MSE + L distill (9)

[0094] Experimental results and analysis

[0095] Experimental data

[0096] SARS-COV-2-CT dataset

[0097] SARS-COV-2-CT comes from, which has a total of 2481 CT scans, including 1252 CT scans of COVID-19 and 1229 CT scans of non-COVID-19 infected patients. In this research experiment, it is the main experimental dataset and the experiments specially mentioned are all carried out on this dataset by default. We divide the dataset into training set and validation set according to the ratio of 8:2, and the data division is shown in Table 1. In the self-supervised pre-training stage, we use the training set. In the subsequent fine-tuning stage, we use the training set and the validation set for experiments.

[0098] COVID-CT dataset

[0099] COVID-CT dataset comes from COVID-CT, which contains 349 COVID-19 CT images and 397 non-COVID-19 CT images, a total of 746 CT images. We divide it into training set and validation set in the ratio of 8:2, and the data division is shown in Table 2. The same as the previous data set processing, the training set is used in the pre-training stage, and the validation set does not participate in the pre-training. Then the training set and the validation set are used in the downstream task for fine-tuning.

[0100] Table 1SARS-COV-2-CT dataset distribution

[0101]

[0102] Table 2COVID-CT dataset distribution

[0103]

[0104]

[0105] Experimental setup and evaluation criteria

[0106] In the experiment, we use the same parameter setting as MAE to train our model. The input image is resized to 224*224, and the pre-training is 200 rounds, and the fine-tuning stage is 50 rounds. We divide the image into 14*14 patches with a size of 16*16. As with most generative methods, RandomResizedCrop is the only enhancement strategy for pre-training. Compared with CNN with induction bias, ViT as a large model has stronger generalization ability, which is conducive to building a general representation model in the field of pathology. Therefore, we use ViT-B (12 tansformer blocks) as the encoder as the backbone network of our model for self-supervised learning. In the pre-training stage, we use the cross-entropy loss function as the self-distillation loss calculation. And we further conduct ablation experiments on the number of linear layers in the encoder and decoder, and the results are shown in Table 6.

[0107] The experimental goal of the present application is mainly aimed at COVID-19 infection classification. In this study, we use accuracy, precision, recall, and F1 score as evaluation indicators. Accuracy is a basic indicator of the model, which is the ratio of correctly predicted results (including positive and negative) to the total sample results. The proportion of correctly predicted positives to all predicted positives is called precision. The recall is the proportion of true positives in all predicted positives to the total actual positives. The weighted average of precision and recall is called F1 score, and its maximum value is 1 and the minimum value is 0.

[0108]

[0109]

[0110]

[0111]

[0112] Mask ratio

[0113] Table 3 shows the performance of our model in downstream fine-tuning tasks at different mask ratios. As can be seen from the figure, unlike the MAE optimal mask ratio of 75%, our model achieves the best performance in downstream fine-tuning tasks at a mask ratio of 40%, indicating that there are significant differences between COVID-19 medical images and natural images. For COVID-19 medical images, they have more abundant tissue features and higher information density than natural images, so more information is needed to guide image reconstruction. Furthermore, the distillation operation can help the encoder improve the feature extraction capability and provide more information-rich features to the decoder in the next stage, thus guiding the image reconstruction more smoothly. From this we can also see that when the mask ratio is 90%, the model can still achieve good performance in downstream fine-tuning, which also shows that even if the visible patch used for training in the pre-training stage is very small, our model can still extract useful information for downstream fine-tuning tasks.

[0114] Table 3 Fine-tuning at different mask ratios on COVID-19

[0115]

[0116] Classification model comparison

[0117] We conducted experiments on ViT, DenseNet169, VGG19, BEiT, MAE, and the experimental results are shown in Table 4. Among them, we use ViT as the baseline, which is not pre-trained by self-supervised learning, so it can represent the effectiveness of self-supervised pre-training. From Table 3, we can see that after self-supervised pre-training, the model performance is improved to a certain extent, such as MAE, BEiT, which improves the classification performance of ViT by 3.43% and 3.64% respectively, indicating that the use of self-supervised pre-training method is helpful to the performance of the model, and further proves the feasibility of the self-supervised learning method adopted in this study. On the basis of self-supervised learning, we further introduce the self-distillation scheme and construct the SSL-Covid model, which further optimizes the ability of the model to extract features and provides more powerful help for fine-tuning. Compared with ViT, ResNet101, DenseNet169, VGG19, BEiT and MAE, the classification performance of our model is improved by 8.89%, 9.5%, 6.67%, 9.9%, 5.25% and 5.46% respectively. Moreover, compared with other models, the area under the curve (Auc) of our model reaches 99.53%, indicating that the performance of our model is more superior for COVID-19 classification. And the recall rate of our model also reaches 98%, which means that the risk of missing positive cases is greatly reduced, and it further provides more powerful guarantee for case screening, greatly reducing the consumption of manpower and material resources.

[0118] In addition, inspired by the self-supervised rotation task, we introduced it into MAE to build a mixed task model MAE_rotate that combines image reconstruction and rotation angle prediction. However, through experimental comparison, we found that its classification performance is equivalent to that of a single image reconstruction task, and it does not enhance the classification performance as we expected. The reason we speculate is that after the rotation of medical images, the mask segmentation is performed, and the difference between the patches of each angle category is not large, making it difficult for the model to accurately predict the angle. Moreover, since the data volume is 4 times that of the original, the data size is large, causing the prediction error to further increase, making it difficult for the model to learn and suppressing the ability of the model to extract features.

[0119] Table 4 Comparison of different methods

[0120]

[0121] Comparison of different classification models on COVID-CT dataset

[0122] As the same as the experiment process on the SARS-COV-2-CT dataset, we also conducted relevant experiments on the COVID-CT dataset. The experimental results are shown in Table 5. As can be seen from Table 5, when the data size is small, the performance of ViT lags behind DenseNet169. The result is predictable because the Transfomer lacks some inductive biases inherent in CNN structures, such as locality and shift invariance, which leads to poor generalization, which is also proved in Table 4. In the same way in this small COVID-19 dataset, the superiority of our model is also reflected, and the classification accuracy is improved by 8.79% and 11.49% compared with MAE and BEiT respectively. Moreover, there is an interesting phenomenon: the effect of MAE_rotate is greatly improved on the small dataset (the classification accuracy is improved by 5.41% compared with MAE). The performance improvement we think the main reason is: the data size is small, and the gain effect of the encoder feature extraction capability optimization by adding the rotation prediction branch is greater than its suppression effect, so the overall performance of the model is improved, and it shows that in a certain size of data (such as the small COVID-CT dataset in the present application), the gain effect of adding the rotation self-supervised auxiliary task is greater than the suppression effect.

[0123] Table 5 Comparison of different methods on COVID-CT dataset

[0124]

[0125] Ablation experiment

[0126] In this subsection, in order to verify the effectiveness of the components we proposed, we conducted ablation experiments on SARS-COV-2-CT. The experimental settings and results are shown in Table 6.

[0127] In Table 6(a), we studied the influence of different pre-training objectives on the experimental results. Among them, rec represents mask image reconstruction, kd represents self-distillation, and ft represents fine-tuning in the downstream task. From the experimental results, for single pre-training objectives, the advantage of mask image reconstruction is slightly greater than that of self-distillation. The main reason we think is that only 60% of the visible image blocks are trained in the encoder, which can further promote the release of the potential of the model and make the model extract effective features as much as possible. Moreover, it can also reduce the redundancy information of the picture to a certain extent, thereby reducing the interference to the model and improving the performance of the model. When using two combinations, the advantages of the two can be combined, so the effect is good.

[0128] In Table 6(b) and (c), we conduct an experimental study on the number of MLP layers at the encoder and decoder ends. For the encoder end, but we use two linear layers, the accuracy of the model decreases, we speculate that when only 60% of the visible image is input into the encoder to obtain the vector containing relatively more key information, using two linear layers to transform the dimension causes the loss of information, which in turn inhibits the ability of the encoder to extract features. In the decoder, since it is involved in the encoding process of the entire image, it contains more global information, and using one linear layer can obtain linearly separable vectors. When it uses two linear layers, its ability is largely wasted in prediction, which will cause a certain loss of information transmission.

[0129] Class sensitivity analysis and visualization

[0130] From the confusion matrix of various methods in Figure 6 We can see that for several methods in the new crown infection identification method, the performance of the SSL-Covid we proposed is more superior, and the performance improvement of other models is restricted, which mainly occurs in the identification of positive cases, and above it produces a relatively large error. For this, we think that the reason why the performance of other models is restricted is that the CT image of the positive case is more complex, and the difference between each CT image is small, which makes it difficult to distinguish the lesion. For the identification of new crown infection positive cases, compared with other models, our model has great advantages. In other words, under the condition that other conditions are not much different, our model can strive to avoid the occurrence of false negatives (actually new crown infection positive cases, but predicted as negative), which can further block the spread of the spread of the new crown infection, and save the consumption of human and material resources.

[0131] From the heat map in Figure 7 We can see that the color in the lung area is relatively deep, which means that the model pays more attention to this part. In other words, our model pays more attention to a certain area of the lung, which is also what we expect, because the lesion of the new crown infection is distributed in the lung area.

[0132] Accuracy curve when fine-tuning on SARS-COV-2-CT

[0133] In Figure 8In this section, we further report the accuracy of each round of our model in the fine-tuning stage on the SARS-COV-2-CT dataset. MAE and SSL-Covid both use ViT-B / 16 and both of them have 200 rounds of pre-training. Following the experimental setup of He et al. in the fine-tuning, we load the pre-trained weights to the model and further conduct 50 rounds of fine-tuning experiments on SARS-COV-2-CT. We can see that even in the first round, our model has achieved a very high accuracy. This is because our SSL-Covid can first learn better COVID-19 related feature information through image reconstruction in the pre-training stage, and then further learn better global characteristics through self-distillation, and then gradually optimize the encoder feature extraction capability, thereby providing a better initialization for fine-tuning. And because of a good initialization, the performance of our SSL-Covid has been relatively stable, with less fluctuation.

[0134] Different pre-training rounds

[0135] In Table 7, we compare the impact of different pre-training rounds on our model. As shown in the table, we only change the pre-training rounds, and the fine-tuning task is 50 rounds. When the pre-training rounds are fewer, such as 100 rounds, we can see that SSL-Covid achieves an accuracy of 92.93%. When the pre-training rounds reach 200 rounds, the classification performance of SSL-Covid is greatly improved, reaching an accuracy of 97.78%. This shows that when the pre-training rounds are not sufficient, the model cannot fully exploit its potential. We further increase the rounds to 300, unfortunately, the performance of the model does not increase, and the classification performance is flat with 200 rounds. After that, due to the limitation of computing resources, we did not further increase the pre-training rounds for experiments. And for the purpose of saving computing resources and time cost, we use 200 rounds as our pre-training rounds.

[0136] Table 6 SSL-Covid ablation experiment pre-training target

[0137]

[0138]

[0139] Table 7 Fine-tuning results under different pre-training rounds

[0140]

[0141] 5CONCLUSION

[0142] In this study, we noticed the problem of data scarcity in the medical field, we introduced MAE into this field and made corresponding improvements, added a self-distillation module, and constructed a new coronavirus infection medical image classification model based on self-supervised learning (SSL-Covid), aiming to shorten the diagnosis time of COVID-19 patients by using computed tomography (CT) to reduce the duration and number of disease transmission and provide a decision support system for radiologists to diagnose COVID-19. In this experimental study, by adding labels, we introduced global information into self-supervised learning, and further self-distilled, migrating the global information learned by the decoder to the encoder, optimizing and improving the feature extraction ability of the encoder. Through relevant experiments, we proved that the SSL-Covid model has more advantages than MAE in the recognition of COVID-19 infection in CT images, and is easy to implement. And for the problem of large-scale data demand of ViT architecture, SSL-Covid also alleviates it well. In short, these experimental results show that SSL-Covid can further improve the performance of ViT in medical image classification.

[0143] In the medical field, data scarcity is an unavoidable problem, and self-supervised learning aims to extract the representation characteristics of data itself as supervision information by designing proxy tasks for unlabeled data to improve the feature extraction ability of the model, so the application of self-supervised learning in this field has great potential. In future work, we will explore the effectiveness of SSL-Covid in other medical tasks, such as medical image segmentation tasks.

[0144] The above is only the preferred specific implementation of the present application, but the protection scope of the present application is not limited to this, any skilled person in the art can make equivalent replacement or change according to the technical solution and concept of the present application within the technical range disclosed by the present application, which should be covered within the protection scope of the present application.

Claims

1. A method for classifying medical images of COVID-19 infection based on self-supervised learning, characterized by, Based on the mask autoencoder (MAE) self-supervised method, and supplemented by self-distillation learning, the SSL-Covid model is constructed. The SSL-Covid model learns domain knowledge from data through self-supervised learning, and improves the feature extraction capability of the model through self-distillation learning, providing auxiliary information for new crown infection classification. The overall model framework is divided into two stages: pre-training stage and fine-tuning stage, and both stages are carried out on the target data set without introducing external data. The pre-training stage consists of two auxiliary tasks: image reconstruction and self-distillation modeling. Image reconstruction completes the reconstruction task of the masked image, promoting the encoder to better learn the domain knowledge. The image modeling process with mask is to divide the input image into image blocks and randomly mask them. The visible image blocks are input into the encoder to obtain the latent representation. The decoder combines the encoder output and the mask token to restore the masked image blocks. Then the reconstruction loss is calculated through the mean square error. After pre-training, the learned encoder weights are migrated to the fine-tuning stage. The visual self-attention model (ViT) is used to convert the image into a sequence form and add position information. After the Transformer Blocks operation, the classtoken is extracted for subsequent processing. Then the classification of new crown infection is performed. In the pre-training stage, two supervised branches are introduced, which are located after the encoder and the decoder respectively. The branches are Student and Teacher. Compared with the local processing of patches in the encoder part, the decoder processes all patches with more globality. The results after the decoder are mapped to the results after the encoder through the Teacher branch and the Student branch, and the distillation constraint is applied. The global information is added to the encoder to improve the feature extraction capability of the encoder. For both branches, the Student branch at the encoder side and the Teacher branch at the decoder side only use one layer of MLP to match the output vector from the encoder and the output vector from the decoder respectively, through the branches at their respective sides, to obtain two high-dimensional vectors and respectively, and finally match the distributions of these two vectors by minimizing the cross-entropy loss, which is formulated as: (6) (7) (8) wherein denotes a cross-entropy loss function, is the true label, denotes a distillation loss; The total loss calculation formula is as follows: (9)。 2. The self-supervised learning-based COVID-19 infection medical image classification method of claim 1, wherein, A visual self-attention model (ViT) is composed of a patch embedding layer, a position embedding layer, and a Transformer Blocks; Patch embedding layer: ViT is inputted as a sequence, patch converts any high-dimensional data into the form of a sequence. For a multi-modal image where X, , H, W, C are the vector representation of the image, real matrix, image height, image width and image channel number, after a series of image preprocessing, the image is processed into a unified resolution size of 224 224, which is passed through a convolution layer and straightened into where P is the resolution size of each image block, is the number of image blocks, is the reshaped vector representation, and then they are mapped to D-dimensional vectors using a trainable linear projection, i.e. the length ; Position embedding layer: add position embedding information at patch embedding, standard ViT uses 1D learnable position embedding, i.e., learnable position embedding ; Self-attention block Transformer Blocks: A Transformer Block is composed of multi-head attention (MSA), multi-layer perceptron (MLP), and layer normalization (LN) in alternation; Firstly, the input image is Through the Patch Embedding layer, then concatenate it with a class token for classification, the obtained patches plus position embedding are sent into the Transformer Encoder module for operation, and finally the class token in the result is extracted and sent into the MLP module for classification, and the calculation process is as follows: (1) (2) (3) (4)； wherein represents the vector representation of the image after patch embedding layer and positional embedding processing, represents the vector representation after processing by the Transformer Blocks, represents the vector representation after processing by the layer normalization and multi-layer perceptron, represents the final prediction output. 3.The self-supervised learning-based COVID-19 infection medical image classification method of claim 1, wherein, The image modeling process with mask is as follows: the encoder in MAE uses the Transformer Block in ViT, which first divides the input image into non-overlapping image blocks , and then obtains patch embedding through linear mapping, and then randomly masks the patches in proportion to define visible patches and invisible patches as , , and the visible patches are retained Add position information as the input of the encoder for processing, and then obtain the corresponding latent representation for image reconstruction. The decoder in MAE uses a complete token set as input, which is composed of the encoding vector output by the encoder and the mask token. For mask token, it is a shared and learnable vector representing the missing patch to be predicted. For all tokens, the corresponding position embedding is added to recover the patch at each specific masked position through the decoder model. MAE by predicting the pixel values of the masked image block and the original unmasked region pixel values The reconstruction loss computation is done using mean squared error, MSE: (5) wherein represents the mean squared error loss between the real and predicted pixels.

Images