Semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning

Abstract

The application relates to a kind of semi-supervised medical image segmentation methods based on mutual correction and pixel-level contrast learning, including constructing the cross pseudo-supervised segmentation framework consisting of two differentiated subnetworks, and calculating supervision loss and cross pseudo-supervised loss for labeled data and unlabeled data for training;For the inconsistent area of the prediction output by the two subnetworks, a correction loss module is constructed to correct potential errors;A pixel-level contrast learning module is constructed to optimize the feature representation of low-uncertainty region pixels.Compared with the prior art, the correction module of the application effectively alleviates the confirmation bias problem existing in semi-supervised medical image segmentation, improves the quality of pseudo-labels, and the contrast learning module models the relationship between pixels while controlling the consumption of computing resources and optimizing the feature representation, ultimately significantly improving the segmentation accuracy of medical images.

Description

Technical Field

[0001] This invention relates to the field of medical image processing technology, and in particular to a semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning. Background Technology

[0002] Accurate segmentation of medical images is a key research area in computer-aided diagnosis. In recent years, the emergence of deep learning models has greatly improved the performance of medical image segmentation. However, these high-performance deep learning models rely heavily on large amounts of high-quality labeled datasets. At the same time, accurate pixel-level annotation of medical images is time-consuming, labor-intensive, and requires a great deal of specialized knowledge.

[0003] Existing semi-supervised methods for medical image segmentation can be broadly categorized into pseudo-label-based self-training, co-training, and consistency regularization methods. In implementing these methods, addressing the inherent confirmation bias of semi-supervised learning is crucial. Some methods employ thresholding or uncertainty estimation to exclude supervision signals from unreliable pseudo-labels, while others introduce additional correction losses to address inconsistencies in labeled data predictions and correct potential errors. However, these methods either fail to consider multi-class scenarios or neglect how to handle potential errors in unlabeled data predictions.

[0004] Furthermore, in pixel-level contrastive learning, noisy pseudo-labels can cause the pixel-level contrastive learning based on pseudo-labels to optimize feature representations in the wrong direction, i.e., the noise sampling problem. Although some recent methods use uncertainty estimation to exclude pixels in high uncertainty regions, which alleviates the noise sampling problem to some extent, this may still result in some pixel features not being optimized, limiting the benefits that the contrastive learning module brings to the overall model segmentation performance.

[0005] In summary, there is currently a lack of a semi-supervised medical image segmentation method based on mutual calibration and pixel-level contrastive learning to solve or partially solve the aforementioned problems. Summary of the Invention

[0006] The purpose of this invention is to overcome the shortcomings of the existing technology by providing a semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning, so as to alleviate the inherent confirmation bias problem of semi-supervised methods and the insufficient optimization of pixel-level contrastive learning based on pseudo-labels.

[0007] The objective of this invention can be achieved through the following technical solutions:

[0008] This invention provides a semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning. It utilizes a pre-trained semi-supervised medical image segmentation model comprising two sub-networks to segment a given medical image. The training process of the semi-supervised medical image segmentation model includes the following steps:

[0009] Labeled and unlabeled data are used as inputs to the semi-supervised medical image segmentation model to obtain predictions and pseudo-labels, and the supervision loss and cross-pseudo-supervision loss are calculated.

[0010] For the predictions and pseudo-labels corresponding to labeled data and unlabeled data respectively, calculate the prediction inconsistency region of category or pseudo-label, and calculate the correction loss within the prediction inconsistency region.

[0011] Based on the predicted mean of the two sub-networks, low-uncertainty pseudo-labels are screened through uncertainty estimation, and the comprehensive contrast loss between anchor points and prototypes is calculated through anchor sampling and prototype generation.

[0012] The semi-supervised medical image segmentation model is trained based on the aforementioned supervised loss, cross-spurious supervision loss, correction loss, and comprehensive contrast loss.

[0013] As a preferred technical solution, the process of calculating the correction loss in the region of inconsistency in the prediction for the labeled data includes the following steps:

[0014] Select the mask of the foreground category corresponding to the two pseudo-labels predicted by the model, and obtain the region of inconsistency in the prediction of the foreground category through XOR operation;

[0015] In the region of inconsistency in prediction, the mean square error between the predicted and the mask corresponding to the true label is calculated as the correction loss.

[0016] As a preferred technical solution, for the labeled data, the correction loss in the prediction inconsistency region is calculated. Calculate using the following formula:

[0017]

[0018] in, These represent the two predicted pseudo-labels. This represents the XOR operation. This represents the prediction of the i-th subnetwork with prediction category c, where i = 1, 2 indicates the first or second subnetwork, and MSE represents the mean squared error loss. This represents the true mask corresponding to the region where the prediction for category c is inconsistent.

[0019] As a preferred technical solution, the process of calculating the correction loss in the region of prediction inconsistency for the unlabeled data includes the following steps:

[0020] Identify the inconsistency region of the pseudo-label, and calculate the combined loss of the positive and negative cross-entropy between the current sub-network pseudo-label and the prediction of another sub-network within the inconsistency region as the correction loss.

[0021] As a preferred technical solution, for the unlabeled data, the correction loss in the prediction inconsistency region is calculated. Calculate using the following formula:

[0022]

[0023] in, The label represents the current sub-network's prediction for the i-th pixel. w represents the pseudo-label obtained from another subnetwork. i represents the confidence level of the current subnetwork in predicting the pseudo-label of the i-th pixel, and m represents the number of pixels in the region of inconsistent prediction.

[0024] As a preferred technical solution, the process of screening for pseudo-labels with low uncertainty based on the prediction mean of two sub-networks through uncertainty estimation includes the following steps:

[0025] Based on the mean of the predictions from the two sub-networks, new pseudo-labels are obtained using the Argmax operation. The entropy of the mean prediction is then calculated as an uncertainty estimate. A dynamically growing threshold is used to distinguish between high uncertainty regions and uncertain regions, and pseudo-labels with high uncertainty are excluded.

[0026] As a preferred technical solution, the anchor point sampling and prototype generation process includes the following steps:

[0027] Based on the filtered low-uncertainty pseudo-labels, each pixel feature corresponds to a pseudo-label by downsampling and matching the size of the sample feature map. For the category appearing in the pseudo-label map, a pixel feature is sampled from the corresponding feature map for each category as an anchor point.

[0028] Based on the filtered low-uncertainty pseudo-labels and sample feature maps, a prototype is calculated for each category appearing in the pseudo-label map through a mask mean pooling operation, and the calculated prototypes are stored in a first-in-first-out memory queue.

[0029] As a preferred technical solution, the calculation process for the comprehensive comparison loss between the anchor point and the prototype includes the following steps:

[0030] Calculate the first contrastive loss for each anchor point against all prototypes in the memory queue;

[0031] For each pixel in the low uncertainty region, each pixel is sampled once and added to the anchor point. A confidence-weighted global prototype is calculated on the sample set of the current training batch. A second contrast loss between each anchor point and the global prototype is calculated.

[0032] The comprehensive contrast loss is obtained based on the first contrast loss and the second contrast loss.

[0033] As a preferred technical solution, the first contrast loss Calculate using the following formula:

[0034]

[0035] Where A1 represents the number of sampling anchor points, P i and N i Let i represent the set of positive samples and the set of negative samples corresponding to anchor point i, respectively. V represents positive and negative samples respectively. i The anchor point for sampling is represented by τ = 0.1, which represents the temperature coefficient used to control the penalty for difficult negative samples, and cos(,) represents the cosine similarity operator.

[0036] As a preferred technical solution, the second contrast loss Calculate using the following formula:

[0037]

[0038] For unlabeled data, r c,i R represents the confidence level of the i-th pixel belonging to category c. For labeled data, r c,i =1, v i Ω represents the projected feature of the i-th pixel. c A1 represents the set of all pixels in the low-uncertainty region whose predicted pseudo-label is c, and A2 represents the number of pixels in the low-uncertainty region. v represents the set of C prototypes calculated in the current training batch. i Indicates the anchor point for sampling. P represents ′ The corresponding pseudo-tag in the middle is The prototype.

[0039] Compared with the prior art, the present invention has at least one of the following beneficial effects:

[0040] (1) Alleviating the inherent confirmation bias problem in semi-supervised data: This invention addresses labeled and unlabeled data respectively by calculating two correction losses to correct potential errors in inconsistent prediction regions, effectively alleviating the confirmation bias problem in semi-supervised medical image segmentation, thereby improving the quality of pseudo-labels.

[0041] (2) Improve the problem of insufficient feature optimization in semi-supervised learning: This invention obtains a more complete sample distribution and optimizes all low-uncertainty pixel features through two pixel-prototype comparison learning, effectively models the relationship between pixels, promotes the generation of features with class discriminative ability, and enables the decision boundary to pass through low-density regions, ultimately improving the overall performance of the segmentation model.

[0042] (3) Improve the generalization ability of the model: This invention uses UNet decoder and SwinUNet decoder to build a cross-pseudo-supervision framework. The former emphasizes local modeling, while the latter captures long-range dependencies. By observing the same sample independently from different angles through different but mutually beneficial inductive biases, differentiated predictions are naturally generated. At the same time, the consistency constraints brought by cross-pseudo-supervision promote the robustness and generalization performance of the model.

[0043] (4) Optimize the training process: This invention uses a preheating equation to control the weights of supervised loss and cross pseudo-supervised loss. The weight of cross pseudo-supervised loss increases continuously with training iterations, preventing the model from over-relying on erroneous pseudo-labels in the early stages of training, thus preventing misfitting risks. It also promotes the model to smoothly adapt to different loss terms, improves the model's generalization ability, and enables it to learn more robust feature representations. Attached Figure Description

[0044] Figure 1 This is a flowchart illustrating the semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning in the embodiments.

[0045] Figure 2 This is a schematic diagram of a semi-supervised medical image segmentation model in an example. Detailed Implementation

[0046] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort should fall within the scope of protection of the present invention.

[0047] Example 1

[0048] To address the problems existing in the aforementioned prior art, this embodiment provides a semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning. Starting from two aspects, it designs an improved correction module and a contrastive learning module to alleviate the confirmation bias problem inherent in semi-supervised learning and improve the current insufficient optimization of pixel-level contrastive learning.

[0049] See Figure 1 This method includes the following steps:

[0050] S1. Construct a segmentation model consisting of a ResNet50 encoder initialized with ImageNet pre-trained weights, a UNet decoder based on convolutional operations and a SwinUnet decoder based on attention mechanism, with the two decoders corresponding to two sub-networks respectively. Feed the training data containing labeled and unlabeled data into the segmentation model to obtain predictions and corresponding pseudo-labels. Calculate the supervision loss and cross-pseudo-supervision loss for training.

[0051] Specifically, step S1 includes the following sub-steps:

[0052] S11. Construct a ResNet50 encoder initialized with ImageNet pre-trained weights, a UNet decoder based on convolution operations and a SwinUnet decoder based on attention mechanism, which are randomly initialized. The two decoders and their corresponding self-networks form a semi-supervised medical image segmentation model.

[0053] S12. Normalize all data to [0,1], divide the dataset into training, validation, and test sets, and perform data augmentation on the training data, including random flipping and rotation. Randomly sample a small batch of 2D medical images from the training set, with a batch size of 16, containing 8 labeled and 8 unlabeled images, and feed them into the segmentation model to obtain predictions. and corresponding pseudo tags

[0054] S13. For labeled data, calculate the cross-entropy between the predictions and the actual labels of the two sub-networks. and Dice loss As a supervisory loss The definition is as follows:

[0055]

[0056] in, Y represents the network's output prediction for labeled data, where Y represents the corresponding true label.

[0057] S14. For unlabeled data, calculate the cross-entropy and Dice loss between the pseudo-label of the current sub-network and the prediction of another sub-network, swap their positions, calculate the loss again, and sum them as the cross-pseudo-supervision loss. The calculation is as follows:

[0058]

[0059] in, This represents the output prediction of the i-th subnetwork for unlabeled data. This represents the pseudo-label output by the i-th subnetwork for unlabeled data. It is a combination of cross-entropy and Dice loss.

[0060] S2. Construct a correction module for regions with inconsistent predictions. Regions with different pseudo-labels are considered as regions with inconsistent predictions. For labeled data and unlabeled data, a combination loss of supervision loss from the real label and positive cross-entropy and negative cross-entropy from the pseudo-label are used to correct potential errors.

[0061] Specifically, step S2 includes the following sub-steps:

[0062] S21. For labeled data, extract the masks corresponding to each foreground category from the two pseudo-labels predicted by the model, perform an XOR operation to obtain the prediction inconsistency region for that category, calculate the mean squared error between the network prediction and the mask corresponding to the true label in the prediction inconsistency region as the correction loss to correct potential errors, and finally sum the correction losses for each category, as shown in the following formula:

[0063]

[0064] Where MSE represents the mean squared error loss. This represents the true mask corresponding to the region where the predictions for category c are inconsistent. These represent the two predicted pseudo-labels. This represents the XOR operation. This represents the prediction of the i-th subnetwork with prediction category c.

[0065] S22. For unlabeled data, directly identify the inconsistency regions of pseudo-labels. In the inconsistency regions, calculate the combined loss of the positive and negative cross-entropy between the current sub-network's pseudo-label and the prediction of another sub-network as the correction loss. The calculation formula is as follows:

[0066]

[0067] in, The label represents the current sub-network's prediction for the i-th pixel. w represents the pseudo-label obtained from another subnetwork. i This represents the confidence level of the current sub-network in predicting the pseudo-label for the i-th pixel, where m represents the number of pixels in the inconsistent prediction region. (This is used in the calculation...) Cut off first The minimum value is 1e-4.

[0068] S3. Construct a pixel-level contrastive learning module based on uncertainty. Calculate the entropy of the mean of the predictions of the two sub-networks as uncertainty to exclude pseudo-labels in high-uncertainty regions. Only perform anchor sampling and prototype generation in low-uncertainty regions. Calculate the contrastive loss for each anchor and all prototypes to model the correlation between pixels and optimize feature representation.

[0069] Specifically, step S3 includes the following steps:

[0070] S31. Calculate the mean of the two predictions, use the Argmax operation to obtain a new single pseudo-label different from the one in S1, and then calculate the entropy of the mean prediction as an uncertainty estimate. Use a dynamically growing threshold to distinguish between high uncertainty regions and high uncertainty regions, and eliminate pseudo-labels with high uncertainty. Pseudo-labels after uncertainty filtering. The calculation is as follows:

[0071]

[0072] Among them, ∈=1e-16 prevents the occurrence of log(0), and for labeled data, the real label is used directly.

[0073] The single pseudo-label in this step is obtained by taking the Argmax of the mean of the two predictions. The pseudo-label in the cross pseudo-supervision loss or correction module in S1 is obtained by taking the Argmax of a single prediction. There is one pseudo-label in each of the two sub-networks.

[0074] S32. Based on the filtered pseudo-labels, sample an anchor point for each category appearing in each image and add it to the anchor point set. Simultaneously, calculate a prototype for each category appearing in each image using masked mean pooling and store it in a first-in-first-out memory queue. Calculate the comparison loss between each anchor point and all prototypes in the memory queue, and finally sum them up. The calculation formula is as follows:

[0075]

[0076]

[0077] Where A1 represents the total number of all sampled anchor points, P i and N i These represent the sets of positive and negative samples corresponding to the anchor points, respectively. V represents positive and negative samples respectively. i The anchor point for sampling is represented by τ = 0.1, which represents the temperature coefficient used to control the penalty for difficult negative samples, and cos(,) represents the cosine similarity operator.

[0078] In step S32, only one pixel feature is sampled as an anchor point for each category. In step S33, all pixel features belonging to each category are sampled as anchor points. Therefore, step S33 has more anchor points and includes the anchor points from step S32.

[0079] S33. Sample all pixels in the low-uncertainty region once as anchor points. Calculate a confidence-weighted global prototype on the sample set of the current training batch to obtain global context information. Calculate the comparative loss between each anchor point and the global prototype, and finally sum them. The calculation formula is as follows:

[0080]

[0081]

[0082] Where, r c,i This represents the confidence level of the i-th pixel belonging to category c. For labeled data, r is set as follows: c,i =1, v i Ω represents the projected feature of the i-th pixel. c Let A1 represent the set of all pixels in the low-uncertainty region with predicted pseudo-label c, and A2 represent the total number of pixels in the low-uncertainty region. This represents the C prototypes calculated in the current batch, v i Indicates the anchor point for sampling. P represents ′ The corresponding pseudo-tag in the middle is The prototype.

[0083] Finally, summing all losses, the calculation is as follows:

[0084]

[0085] Where λ1 is the preheating equation

[0086] S4. Build a cross-supervised segmentation framework based on all the above modules, train it, and then predict the segmentation result on any test data.

[0087] To verify the effectiveness of this method, the gradient backpropagation method is used to update the overall segmentation model. When the training reaches the maximum number of iterations, the trained model is used to make predictions on any input image.

[0088] Table 1 shows the segmentation performance of this method compared with existing algorithms when the labeled proportions in the ACDC heart four-class classification dataset are 5% and 10% respectively (↑ indicates that the larger the value of the indicator, the better, ↓ indicates that the smaller the value of the indicator, the better).

[0089] Table 1 Comparison of segmentation performance between this method and existing algorithms

[0090]

[0091] As can be seen from the comparison of various indicators in Table 1, with less labeled data, compared with existing technologies, this method corrects some errors in pseudo-labels by correcting loss, forces the model to focus on challenging learning of blurred organ boundary regions, and improves the quality of pseudo-labels. This enables pseudo-label-based contrastive learning to optimize more accurately classified pixel features, model pixel relationships and generate features with class discriminative power, ultimately significantly improving the accuracy of medical image segmentation.

[0092] Compared with existing methods, this method has the following advantages:

[0093] (1) This method uses two correction losses to process labeled and unlabeled data respectively, corrects potential errors in the prediction of inconsistent regions, effectively alleviates the confirmation bias problem in semi-supervised medical image segmentation, and thus improves the quality of pseudo-labels.

[0094] (2) This method obtains a more complete sample distribution and optimizes all low-uncertainty pixel features through two pixel-prototype contrast learning modules, effectively models the correlation between pixels, promotes the generation of features with class discriminative power, and enables the decision boundary to pass through low-density regions, thereby improving the overall performance of the segmentation model.

[0095] Example 2

[0096] This embodiment provides an electronic device, including: one or more processors and a memory, wherein the memory stores one or more programs, the one or more programs including instructions for executing the semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning as described in Embodiment 1.

[0097] Example 3

[0098] This embodiment provides a computer-readable storage medium, characterized in that it includes one or more programs executable by one or more processors of an electronic device, the one or more programs including instructions for performing a semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning as described in Embodiment 1.

[0099] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any person skilled in the art can easily conceive of various equivalent modifications or substitutions within the technical scope disclosed in the present invention, and these modifications or substitutions should all be covered within the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be determined by the scope of the claims.

Claims

1. A semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning, characterized in that, A pre-trained semi-supervised medical image segmentation model, comprising two sub-networks, is used to segment a given medical image. The training process of the semi-supervised medical image segmentation model includes the following steps: Labeled and unlabeled data are used as inputs to the semi-supervised medical image segmentation model to obtain predictions and pseudo-labels, and the supervision loss and cross-pseudo-supervision loss are calculated. For the predictions and pseudo-labels corresponding to labeled data and unlabeled data respectively, calculate the prediction inconsistency region of category or pseudo-label, and calculate the correction loss within the prediction inconsistency region. Based on the predicted mean of the two sub-networks, low-uncertainty pseudo-labels are screened through uncertainty estimation, and the comprehensive contrast loss between anchor points and prototypes is calculated through anchor sampling and prototype generation. The semi-supervised medical image segmentation model is trained based on the aforementioned supervision loss, cross-pseudo-supervision loss, correction loss, and comprehensive contrast loss. The process of filtering pseudo-labels with low uncertainty based on the prediction mean of two sub-networks through uncertainty estimation includes the following steps: Based on the mean of the predictions from the two sub-networks, new pseudo-labels are obtained using the Argmax operation. Then, the entropy of the mean prediction is calculated as an uncertainty estimate. A dynamically growing threshold is used to distinguish between high uncertainty regions and low uncertainty regions, and pseudo-labels with high uncertainty are excluded. The anchor point sampling and prototype generation process includes the following steps: Based on the filtered low-uncertainty pseudo-labels, each pixel feature corresponds to a pseudo-label by downsampling and matching the size of the sample feature map. For the category appearing in the pseudo-label map, a pixel feature is sampled from the corresponding feature map for each category as an anchor point. Based on the filtered low-uncertainty pseudo-labels and sample feature maps, a prototype is calculated for each category appearing in the pseudo-label map through mask mean pooling operation, and the calculated prototypes are stored in a first-in-first-out memory queue. The calculation process for the overall contrast loss between the anchor point and the prototype includes the following steps: Calculate the first contrastive loss for each anchor point against all prototypes in the memory queue; For each pixel in the low uncertainty region, it is sampled once and added to the anchor point. A confidence-weighted global prototype is calculated on the sample set of the current training batch. The second contrast loss between each anchor point and the global prototype is calculated. The comprehensive contrast loss is obtained based on the first contrast loss and the second contrast loss; The process of calculating the correction loss in the region of inconsistency within the labeled data includes the following steps: Select the mask of the foreground category corresponding to the two pseudo-labels predicted by the model, and obtain the region of inconsistency in the prediction of the foreground category through XOR operation; In the region of inconsistency in prediction, the mean square error between the predicted and the mask corresponding to the true label is calculated as the correction loss. The process of calculating the correction loss in the region of inconsistency in the prediction for the unlabeled data includes the following steps: Regions with different pseudo-labels are designated as regions of inconsistent prediction. The combined loss of the positive and negative cross-entropy between the current sub-network pseudo-label and the prediction of another sub-network within the region of inconsistent prediction is calculated as the correction loss.

2. The semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning according to claim 1, characterized in that, For the labeled data, calculate the correction loss in the predicted inconsistency region. Calculate using the following formula: in, , These represent the two predicted pseudo-labels. This indicates the XOR operation. Indicates the prediction category is The Prediction of individual subnetworks, Indicates the first or second subnetwork. Indicates the mean square error loss. Indicates category The true mask corresponding to the region of inconsistency in prediction.

3. The semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning according to claim 1, characterized in that, For the unlabeled data, calculate the correction loss in the predicted inconsistency region. Calculate using the following formula: in, For the current subnetwork to the th Predicted label for each pixel. This represents the pseudo-label obtained from another sub-network. Indicates the current subnetwork's relationship to the th The confidence level of the predicted pseudo-label for each pixel. This indicates the number of pixels in the region of inconsistency in prediction.

4. The semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning according to claim 1, characterized in that, The first contrast loss Calculate using the following formula: in, Indicates the number of anchor points sampled. and These represent anchor points. The corresponding sets of positive and negative samples, , These represent positive and negative samples, respectively. Indicates the anchor point for sampling. Indicates the temperature coefficient. This represents the cosine similarity operator.

5. The semi-supervised medical image segmentation method based on mutual correction and pixel-level contrastive learning according to claim 1, characterized in that, The second contrast loss Calculate using the following formula: Specifically, for unlabeled data, Indicates the first The category of each pixel The confidence level for labeled data. , Indicates the first Projection features of each pixel, This indicates that all predicted pseudo-labels for the low-uncertainty region are... The set of pixels, This indicates the number of pixels in the low-uncertainty region. This indicates the result calculated in the current training batch. A collection of prototypes, Indicates the anchor point for sampling. express The corresponding pseudo-tag in the middle is The prototype.

Images