SAM-guided knowledge distillation-based semi-supervised medical image segmentation method and device

Abstract

The invention discloses a semi-supervised medical image segmentation method and device based on SAM guided knowledge distillation, and the method proposes a multi-view collaborative training strategy, employs two different sub-networks to carry out collaborative teaching, and generates a more generalized result. A learnable prompt strategy is provided to dynamically generate prompt information, and an adapter is integrated to finely adjust the SAM and is specially used for a medical image segmentation task; besides, knowledge distillation guided by the SAM is provided, useful knowledge is transferred from the SAM to the two sub-networks, learning can be carried out from prediction of the SAM, and the influence of false labels in the training process is reduced. According to the method, the effect of full supervision can be achieved only through a small amount of annotated data, and the demand quantity of the annotated data is greatly reduced; the generalization ability of a large model is fully utilized, and knowledge distillation of the SAM module is applied to the semi-supervised field for the first time, so that the segmentation performance of the semi-supervised network module is greatly improved.

Description

Technical Field

[0001] The present invention relates to the field of image processing, and in particular to a semi-supervised medical image segmentation method and device based on SAM-guided knowledge distillation. Background Art

[0002] Medical image segmentation plays a pivotal role in various clinical applications. Annotating large-scale pixel-level medical data for fully supervised training is an expensive and time-consuming task. Therefore, this constraint has led to greater interest in semi-supervised learning methods that can leverage large amounts of unlabeled data to enhance model performance.

[0003] Recently, basic models like the Segment Anything Model (SAM) (Kirillov, Alexander, et al. "Segment anything." Proceedings of the IEEE / CVF International Conference on Computer Vision. 2023.) have shown promising progress in segmentation tasks. However, SAM cannot be directly applied to all segmentation scenarios, especially in medical image segmentation tasks. To address this issue, some studies (Chen, Tianrun, et al. "Sam fails to segment anything? –sam-adapter: Adaptingsamin underperformed scenes: Camouflage, shadow, and more." arXiv preprint arXiv:2304.09148 1.2(2023):5.) integrated an adapter into SAM and subsequently fine-tuned SAM for downstream tasks on specific datasets. In addition, some studies (Zhang, Yichi, Yuan Cheng, and Yuan Qi. "Semisam: Exploring sam for enhancing semi-supervised medical image segmentation with extremely limited annotations." arXiv preprint arXiv:2312.06316(2023).) integrated SAM to improve the accuracy of pseudo labels for unlabeled data, thereby improving model performance. These methods provide SAM with precise predicted points or boxes as prompts for generating pseudo labels. However, one disadvantage of this strategy is that providing incorrect prompts will directly lead to a significant drop in the performance of SAM. Although many studies have integrated SAM into medical segmentation tasks, how to best utilize the generalization ability of SAM to enhance semi-supervised medical image segmentation remains an area worthy of further research. Summary of the invention

[0004] The purpose of the present invention is to provide a semi-supervised medical image segmentation method and device based on SAM guided knowledge distillation.

[0005] The technical solution to achieve the purpose of the present invention is: a semi-supervised medical image segmentation method based on SAM-guided knowledge distillation, comprising the following steps:

[0006] Step 1: collect various types of annotated and unannotated medical images, and then perform data augmentation to form training, validation, and test sample sets;

[0007] Step 2: construct a multi-view collaborative training semi-supervised network module from two sub-networks;

[0008] Step 3, construct a learnable prompt information SAM module;

[0009] Step 4: For unlabeled data, the output of the semi-supervised network module will be used as mask prompt information for the SAM module. After that, the results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module. For labeled training samples, not only will the SAM module perform distillation training on the two sub-networks, but also labels will be used to perform supervised training on the SAM module and the semi-supervised network module.

[0010] Step 5: During training, the parameters of the semi-supervised network module and the SAM module are updated through the loss function and the optimizer; during application, the SAM module does not participate in the calculation, and only needs to send the collected medical images to the semi-supervised network module to obtain the segmentation results after calculation.

[0011] A semi-supervised medical image segmentation device based on SAM-guided knowledge distillation, used to implement the above-mentioned semi-supervised medical image segmentation method based on SAM-guided knowledge distillation, the device includes a medical image acquisition module, a semi-supervised network module and a SAM module;

[0012] The medical image acquisition module is used to collect various types of annotated and unannotated medical images, and then perform data enhancement to form training, verification and test sample sets;

[0013] For unlabeled data, the output of the semi-supervised network module will be used as mask prompt information for the SAM module; the results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module; and for labeled training samples, not only will the SAM module perform distillation training on the two sub-networks, but also labels will be provided to the SAM module and the semi-supervised network module for supervised training;

[0014] During training, the parameters of the semi-supervised network module and the SAM module are updated through loss functions and optimizers; during application, the SAM module does not participate in the calculation, and only needs to send the collected medical images to the semi-supervised network module to obtain the segmentation results after calculation.

[0015] A computer device comprises a memory, a processor and a computer program stored in the memory and executable on the processor, wherein the steps of the above method are implemented when the computer program is loaded into the processor.

[0016] A computer-readable storage medium stores a computer program, which implements the steps of the above method when executed by a processor.

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

[0018] (1) The present invention only requires a small amount of labeled data to achieve the effect of full supervision, which greatly reduces the demand for labeled data;

[0019] (2) By making full use of the generalization ability of the large model, the knowledge distillation of the SAM module is applied to the semi-supervised field for the first time, which greatly improves the segmentation performance of the semi-supervised network module. BRIEF DESCRIPTION OF THE DRAWINGS

[0020] Figure 1 This is the architecture diagram of KnowSAM.

[0021] Figure 2 The following is a visual comparison chart of the results of the present invention and other methods. DETAILED DESCRIPTION

[0022] The present invention proposes a semi-supervised medical image segmentation method (KnowSAM) and device guided by SAM distillation based on learnable prompt information. First, we propose a multi-view co-training strategy, which adopts two different sub-networks for collaborative teaching to produce more generalized results. Secondly, we propose a learnable prompt strategy (Learnable Prompt SAM) to dynamically generate prompt information and integrate an adapter to fine-tune SAM specifically for medical image segmentation tasks. In addition, we propose SAM-induced Knowledge Distillation to transfer useful knowledge from SAM to two sub-networks, enabling them to learn from SAM's predictions and mitigate the impact of erroneous pseudo-labels during training. It is worth noting that the predictions generated by our sub-networks are used to generate mask prompts for SAM, which effectively promotes information exchange between modules.

[0023] A semi-supervised medical image segmentation method based on SAM guided knowledge distillation, comprising the following steps:

[0024] Step 1: Collect various types of annotated and unannotated medical images, and then perform data augmentation to form training, validation, and test sample sets.

[0025] Step 2: Construct a multi-view collaboratively trained semi-supervised network module from two sub-networks.

[0026] Step 3: Build a learnable hint SAM module.

[0027] Step 4: For unlabeled data, the output of the semi-supervised network module will be used as mask prompt information to the SAM module. After that, the results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module. For labeled training samples, not only will the SAM module perform distillation training on the two sub-networks, but there will also be labels for supervised training of the SAM module and the semi-supervised network module.

[0028] Step 5: During training, the parameters of the semi-supervised network module and the SAM module are updated through the loss function and the optimizer. However, when applied, the SAM module does not participate in the calculation, which greatly reduces the inference overhead. Simply send the collected medical image to the semi-supervised network module and obtain the segmentation result after calculation.

[0029] Furthermore, in step 1, the image size is first adjusted to 224×224 in width and height. Then, the training data is horizontally flipped, vertically flipped, rotated, and noise is added to the data enhancement method.

[0030] Furthermore, the semi-supervised network module for multi-view collaborative training in step 2 is specifically:

[0031] First, the input image passes through two sub-networks and Then we get the prediction logic diagram, recorded as and Then, the results of the two sub-networks are fed into the aggregation module. There are three perspectives for analysis in the aggregation module: and The first perspective is entropy information, which uses entropy as an uncertainty map and focuses on solving the inherent uncertainty in the two subnets. The entropy acquisition method is defined as follows in represents the prediction result of the cth category of the subnetwork, and C represents the number of categories. and Substituting into the above formula, we can obtain the entropy graph of the two sub-models, recorded as and In order to improve the confidence of the model prediction, an unsupervised entropy loss function is used for the two sub-networks. in and Respectively and The kth pixel of the cth class in , H and W represent the height and width of the logical prediction map respectively. The second perspective information is the area where the two network predictions are inconsistent, denoted as It is calculated by and The XOR of in and Represents the binarization operation and XOR operation respectively. The reason for binarization is that and It is a decimal, and can only be XORed after being binarized. The third perspective information is the information of the sub-network prediction logic graph itself. Finally, we send all the information from the three perspectives to the fusion module to get the aggregated result. in represents the convolution-batch normalization and ReLU activation operations, It means that the splicing operation is performed at the channel level. Here, It contains the prediction results of labeled and unlabeled data, which are recorded as and Aggregate results It will be supervised by labeled data, so the loss function for the aggregation module can be defined as where Y l It represents the label of the annotated data. It is represented by the cross entropy loss function And dice loss function The supervised loss function is composed of and The prediction of labeled data in is also supervised by the label, so the supervised loss function of the entire semi-supervised network module can be written as in and express and In addition, we also use a mutual supervision strategy to ensure the consistency of the results generated by the two sub-networks. This consistency function is expressed as

[0032]

[0033] Furthermore, the step 3 of constructing a learnable prompt information SAM module is specifically as follows:

[0034] First, the input image passes through the encoder of the SAM module to generate the feature information of the image, denoted as Z. Then, the image feature Z passes through a learnable lightweight feature decoder ψ(·) to obtain the prompt information. This learnable lightweight decoder consists of two basic blocks consisting of convolution, batch normalization and ReLU activation layers, and a fully connected layer. The prompt information is a dense point or box, denoted as and Where N p and N b represents the number of points or boxes, and L represents the length of the dense prompt. Here L is set to 128, so the points and boxes here are not in the traditional sense, but a high-dimensional vector information. It can be expressed as P p ,P b =ψ(Z;θ s ), where θ s The parameters of the ψ(·) decoder are represented by . Finally, the image feature information Z and the dense prompt information Ρ p and P b And the aggregation module gets Send it to the decoder of the SAM module, and you can get the prediction result of the SAM module, which is expressed as in It represents the decoder of the SAM module, θ m represents the decoder parameters. In addition, we introduce an adapter tuning strategy in the encoder and decoder of the SAM module to fine-tune the SAM module to adapt the SAM module to this new prompting method; the adapter tuning strategy is to freeze the parameters in the SAM module and only add learnable components between the layers of the SAM module, each component consists of two linear layers; the SAM module also has only the label as the supervision signal, so the loss function is expressed as in express There are prediction results of labeled data in , Already explained in step 1.

[0035] Furthermore, the supervised training of the labeled training samples in step 4 has been described in detail in steps 2 and 3. The results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module, specifically:

[0036] We use distillation learning to extract the generalization ability of the SAM module to the two sub-networks, thereby enhancing the inherent ability of each sub-network and reducing the impact of the sub-network on semi-supervised data noise. According to the distillation learning approach, we need to use temperature-calibrated Softmax to soften the Kullback-Leibler divergence.

[0037] The probability graph in the function can be expressed as in Represents the prediction logic diagram of c, It represents the softened prediction probability map of the cth class, and T represents the temperature coefficient. and and the output of the SAM module Substituting into the above formula, we can get and We use the SAM module as the teacher model, then the distillation loss can be expressed as Where KL represents the Kullback-Leibler divergence function.

[0038] Furthermore, in step 5, the SAM module is only supervised by the labeled data, so the overall loss function of the SAM module is It has been explained in detail in step 3. The overall loss function of the semi-supervised network module is and In step 1, This is described in step 4. e is set to 0.9. m It will change with the number of iterations, specifically where t max is the maximum number of iterations, and β is 1.

[0039] The specific implementation of the present invention is described in detail below with reference to the accompanying drawings and examples.

[0040] Example

[0041] Combined with Figure 1 , this embodiment proposes a semi-supervised medical image segmentation method guided by SAM distillation that can learn prompt information. The specific steps are as follows:

[0042] 1.1 Collect public medical data sets. In this paper, five polyp data sets and one skin data set are mainly used. 30% of the data are selected as labeled data samples X l , the others are unlabeled samples X u .

[0043] 1.2 The image resolution is adjusted to 224×224 by bilinear interpolation, and a series of data enhancements such as flipping, rotating and mapping are performed on the image.

[0044] 2.1 For the semi-supervised segmentation network, first send the input image to the two sub-networks, and you can get the predictions of the two networks and The prediction results for labeled data will be supervised by the labels.

[0045] 2.2 This and By calculating the entropy formula Get the entropy of the two sub-networks, denoted as and At the same time, and After the sigmoid activation function and binarization (the threshold is 0.5), the XOR method is used to find the area where the two sub-networks predict inconsistently, which is recorded as

[0046] 2.3 pairs and and After the splicing operation, Conv-BatchNorm-LeakyReLU (CBL) is performed, and then the inconsistent areas are CBL is also used for feature extraction. The feature information of these three is concatenated and sent to a CBL layer before outputting the fusion combination.

[0047] 2.4 It contains the prediction results of labeled and unlabeled data, which are recorded as and Then the prediction results for the labeled data will be supervised by the label. The hint information sent to the decoder of the SAM module is used as a mask.

[0048] 3.1 For the distillation guidance network of the learnable hint information SAM, the image is first sent to the encoder part of the SAM module to obtain the semantic information, denoted as Z.

[0049] 3.2 Send Z to the lightweight decoder ψ(·) to generate dense prompt information, which is a macroscopic point and box, denoted by P p and P b (High-dimensional vector information consisting of 128 points). This lightweight decoder consists of two convolution-normalization-activation function blocks and a fully connected layer.

[0050] 3.3 Then the semantic information Z is integrated with the prediction results and dense prompt information P p and P b It is sent to the decoder of the SAM module to generate the prediction result of the SAM module, which is recorded as Similarly, the prediction part of the SAM module for labeled data will be supervised by the labels.

[0051] 4.1 Using the knowledge distillation method, the results of the SAM module Will be distilled into two sub-networks and

[0052]

[0053] 4.2 Calculate the entropy loss function of the two sub-networks themselves to reduce the uncertainty of the model.

[0054] 4.3 Use the cross entropy loss function to calculate the consistency loss function between the two sub-networks.

[0055] 4.4 Calculate the gradient and update the parameters of the SAM module and the semi-supervised segmentation network module. Perform a verification on the verification set every 200 iterations, and finally save the weights that perform best on the verification set.

[0056] 5.1 Test the trained model on the test set and compare it with the classic semi-supervised method MT.

[0057] 5.2 Table 1 is a comparison between the present invention and the famous method Mean-Teacher in terms of mean Dice, mean IoU and 95HD;

[0058] Table 1 Comparison of results with MT methods on five polyp datasets and skin dataset

[0059]

[0060]

[0061] It can be seen from Table 1 that the present invention surpasses the traditional semi-supervised MT model in all evaluation indicators.

[0062] 5.3 Figure 2 The first column is the input image, the second column is the corresponding true label segmentation map, the third column is the segmentation map predicted by the present invention, and the fourth column is the segmentation map predicted by Mean-Teacher (MT). It can be seen that the present invention can better and more completely segment polyps, while the MT method is incomplete for small-sized polyps and discontinuous for large-sized polyps, and there are missed detections and false detections, which reflects the superiority of the present invention over other methods.

[0063] The above embodiments only show several specific implementation methods of the present application. Although described in more detail, they should not be regarded as limiting the scope of the present invention. It should be emphasized that, for those of ordinary skill in the art, various modifications and improvements can be made without departing from the basic concept of the present application, which should be included in the protection scope of the present application. Therefore, the protection scope of this patent should be based on the attached claims.

Claims

1. A semi-supervised medical image segmentation method based on SAM-guided knowledge distillation, characterized in that: The following steps are involved: Step 1: collect various types of annotated and unannotated medical images, and then perform data augmentation to form training, validation, and test sample sets; Step 2: construct a multi-view collaborative training semi-supervised network module from two sub-networks; Step 3, construct a learnable prompt information SAM module; Step 4: For unlabeled data, the output of the semi-supervised network module will be used as mask prompt information for the SAM module. After that, the results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module. For labeled training samples, not only will the SAM module perform distillation training on the two sub-networks, but also labels will be used to perform supervised training on the SAM module and the semi-supervised network module. Step 5: During training, the parameters of the semi-supervised network module and the SAM module are updated through the loss function and the optimizer; during application, the SAM module does not participate in the calculation, and only needs to send the collected medical images to the semi-supervised network module to obtain the segmentation results after calculation.

2. The semi-supervised medical image segmentation method based on SAM-guided knowledge distillation according to claim 1, characterized in that: In step 1, the image size is first adjusted to 224×224 in width and height; then the training data are horizontally flipped, vertically flipped, rotated, and noise is added to the data enhancement method.

3. The semi-supervised medical image segmentation method based on SAM guided knowledge distillation according to claim 1, characterized in that: The semi-supervised network module for multi-view collaborative training in step 2 is specifically: First, the input image passes through two sub-networks and Then we get the prediction logic diagram, recorded as and Then, the results of the two sub-networks are fed into the aggregation module; There are three perspectives for analysis in the aggregation module and The first perspective is entropy information, which uses entropy as uncertainty mapping to solve the inherent uncertainty in the two subnets; the entropy acquisition method is defined as follows in Represents the prediction result of the cth category of the subnetwork, C represents the number of categories; and Substituting into the above formula, we can obtain the entropy graph of the two sub-models, which is recorded as and Unsupervised entropy loss function is used for the two sub-networks in and Respectively and For the k-th pixel of the c-th class, H and W represent the height and width of the logical prediction map respectively; the second perspective information is the area where the two network predictions are inconsistent, denoted as It is calculated by and The XOR of in and Represents the binarization operation and XOR operation respectively; the reason for binarization is because and It is a decimal, and can only be XORed after being binarized; the third perspective information is the information of the sub-network prediction logic graph itself; finally, all the information from the three perspectives is sent to the fusion module to obtain the aggregation result in represents convolution-batch normalization and ReLU activation operations, Indicates that splicing operation is performed at the channel level; It contains the prediction results of labeled and unlabeled data, which are recorded as and Aggregate results It will be supervised by labeled data, so the loss function for the aggregation module can be defined as where Y l Indicates the label of the labeled data, Represented by the cross entropy loss function And dice loss function The supervised loss function constructed; and The prediction of labeled data in is also supervised by the label, so the supervised loss function of the entire semi-supervised network module can be written as in and express and The prediction results for the labeled data in ; a mutual supervision strategy is used to ensure the consistency of the results generated by the two sub-networks. This consistency function is expressed as 4. The semi-supervised medical image segmentation method based on SAM guided knowledge distillation according to claim 3, characterized in that: The threshold of the binarization operation is 0.

5.

5. The semi-supervised medical image segmentation method based on SAM guided knowledge distillation according to claim 3, characterized in that: In step 3, a learnable prompt information SAM module is constructed, specifically: First, the input image is passed through the encoder of the SAM module to generate the feature information of the image, denoted as Z; then, the image feature Z is passed through a learnable lightweight feature decoder ψ(·) to obtain the prompt information. This learnable lightweight decoder consists of two basic blocks consisting of convolution, batch normalization and ReLU activation layers, and a fully connected layer; the prompt information is a dense point or box, denoted as and Where N p and N b represents the number of points or boxes, and L represents the length of the dense prompt; It can be expressed as P p ,P b =ψ(Z;θ s ), where θ s Represents the parameters of the ψ(·) decoder; Finally, the image feature information Z, dense prompt information P p and P b And the aggregation module gets It is sent to the decoder of the SAM module to obtain the prediction result of the SAM module, which is expressed as in It represents the decoder of the SAM module, θ m Represents decoder parameters; An adapter tuning strategy is introduced in the encoder and decoder of the SAM module to fine-tune the SAM module. The adapter tuning strategy is to freeze the parameters in the SAM module and only add learnable components between the layers of the SAM module. Each component consists of two linear layers. The SAM module also has only the true label as the supervision signal, so the loss function is expressed as in express There are prediction results for labeled data.

6. The semi-supervised medical image segmentation method based on SAM-guided knowledge distillation according to claim 5, characterized in that: In step 4, the results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module, specifically: Distillation learning is used to extract the generalization ability in the SAM module into two sub-networks; temperature-calibrated Softmax is used to soften the probability map in the Kullback-Leibler divergence function, expressed as in Represents the prediction logic diagram of c, represents the softened prediction probability map of the cth class, T represents the temperature coefficient; the output of the two sub-networks and and the output of the SAM module Substituting into the above formula, we get and Taking the SAM module as the teacher model, the distillation loss can be expressed as Where KL represents the Kullback-Leibler divergence function.

7. The semi-supervised medical image segmentation method based on SAM guided knowledge distillation according to claim 6, characterized in that: In step 5, the SAM module is only supervised by the labeled data, so the overall loss function of the SAM module is The overall loss function of the semi-supervised network module is λ e is set to 0.9, λ m It will change with the number of iterations, specifically where t max is the maximum number of iterations, and β is 1.

8. A semi-supervised medical image segmentation device based on SAM-guided knowledge distillation, characterized in that: Used to implement the method described in any one of claims 1 to 7, the device includes a medical image acquisition module, a semi-supervised network module and a SAM module; The medical image acquisition module is used to collect various types of annotated and unannotated medical images, and then perform data enhancement to form training, verification and test sample sets; For unlabeled data, the output of the semi-supervised network module will be used as mask prompt information for the SAM module; the results of the SAM module will be used to perform distillation training on the two sub-networks in the semi-supervised network module; and for labeled training samples, not only will the SAM module perform distillation training on the two sub-networks, but also labels will be provided to the SAM module and the semi-supervised network module for supervised training; During training, the parameters of the semi-supervised network module and the SAM module are updated through loss functions and optimizers; during application, the SAM module does not participate in the calculation, and only needs to send the collected medical images to the semi-supervised network module to obtain the segmentation results after calculation.

9. A computer device comprising a memory, a processor and a computer program stored in the memory and executable on the processor, characterized in that: When the computer program is loaded into a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the program is executed by a processor, the steps of the method described in any one of claims 1 to 7 are implemented.

Images