A semi-supervised medical image classification method based on class prototype matching and soft pseudo-labels

Abstract

The present invention provides a semi-supervised medical image classification method based on class prototype matching of soft pseudo labels. The class prototype matching module (CPM) is used to predict the soft pseudo labels of unlabeled data. The module adaptively predicts the corresponding soft pseudo labels based on the class prototype similarity of labeled data in the local neighborhood of the unlabeled data. A dynamic and unbiased updated cache queue is designed to store the labels and embeddings of correctly predicted labeled data, and the mean of the embeddings of multiple labeled data belonging to the same category in the cache queue is used as the prototype of the class, so that the class prototype is closer to the true feature center. The mixup enhancement method is used to mix labeled data and unlabeled data and their corresponding labels to enhance the model's ability to learn intra-class and inter-class features. An additional regularization term is added between the two different types of enhancements for unlabeled data to improve the model's predictive ability. The present invention effectively combines pseudo labels and consistency regularization methods.

Description

Technical Field

[0001] The present invention belongs to the technical field of medical image classification, and in particular relates to a semi-supervised medical image classification method based on class prototype matching with soft pseudo labels. Background Art

[0002] Medical image classification is a crucial task in medical imaging diagnosis, aiming to diagnose patient conditions by analyzing medical images. Currently, deep learning-based medical image classification techniques, through qualitative and quantitative analysis of medical images, can help doctors more quickly identify hidden lesions and complete diagnosis and treatment, reducing the rates of missed and misdiagnosed diagnoses, as well as the cost of diagnosis. However, due to the typically small size and high annotation costs of medical image datasets, traditional supervised learning methods face challenges in medical image classification, such as insufficient data and overfitting. Semi-supervised learning techniques can alleviate these issues to a certain extent. Numerous studies have demonstrated that semi-supervised learning methods, which use large amounts of unlabeled data to supplement training with smaller amounts of labeled data, can significantly improve model classification accuracy, approaching or even achieving the performance of training with only large amounts of labeled data. Furthermore, compared to labeled medical data, unlabeled medical data is not only easier to collect from clinical practice but also reduces the cost of data annotation. Therefore, studying medical image classification methods based on semi-supervised learning is a very meaningful task, which can effectively improve the generalization and robustness of deep learning models, help alleviate the problem of scarcity of labeled data sets, and improve the current situation where massive unlabeled clinical medical data are not effectively utilized. It can reduce medical costs and has important research value and broad application prospects.

[0003] In the field of medical imaging, semi-supervised learning methods have been extensively researched and explored. Currently, the mainstream approaches are consistency regularization and pseudo-labeling. Consistency regularization, based on the smoothness assumption, improves the model's classification performance and generalization ability by constraining the output of unlabeled data under different perturbations. This method effectively utilizes the information in unlabeled data and is simple and easy to implement. However, consistency regularization also has some drawbacks. First, while consistency regularization constrains the output of unlabeled data, unlabeled data may contain noise that can affect model performance. Therefore, consistency regularization can be sensitive to noise. Second, this method relies heavily on the perturbation function being adapted to each new medical image, resulting in weak generalization. Pseudo-labeling methods use a model trained on labeled data to predict unlabeled data and add the predictions as pseudo-labels to the training set, increasing the amount of training data and improving model performance. Pseudo-labeling methods are simple and easy to use, requiring no additional effort, and requiring only the addition of unlabeled data to the training set. Therefore, they can handle large amounts of unlabeled data. However, pseudo-labeling methods also have their own drawbacks. First, since the labels of unlabeled data are unknown, if the prediction is wrong, the wrong labels will be added to the training data, resulting in a decrease in model performance. Secondly, the pseudo-labeling method has limitations in its use of unlabeled data. Since such methods can only use the prediction results of unlabeled data, there may be some useful information in the unlabeled data, such as features, relationships between categories, etc., but this information is not utilized. Furthermore, pseudo-labeling methods usually filter out high-quality (confidence level higher than the threshold) unlabeled samples by pre-setting a fixed threshold. Although such methods can effectively improve the quality of pseudo-labels, there is an obvious drawback, which is that a large amount of data below the set threshold is filtered out, which reduces the utilization rate of unlabeled data. Especially in the early training stage, only a few unlabeled samples are involved in training, resulting in poor overall model performance. Summary of the Invention

[0004] To overcome the shortcomings of the prior art, the present invention aims to provide a semi-supervised medical image classification method based on class prototype matching with soft pseudo-labels. Compared to traditional supervised methods, this method only requires manual labeling of a portion of medical images to generate a classifier with excellent classification capabilities. This effectively reduces the time spent on manual image labeling, significantly reduces medical costs, and alleviates the scarcity of labeled medical datasets. It offers advantages such as high classification accuracy and strong generalization.

[0005] To achieve the above object, the present invention provides the following technical solution: a semi-supervised medical image classification method based on class prototype matching soft pseudo labels, comprising the following steps:

[0006] Step 1: Divide the public medical image classification dataset into a training set, a test set, and a validation set. The training set is divided into a labeled dataset and an unlabeled dataset, and all images and corresponding label data are preprocessed.

[0007] Step 2: Build a medical image classification network model F θ (·), including: input layer, deep convolution layer, downsampling layer, output layer;

[0008] Step 3: Build a mark cache queue Q;

[0009] Step 4: Medical image classification network model F constructed in steps 2 and 3 θ (·) and the labeled cache queue Q design loss function, including: supervised loss Unsupervised loss and consistency loss between unlabeled data

[0010] Step 5: Sample the labeled data set and the unlabeled data set in step 1 in a certain batch size;

[0011] Step 6: performing weak enhancement on the marked image in step 5, and performing weak enhancement and strong enhancement on the unmarked image in step 5;

[0012] Step 7: Input the weakly enhanced labeled image obtained in step 6 and the weakly and strongly enhanced unlabeled images into the medical image classification network model respectively, and obtain corresponding embeddings;

[0013] Step 8: The cache queue is dynamically updated to queue the embeddings and labels of the weakly enhanced labeled data of the current batch. When the queue is full, the embeddings and corresponding labels that were queued earliest are deleted from the cache space, and the embeddings and labels of the new labeled data are queued.

[0014] Step 9: Calculate prototypes of different categories based on the mark buffer queue in step 8 and the mean of the embedded labeled samples;

[0015] Step 10: Predict the pseudo labels of the unlabeled images based on the embedding of the weakly enhanced unlabeled images and the embedding of the strongly enhanced unlabeled images in step 7 and the prototypes of each category in step 9, and obtain the weakly enhanced pseudo labels y′ and the strongly enhanced pseudo labels y″;

[0016] Step 11: The weakly enhanced pseudo-label y′ and the strongly enhanced pseudo-label y″ obtained in step 10 are mixed to obtain As the true label of all unlabeled data;

[0017] Step 11: The unlabeled data after strong and weak enhancement with a batch size of B is recorded as χ u The set and batch size are also the weakly augmented labeled dataset χ of B l After being merged using the Concat(·) function and then randomly shuffled using the Shuffle(·) function, the mixed data is recorded as W;

[0018] Step 12: The weakly enhanced labeled data χ l Linearly mix the unlabeled samples in the mixed data W and their corresponding true labels and pseudo labels into the mixup module, and train the network using the mixed data with continuously updated loss;

[0019] Step 13: Calculate supervised loss

[0020] Step 14: Calculate unsupervised loss

[0021] Step 15: Calculate the consistency loss of unlabeled data

[0022] Step 16: Back-propagation updates the parameters of the medical image classification model and saves the best medical image classification model;

[0023] Step 17: Use the images in the test set to test the classification performance of the best medical image classification model in step 16, and finally obtain the average AUC, sensitivity, specificity, accuracy and F1-Score of all test medical images.

[0024] Compared with the prior art, the present invention has the following beneficial effects:

[0025] 1) This paper proposes a semi-supervised deep learning method for soft pseudo-labeling with prototype matching and consistency regularization for medical image classification. The algorithm predicts soft pseudo-labels for unlabeled data using a prototype matching module (CPM). This module adaptively predicts soft pseudo-labels based on the prototype similarity of labeled data within a local neighborhood of the unlabeled data.

[0026] 2) The present invention designs a dynamic and unbiased updated cache queue to store the labels and embeddings of the correctly predicted labeled data. The mean of the embeddings of multiple labeled data belonging to the same class in the cache queue is used as the prototype of the class, making the class prototype closer to the true feature center, thereby providing strong supervision information for training and improving the accuracy of the predicted pseudo-labels.

[0027] 3) This paper uses a mixup augmentation method to mix labeled and unlabeled data and their corresponding labels to enhance the model's ability to learn intra-class and inter-class features. It also adds an additional regularization term between the two different types of augmentation for unlabeled data to improve the model's predictive power. Experimental results on two benchmark datasets, the ISIC2018 dataset and the CheXpert dataset, demonstrate that this paper effectively combines pseudo-labeling and consistency regularization methods, achieving superior classification performance and generalization capabilities on medical image classification tasks compared to other state-of-the-art semi-supervised classification methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0028] Figure 1 It is the overall framework diagram of the present invention;

[0029] Figure 2 This is a schematic diagram of the class prototype matching module;

[0030] Figure 3 Schematic diagram of data mixing;

[0031] Figure 4 This is a further qualitative comparison of the method of the present invention and other methods on the confusion matrix. DETAILED DESCRIPTION

[0032] In order to make the purpose, technical solution and advantages of the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and embodiments. The specific embodiments described herein are only used to explain the technical solution of the present invention and are not intended to limit the present invention.

[0033] A semi-supervised medical image classification method based on class prototype matching and soft pseudo labels. The overall framework of this method is a teacher-student model. The student network and the teacher network are trained simultaneously. The teacher network parameters are generated by exponential moving average (EMA) iterative update using the student network parameters to achieve the goal of dynamically generating reliable consistency. Figure 1 As shown in Figure 2, the overall network structure includes a teacher-student module, a class prototype matching module, a mixup module, and a consistency regularization module. The overall algorithm combines the advantages of multiple advanced algorithms in semi-supervised deep learning. Specifically, Figure 1 、 Figure 2 as well as Figure 3 As shown, it includes the following steps, which should be performed in sequence:

[0034] Step 1: Divide the public medical image classification dataset into training set, test set and validation set, and divide the labeled dataset into the training set and unlabeled datasets All images and their corresponding labeled data were preprocessed. The ISIC2018 dataset, consisting of 10,015 dermoscopic images labeled for seven different disease categories, was randomly partitioned into 70% for training, 10% for validation, and 20% for testing. All images were preprocessed and resized to 299×299 pixels. The CheXpert large-scale chest X-ray dataset was used for multi-label chest disease diagnosis. Preprocessing began by removing all data labeled Uncertain and Absent, as well as lateral views, from the CheXpert dataset and resizing the images to 128×128 pixels. The dataset was partitioned to create small labeled data training sets by balancing each disease category. Five labeled data training sets (1400, 2800, 4200, 5600, and 7000) were constructed, each containing an equal number of data between 100 and 500 (rounded to the nearest hundred) for each disease category. The validation set and test set are 5,000 and 50,000 randomly selected data from the remaining data after the labeled data training set is constructed. The remaining data will be used as the unlabeled data training set.

[0035] Step 2: Use the InceptionV3 network to build a medical image classification network model F θ (·) = g(f(·)), where f is the feature extractor and g is the classifier. The feature extractor f consists of an input layer, deep convolutional layers, a downsampling layer, and an output layer. The input layer consists of convolutional layers and max pooling layers to extract low-level image features. The deep convolutional layers consist of multiple Inception modules to extract higher-level features. The downsampling layers consist of max pooling layers and average pooling layers to reduce the size of the feature map.

[0036] Step 3: Build a tag cache queue in The storage size is m×k, where m and k are the number of categories and the maximum storage capacity of each category, respectively.

[0037] Step 4: Based on the network model F constructed in steps 2 and 3 θ (·) and the labeled cache queue Q design loss function, including: supervised loss Unsupervised loss and consistency loss between unlabeled data The overall objective function L is expressed as follows:

[0038]

[0039]

[0040] Among them, t is the number of iterations of the current training, T is the total number of training iterations, and β is the balance consistency loss hyperparameters.

[0041] Step 5: Sample the labeled dataset and unlabeled dataset in step 1 in a certain batch size to obtain labeled images. and unlabeled images

[0042] Step 6: weakly enhance the marked image in step 5 to obtain x l′ =ψ(x l ), perform weak enhancement and strong enhancement on the unlabeled image in step 5 to obtain x u′ =ψ(x u ) and x u″ =ψ′(x u ), where weak enhancement is achieved by randomly rotating the image in the range of (-10°, 10°) and randomly translating the image horizontally and vertically in the range of (0, 10%). Strong enhancement uses the RandAugment method, which consists of a series of randomly selected transformations, such as image contrast maximization, image histogram equalization, and color channel reduction.

[0043] Step 7: The weakly enhanced labeled image x obtained in step 6 is l′ , unlabeled image x after weak enhancement and strong enhancement u′ and x u″ Input medical image classification network model F θ , and obtain the embedding of the weakly enhanced labeled image Weakly enhanced embedding of unlabeled images Embedding of unlabeled images with strong augmentation

[0044] Step 8: The tag cache queue Q dynamically stores the embeddings of the k weakly enhanced tag images corresponding to the m classes. and label y i ,in is the exponential moving average (EMA) form of the feature extractor f(·). The cache queue is dynamically updated by queueing the embeddings and labels of the weakly enhanced labeled data of the current batch. When the queue is full, the earliest queued embeddings and corresponding labels are deleted from the cache space, and the embeddings and labels of the new labeled data are queued. It is expressed as:

[0045]

[0046] in, is the embedding of the labeled data, ν[m] represents the index of the largest element in the vector m, → represents the queue operation, and f(x) is the probability distribution of the labeled data output.

[0047] Step 9: Calculate the mean embedding value of k labeled samples belonging to the same category according to the label cache queue Q in step 8. And as a prototype of this category, it is calculated as follows:

[0048]

[0049] Step 10: Embedding the weakly enhanced unlabeled image according to step 7 Embedding of unlabeled images with strong augmentation And the prototypes of each category in step nine Predict pseudo labels for unlabeled images. The prediction method is as follows:

[0050]

[0051]

[0052] Among them, m represents the total number of categories, and ω(·,·) is used to calculate the embedding of unlabeled samples. Prototypes with labeled samples for each class The distance between the two classes is calculated, and a similarity weight vector is ultimately returned to measure the similarity between the unlabeled sample and each class. A larger weight indicates a higher similarity to the class. The soft pseudo label of each unlabeled sample is generated by linearly weighting the similarity weight vector and the class label.

[0053] The ω function is defined as a softmax function with a temperature coefficient τ, where sim(·,·) actually calculates the cosine similarity of the prototypes of each class of the unlabeled sample, which is expressed as follows:

[0054]

[0055] Step 11: The weakly enhanced pseudo-label y′ and the strongly enhanced pseudo-label y″ obtained in step 10 are mixed to obtain As the “true label” of all unlabeled data (including strong and weak enhancement data), the mixing method is as follows:

[0056] λ~Beta(α,α)

[0057] λ′=max(λ,1-λ)

[0058]

[0059] Among them, the mixing coefficient λ follows the Beta distribution and α is a positive hyperparameter.

[0060] Step 11: The unlabeled data after strong and weak enhancement with a batch size of B is recorded as χ u The set and batch size are also the weakly enhanced labeled dataset χ of B l After merging using the Concat(·) function and then randomly shuffling using the Shuffle(·) function, the final set of label data pairs is recorded as W. The operation is as follows:

[0061]

[0062]

[0063] W=Shuffle(Concat( u ,χ l ))

[0064] Step 12: The weakly enhanced labeled data χ l With the previous |χ in the mixed data W l | The data and labels are mixed up separately, and the mixed data is added as the labeled data to χ′ l Similarly, the unlabeled data is compared with the label χ u With the post|χ in W u | data are mixed up, and the mixed up data are added as unlabeled data and labels to χ′ u The specific method is as follows:

[0065]

[0066]

[0067] Step 13: Calculate supervised loss The specific expression is as follows:

[0068]

[0069] in, is the mixed data of the labeled data, p is the mixed label, Represents the mixed data output by the student model (θ parameterization) The probability distribution of represents the cross entropy between two probability distributions p and q.

[0070] Step 14: Calculate unsupervised loss The specific expression is as follows:

[0071]

[0072] in, is the mixed data of unlabeled data, and q is the mixed label. Unsupervised loss It can be expressed as the mixed data output by the student model (θ parameterization) The probability distribution of q and the mean squared error of the mixed label q.

[0073] Step 15: Calculate the consistency loss of unlabeled data The specific expression is as follows:

[0074]

[0075] Among them, x u represents the unlabeled data for each training, B represents the number of unlabeled data in the batch, ψ(·) is denoted as the weak enhancement function and ψ′(·) is denoted as the strong enhancement function, and f(·) is the feature extractor of the teacher model (parameterized by θ′).

[0076] Step 16: Back-propagation updates the medical image classification model parameters and saves the best medical image classification model.

[0077] Step 17: Use the images in the test set to test the classification performance of the best medical image classification model in step 16, and finally obtain the average AUC, sensitivity (Sensitivity), specificity (Specificity), accuracy (Accuracy) and F1 value (F1-Score) of all test medical images.

[0078] Table 1 lists the experimental results of the method of the present invention and the latest semi-supervised learning methods on the ISIC2018 skin lesion dataset, including self-training method (Self-training), semi-supervised deep convolutional adversarial network (SS-DCGAN), temporal integration method (TE), GLM method, NM method, Mean Teacher method, and SRC-MT method. Among them, UpperBound refers to the supervised model trained with 100% (7000) labeled data, which is used as the upper limit of the network model performance. Baseline is the baseline of the experiment, which is the supervised model trained with only 20% (1400) labeled data, which is used as the baseline of the network model performance. Except for UpperBound, the training sets of all comparison methods in this experiment are divided into 20% labeled data and 80% unlabeled data. As shown in Table 1, the evaluation indicators of the semi-supervised medical image classification method proposed in this invention on the skin lesion dataset are all the best compared with the existing advanced methods. Therefore, compared with these latest semi-supervised medical image classification methods, the method of the present invention has certain advantages.

[0079] Table 1 Comparison with the latest semi-supervised methods on the ISIC 2018 dataset (%)

[0080]

[0081] Further qualitative comparison of the method of the present invention with other methods on the confusion matrix is ​​as follows Figure 4 As shown. 20% of the training set data is randomly divided as the labeled training set and applied to the confusion matrix. The number of 7 different types of skin lesion images in the labeled training set is: 142, 989, 81, 45, 146, 16, and 23. The confusion matrix can intuitively display the accuracy of the classification model corresponding to each category, where the diagonal elements present the probability of correct prediction for each category. As can be seen from the figure, compared with other methods, the prediction accuracy of all categories of the method of the present invention has achieved a relatively balanced improvement. It can be seen that the method of the present invention can generate more accurate pseudo labels for unlabeled data through the dynamic unbiased update queue in the class prototype matching and the consistency regularization method, which effectively improves the performance of the semi-supervised classification model.

[0082] To evaluate the effectiveness of our method in multi-label classification, we conducted comparative experiments on the Chexpert dataset, a public dataset for multi-label chest disease diagnosis, with state-of-the-art semi-supervised medical image classification methods, including the LSSE method, the NM method, and the GLM method. The dataset was partitioned using the GLM method, with 500 labeled data points for each disease. Table 2 compares the AUC values ​​of these methods for 14 chest disease image categories. The data in this table show that our method achieved an average AUC of 69.08% on the Chexpert dataset, significantly exceeding the performance of other methods. Compared to the LSSE method, the NM method, and the GLM method, the average AUC values ​​improved by 2.35%, 1.36%, and 1.03%, respectively, and by 4.22% compared to the baseline model. Furthermore, our method outperformed other methods in 10 of the 14 chest disease image categories, further demonstrating the excellent performance of our method in multi-label classification.

[0083] Table 2 Comparison with the latest semi-supervised methods on the ISIC 2018 dataset (%)

[0084]

[0085] The above description merely represents the preferred embodiments of the present invention, and while the description is relatively specific and detailed, it should not be construed as limiting the scope of the present invention. It should be noted that a person skilled in the art would be able to make various modifications, improvements, and substitutions without departing from the spirit of the present invention, all of which fall within the scope of protection of the present invention. Therefore, the scope of protection of the present invention shall be determined by the appended claims.

Claims

1. A semi-supervised medical image classification method based on class prototype matching and soft pseudo-labels, characterized by: The following steps are involved: Step 1: Divide the public medical image classification dataset into a training set, a test set, and a validation set. The training set is divided into a labeled dataset and an unlabeled dataset, and all images and corresponding label data are preprocessed. Step 2: Build a medical image classification network model , including: input layer, deep convolution layer, downsampling layer, output layer; Step 3: Build a tag cache queue ; Step 4: Medical image classification network model built in steps 2 and 3 and tag cache queues Design loss functions, including supervised loss , unsupervised loss and consistency loss between unlabeled data ; Step 5: Sample the labeled data set and the unlabeled data set in step 1 in a certain batch size; Step 6: performing weak enhancement on the marked image in step 5, and performing weak enhancement and strong enhancement on the unmarked image in step 5; Step 7: Input the weakly enhanced labeled image obtained in step 6 and the weakly and strongly enhanced unlabeled images into the medical image classification network model respectively, and obtain corresponding embeddings; Step 8: The cache queue is dynamically updated to queue the embeddings and labels of the weakly enhanced labeled data of the current batch. When the queue is full, the embeddings and corresponding labels that were queued earliest are deleted from the cache space, and the embeddings and labels of the new labeled data are queued. Step 9: Calculate the number of cached items belonging to the same category according to the mark cache queue in step 8. The mean of the labeled sample embeddings , and as the prototype of this category, is calculated as follows: ； Step 10: According to the embedding of the weakly enhanced unlabeled image and the embedding of the strongly enhanced unlabeled image in step 7 and the prototypes of each category in step 9, the pseudo label of the unlabeled image is predicted to obtain the weakly enhanced pseudo label. With strong enhanced pseudo labels ; Step 11: weakly enhanced pseudo labels obtained according to step 10 With strong enhanced pseudo labels , obtained by mixing the two As the true label of all unlabeled data; Step 11: Set the batch size to The unlabeled data after strong and weak enhancement are recorded as The set and batch sizes are also Weakly enhanced labeled dataset use After the functions are merged, use The mixed data after random disruption of the function is recorded as ; Step 12: Weakly enhanced labeled data With mixed data The unlabeled samples in the , and their corresponding true labels and pseudo labels are input into the mixup module for linear mixing, and the network is trained using the mixed data with continuously updated loss; Step 13: Calculate supervised loss ; Step 14: Calculate unsupervised loss ; Step 15: Calculate the consistency loss of unlabeled data ; Step 16: Back-propagation updates the parameters of the medical image classification model and saves the best medical image classification model; Step 17: Use the images in the test set to test the classification performance of the best medical image classification model in step 16, and finally obtain the average AUC, sensitivity, specificity, accuracy and F1-Score of all test medical images.

Images