Semi-supervised medical image segmentation method based on uncertainty and attention mechanism

Abstract

The invention relates to a semi-supervised medical image segmentation method based on uncertainty and an attention mechanism, and the method comprises the steps: inputting a medical image into a trained segmentation model to achieve image segmentation, and the segmentation model comprises an encoder and three attention-enhanced UNet decoders; the training process of the segmentation model comprises the following steps: obtaining labeled medical image data and label-free medical image data, and obtaining a corresponding weak enhancement view and a strong enhancement view through weak enhancement operation and strong enhancement operation respectively; inputting the weakly enhanced view and the strongly enhanced view into a segmentation model, at the low resolution of a decoder, performing feature enhancement in a boundary region, at the full resolution of the decoder, obtaining a new enhanced sample through a CutMix strategy, and respectively calculating supervision loss and cross pseudo supervision loss at the low resolution and the full resolution; the sum of all losses of the above steps is used as a training target of each iteration. Compared with the prior art, the method has the advantages of enhancing boundary feature extraction, improving segmentation accuracy and the like.

Description

Technical Field

[0001] The present invention relates to the technical field of image processing, and in particular to a semi-supervised medical image segmentation method based on uncertainty and attention mechanism. Background Art

[0002] Medical image segmentation aims to accurately identify and locate anatomical structures, such as organ tissues or lesion regions, from medical images, and is an important part of computer-aided treatment. By automatically and quickly segmenting organ tissues or lesion regions, the efficiency of medical tasks such as pre-operative disease localization and post-operative treatment effect inspection has been greatly improved. Although existing fully supervised methods can show reliable segmentation performance on specific datasets, limited by the high annotation cost and inconsistent quality of the datasets, researchers have to consider mining potential structural information and position prior information in unlabeled data. Semi-supervised methods train the model to represent features containing organ structure and position information by mining the information of a small amount of labeled data and a large amount of unlabeled data, thus significantly reducing the dependence on labeled data and gradually approaching the segmentation performance under full supervision.

[0003] Existing semi-supervised algorithms can generally be divided into methods based on pseudo-label self-training, methods based on consistency learning, methods based on co-training, data augmentation methods, and contrastive learning methods. Among them, the method based on pseudo-label self-training uses the pseudo-labels predicted by the model for subsequent iterative training; the method based on consistency learning usually applies different perturbations at the data level, model level, and feature level, and imposes consistency constraints on different outputs of the same sample to achieve a regularization effect; the method based on co-training improves the robustness of feature representation by using different models to perceive different information of the same sample. In addition, strong data augmentation methods have been widely verified to be effective in semi-supervised learning, and complex and diverse perturbations can promote the alignment of decision boundaries and low-density regions. The contrastive learning methods in the field of semi-supervised medical image segmentation have strategies for defining positive and negative samples from the pixel, slice level, and patch level, and also have strategies for defining positive and negative samples based on pseudo-labels. Applying supervised contrastive learning according to pseudo-labels can utilize the prior knowledge of categories and help smooth the decision boundary.

[0004] Due to the blurred pixels in the boundary region of medical images, it is difficult for the model to learn the nearby information with only a small amount of labeled data. There are many existing attention variants that can perceive long-distance dependencies and class distribution information, while most of the existing semi-supervised segmentation methods are based on UNet with simple convolutional operations and lack a strategy to efficiently apply the attention mechanism to the semi-supervised scenario. In addition, considering that the amount of information contained in each unlabeled sample is different and the contribution to model training is also different, although the existing data augmentation methods apply a large amount of random and strong data augmentation to the unlabeled samples to expand the data distribution, the indiscriminate data augmentation may introduce difficult perturbation samples, resulting in low-quality pseudo-labels having a certain negative impact on model training, which further leads to the problem of misfitting of pseudo-labels.

[0005] In summary, there is currently a lack of a semi-supervised medical image segmentation method based on uncertainty and attention mechanism to solve or partially solve the aforementioned problems.

[0006] After retrieval, Chinese Patent Application Publication No. CN118587438A discloses a semi-supervised medical image segmentation method based on mutual correction and pixel-level contrast learning, which uses a pre-trained semi-supervised medical image segmentation model including two sub-networks to implement the segmentation of a given medical image. The training process of the semi-supervised medical image segmentation model includes the following steps: obtaining labeled data and unlabeled data as the input of the semi-supervised medical image segmentation model, obtaining predictions and pseudo-labels, and calculating the supervised loss and cross-pseudo-supervised loss; respectively calculating the prediction inconsistent regions of the classes or pseudo-labels for the predictions and pseudo-labels corresponding to the labeled data and unlabeled data, and calculating the correction loss within the prediction inconsistent regions; based on the prediction means of the two sub-networks, screening low-uncertainty pseudo-labels through uncertainty estimation, and calculating the comprehensive contrast loss between the anchor points and prototypes through anchor point sampling and prototype generation; based on the supervised loss, cross-pseudo-supervised loss, correction loss and comprehensive contrast loss, realizing the training of the semi-supervised medical image segmentation model. This existing patent application has problems such as insufficient ability to extract boundary features, the segmentation model relying on the confidence threshold resulting in loss of some information, and the complex random position and random degree of data augmentation being prone to introducing difficult samples and causing model training deviation.

[0007] How to effectively extract boundary features, improve the quality of pseudo-labels, and thus improve the accuracy of medical image segmentation has become a technical problem to be solved. Summary of the Invention

[0008] The objective of the present invention is to overcome the defects existing in the above-mentioned prior art and provide a semi-supervised medical image segmentation method based on uncertainty and attention mechanism, so as to alleviate the problems that it is difficult to accurately locate and extract fuzzy boundary regions and random strong data augmentation may introduce difficult samples, resulting in model training deviation.

[0009] The objective of the present invention can be achieved by the following technical solutions:

[0010] According to one aspect of the present invention, there is provided a semi-supervised medical image segmentation method based on uncertainty and attention mechanism. The medical image is input into a trained semi-supervised medical image segmentation model to implement image segmentation. The segmentation model includes an encoder and three attention-enhanced UNet decoders, and the three decoders respectively correspond to three sub-networks. The training process of the segmentation model includes the following steps:

[0011] Step 1: Obtain labeled and unlabeled medical image data, and respectively obtain corresponding weakly augmented views and strongly augmented views through weak augmentation operations and strong augmentation operations;

[0012] Step 2: Input the weakly augmented view and the strongly augmented view into the segmentation model, and calculate the supervised loss and the cross pseudo-supervised loss at the low resolution of the decoder respectively;

[0013] Step 3: Calculate the uncertainty based on the prediction means of the three sub-networks, identify the boundary regions, enhance the features in the boundary regions and then decode, and calculate the supervised loss and the cross pseudo-supervised loss at the full resolution of the decoder respectively;

[0014] Step 4: Calculate the average uncertainty at the patch level based on the prediction means of the three sub-networks, obtain new augmented samples through the uncertainty-guided CutMix strategy and then input them into the segmentation model, and calculate the supervised loss and the cross pseudo-supervised loss at the low resolution and the full resolution of the decoder respectively;

[0015] Step 5: Sum up all the losses in Steps 2 to 4 as the training objective for each iteration until the maximum number of iterations is reached.

[0016] Preferably, the Step 2 includes:

[0017] For the weakly augmented view and the strongly augmented view of the labeled data, calculate the predictions at the low resolution and the cross-entropy loss L ce and the Dice loss L dice with the downsampled ground truth labels as the supervised loss L ds_sup at the low resolution, which is defined as follows:

[0018]

[0019] where λ1 and λ2 are the weighting coefficients of two different losses, and Y ds represents the true label of the downsampling;

[0020] For the weakly augmented view and strongly augmented view of the unlabeled data, calculate the cross-entropy L ce and the Dice loss L dice as the cross pseudo-supervision loss L ds_cps at the low resolution, which is calculated as follows:

[0021]

[0022] where is the feature prediction before the i-th branch attention enhancement, represents the pseudo-label obtained by the prediction competition of the j-th and k-th branches, and λ1 and λ2 are the weighting coefficients of two different losses.

[0023] More preferably, for the weakly augmented view and strongly augmented view of the unlabeled data, the process of obtaining the pseudo-label of the winning competition is specifically: in the predicted class probability distribution of each pixel, select the predicted class probability with the largest value, and compare the predicted class probabilities of any two branches, and select the class corresponding to the larger one as the pseudo-label of the winning competition.

[0024] Preferably, the process of calculating the uncertainty and identifying the boundary region in step 3 includes:

[0025] Calculate the entropy of the predicted mean of any two branches as the uncertainty u, and use a dynamically growing threshold to determine the mask M of the high uncertainty region, and the formula is as follows:

[0026]

[0027] where ∈ = 1e -16 is a hyperparameter, t represents the current epoch number, and t max represents the preset maximum training epoch number, represents the average value of the two predicted class probability maps obtained by processing the input view through the j-th and k-th branches the predicted probability of the c-th class in i u represents the uncertainty map obtained by calculating the entropy from the predicted probability

[0028] Divide the mask M of the high-uncertainty region into windows of a fixed size. If any part of the high-uncertainty region is included in a window, then this window is identified as a boundary window.

[0029] Preferably, the enhancement of features in the boundary region in step 3 is specifically as follows:

[0030] Calculate the features after depthwise separable convolution and channel attention enhancement of the original features, and fuse the features before and after attention enhancement on the boundary window to represent the sparse enhanced features. The formula is as follows:

[0031]

[0032] where, f i represents the original feature at the i-th position in the feature map, represents the attention-enhanced feature at the i-th position in the feature map, f i ′ represents the feature at the i-th position after passing through the uncertainty-guided attention module. DepthWise_Conv() is the depthwise separable convolution operation, and Channel_Attention() is the channel attention operation.

[0033] Preferably, in step 3, calculating the supervised loss and the cross pseudo-supervised loss at the full resolution of the decoder includes:

[0034] For the weakly augmented view and the strongly augmented view of the labeled data, calculate the prediction probability and the cross-entropy loss L ce and the Dice loss L dice with the true label Y as the supervised loss L sup , which is defined as follows:

[0035]

[0036] where, λ1 and λ2 are the weighting coefficients of the two different losses, and Y represents the true label;

[0037] For the weakly augmented view and the strongly augmented view of the unlabeled data, calculate the cross-entropy loss L ce and the Dice loss L dice between the winning pseudo-label of the competition of any two sub-networks and the prediction of the remaining sub-network as the cross pseudo-supervised loss L cps , which is calculated as follows:

[0038]

[0039] where, is the prediction of the unlabeled data of the i-th branch, is the pseudo-label obtained by the prediction competition of the j-th and k-th branches.

[0040] Preferably, in step 4, the process of calculating the average uncertainty at the Patch level includes:

[0041] Calculate the entropy of the predicted mean of the three branches for the weakly augmented view and the strongly augmented view respectively, and use it as their uncertainty. The formula is as follows:

[0042]

[0043] where represents the predicted class probability map obtained by processing any input view through the i-th branch, represents the mean of the three predicted class probability maps of the three branches, represents the mean of the predicted probability maps obtained by processing the weakly augmented view through the three branches The predicted probability of the c-th class in represents the mean of the predicted probability maps obtained by processing the weakly augmented view through the three branches The predicted probability of the c-th class in, U w represents the predicted uncertainty map corresponding to the weakly augmented view, U s represents the predicted uncertainty map corresponding to the strongly augmented view, ∈ = 1e -16 is a hyperparameter;

[0044] Divide each view into multiple non-overlapping image patches, and calculate the mean uncertainty of each image patch as the uncertainty at the image patch level. The formula is as follows:

[0045]

[0046] where n×n is the number of non-overlapping image patches, represents the average uncertainty of the j-th image patch in the weakly augmented view, U w,i represents the uncertainty of the i-th pixel in the j-th image patch in the weakly augmented view, represents the average uncertainty of the j-th image patch in the strongly augmented view, U s,i represents the uncertainty of the i-th pixel in the j-th image patch in the strongly augmented view.

[0047] Preferably, in step 4, the process of obtaining new augmented samples through the uncertainty-guided CutMix strategy includes: According to the average uncertainty of the image patches in the weakly augmented view and the average uncertainty of the image patches in the strongly augmented view respectively determine the image patch mask M with the highest uncertainty in the weakly augmented view max and the image patch mask M with the lowest uncertainty in the strongly augmented viewmin Cover the high-uncertainty image patches of the weakly augmented view with the low-uncertainty image patches of the current strongly augmented view to construct a new augmented sample X s→w :

[0048]

[0049] where X w and X s are the weakly augmented and strongly augmented samples of the same image respectively, and are the high-uncertainty image patch masks of the weakly augmented sample X w and the low-uncertainty image patch masks of the strongly augmented sample X s respectively, and ⊙ represents element-wise multiplication operation.

[0050] Preferably, in step 4, calculating the supervised loss and the cross pseudo-supervised loss at the low resolution and full resolution of the decoder includes:

[0051] Send the uncertainty-guided CutMix augmented sample into the segmentation model. At the low resolution of the second-to-last layer of the decoder, for the labeled data, calculate the supervised loss L d ′ s_sup between the prediction and the downsampled CutMix ground truth label; for the unlabeled data, calculate the cross pseudo-supervised loss L d ′ s_cps , and the formula is as follows:

[0052]

[0053] where Y d ′ s is downsampled from the ground truth label Y′ of the corresponding CutMix, represents the winning pseudo-label of the i-th branch at the low resolution, represents the feature prediction at the low resolution of the i-th branch of the labeled data, represents the feature prediction before attention enhancement of the i-th branch of the unlabeled data, and λ1 and λ2 are the weighting coefficients of two different losses;

[0054] At the full resolution, for the labeled data, calculate the supervised loss L s ′ up between the prediction and the CutMix ground truth label; for the unlabeled data, calculate the cross pseudo-supervised loss L c ′ ps , and the formula is as follows:

[0055]

[0056] Among them, Y' represents the true label for implementing the corresponding CutMix strategy. represents the winning pseudo-label of the i-th branch. represents the CutMix augmented view of the labeled data The predicted class probability map obtained via the i-th branch. represents the CutMix augmented view of the unlabeled data The predicted class probability map obtained via the i-th branch.

[0057] Preferably, the sum of all losses in steps 2 to 4 is specifically:

[0058] L = L sup + L ds_sup + β1(L cps + L ds_cps )

[0059] L' = L s ′ up + L d ′ s_sup + β1(L c ′ ps + L d ′ s_cps )

[0060] L overall = L + L'

[0061] Among them, is the warm-up equation; t refers to the current epoch number; t max refers to the total number of epochs; L represents the weighted loss calculated for the weakly augmented views and strongly augmented views of the labeled data and unlabeled data; L ds_sup and L ds_cps are respectively the supervised loss and cross pseudo-supervised loss at low resolution in step 2; L sup and L cps are respectively the supervised loss and cross pseudo-supervised loss at full resolution in step 3; L' represents the weighted loss calculated for the CutMix augmented views of the labeled data and unlabeled data; L d ′ s_sup and L d ′ s_cps are respectively the supervised loss and cross pseudo-supervised loss at low resolution in step 4; L s ′ up and L c ′ ps are respectively the supervised loss and cross pseudo-supervised loss at full resolution in step 4; L overall refers to the sum of all losses, that is, the training objective for each iteration.

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

[0063] (1) Through the uncertainty-guided attention mechanism, the present invention identifies the fuzzy boundary region according to the uncertainty predicted by different branches, alleviating the problem that the model is difficult to effectively extract the features of the fuzzy organ boundary region; then, depthwise separable convolution and channel attention are used to enhance the feature extraction of the boundary region, promoting the feature representation ability in a sparse injection manner, effectively perceiving local information and class distribution information, and reducing the possibility of network overfitting; the combined use of multiple loss functions at different stages improves the overall performance of the segmentation model, including the segmentation accuracy rate.

[0064] (2) Perceiving a more complete data distribution: Through the uncertainty-guided CutMix enhancement mechanism, the present invention dynamically locates the perturbation addition region according to the uncertainty predicted by different branches, and uses low-uncertainty image patches in different enhanced views to replace high-uncertainty image patches, thereby using high-quality pseudo-labels to supervise difficult perturbation samples, forcing the model to learn new information from difficult samples, stabilizing the model training process while further expanding the data perception distribution, promoting the alignment of the decision boundary and the low-density region, and significantly improving the accuracy of medical image segmentation of the present application compared with the prior art with less labeled data.

[0065] (3) Improving the quality of pseudo-labels and stabilizing the model training: The present invention compares the prediction probabilities of multiple branches and selects the category corresponding to the larger one as the winning pseudo-label for pseudo-supervising the prediction of the remaining one branch. Similar to the way that ensemble learning obtains reliable prediction results through majority voting of multiple weak learners, this method selects high-confidence pseudo-labels, greatly alleviating the challenge that setting a high threshold to filter pseudo-labels will limit the model from generating reliable true positive pseudo-labels, while setting a low threshold to filter pseudo-labels cannot effectively filter out unreliable pseudo-labels, thereby stabilizing and optimizing the model training process.

[0066] (4) Optimizing the model training process: The present invention uses a warm-up equation to control the weights of the supervised loss and the cross pseudo-supervised loss. The weight of the cross pseudo-supervised loss continuously increases with the training iteration, preventing the risk of incorrect fitting caused by the model relying too much on incorrect pseudo-labels in the initial stage of training, and promoting the model to smoothly adapt to different loss terms, improving the model generalization ability, and enabling it to learn more robust feature representations. BRIEF DESCRIPTION OF THE DRAWINGS

[0067] Figure 1 It is a schematic flowchart of the semi-supervised medical image segmentation method based on uncertainty and attention mechanism in the present invention;

[0068] Figure 2Schematic diagram of the training framework of the semi - supervised medical image segmentation method based on uncertainty and attention mechanism in the present invention;

[0069] Figure 3 Schematic diagram of the segmentation model embedded with uncertainty - guided attention mechanism in the present invention. Detailed implementation manners

[0070] Next, the technical solutions in the embodiments of the present invention will be clearly and completely described in conjunction with the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are part of the embodiments of the present invention, rather than all embodiments. All other embodiments obtained by those of ordinary skill in the art based on the embodiments of the present invention without creative efforts shall fall within the protection scope of the present invention.

[0071] Embodiment 1

[0072] This embodiment relates to a semi - supervised medical image segmentation method based on uncertainty and attention mechanism. Starting from two aspects: the fuzziness of the organ boundary region in medical images and the insufficient perception of the sample distribution, an uncertainty - guided attention mechanism and an uncertainty - guided CutMix strategy are designed, and based on this, a segmentation model that obtains pseudo - labels by competing for victory for cross - pseudo - supervision is built to effectively alleviate the problem that it is difficult to effectively extract features due to the fuzzy organ boundary, and improve the problem of training deviation caused by introducing difficult samples through random data augmentation.

[0073] See Figure 1 , this method includes the following steps:

[0074] S1. Construct a segmentation model composed of a ResNet50 shared encoder initialized with ImageNet pre - trained weights and three randomly initialized UNet decoders with attention enhancement. The three decoders correspond to three sub - networks respectively. Obtain labeled data and unlabeled data, and respectively obtain the corresponding weakly augmented views and strongly augmented views through weak augmentation operations and strong augmentation operations, and input them into the constructed segmentation model. Calculate the supervised loss and cross - pseudo - supervision loss respectively at the low resolution of the decoder.

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

[0076] S11. Construct a model composed of a ResNet50 shared encoder initialized with ImageNet pre - trained weights and three randomly initialized UNet decoders with attention enhancement. The three decoders correspond to three sub - networks respectively, forming a semi - supervised medical image segmentation model.

[0077] S12. Apply weak data augmentation operations to the labeled data and unlabeled data, including random rotation, randomly selecting regions for scaling, random affine transformation, and elastic transformation, to obtain weakly augmented views. On this basis, apply strong data augmentation Color Jitter to adjust brightness, contrast, saturation, and hue to obtain strongly augmented views.

[0078] S13. As Figure 2 , feed the weakly augmented views and strongly augmented views into the segmentation model, and obtain the predicted class probability map via the classification head at the feature map of the penultimate layer of the decoder For the unlabeled data, select the predicted probability with the largest value in the predicted class probability distribution of each pixel, and compare the predicted probabilities of any two branches, and select the larger one. The corresponding class is used as the pseudo-label that wins the competition. The formula is as follows:

[0079]

[0080] Among them, i, j, k represent different branches, indicating different value combinations of the three branches; C represents the total number of classes of the current data; represents the pseudo-label obtained by the prediction competition victory of the j and k branches, and is used to supervise the prediction of the i-th branch; represents the predicted probability of the c-th class; and are the predicted probabilities of the j-th branch and the k-th branch in the c-th class respectively; Argmax is a function used to find the index of the maximum value in an array or matrix.

[0081] S14. For the weakly augmented views and strongly augmented views of the labeled data, calculate the prediction at the low resolution and the cross-entropy loss L ce and Dice loss L dice between the downsampled ground truth labels as the supervision loss L ds_sup at the low resolution, which is defined as follows:

[0082]

[0083] Among them, λ1 and λ2 are the weighting coefficients of the two different losses, and Y ds represents the downsampled ground truth label.

[0084] S15. For the weakly augmented views and strongly augmented views of the unlabeled data, calculate the cross-entropy and Dice losses between the pseudo-labels that win the competition of any two sub-networks and the predictions of the remaining sub-networks as the cross-pseudo-supervision loss L ds_cps at the low resolution, which is calculated as follows:

[0085]

[0086] Among them, is the feature prediction before the enhancement of the i-th branch attention, represents the pseudo-label obtained by the prediction competition winning of the two branches of j and k. The pseudo-label of the winning competition of the weak enhancement view is adopted for different enhanced views to impose the consistency constraint between different enhanced views. λ1 and λ2 are the weighting coefficients of two different losses.

[0087] S2. As Figure 3 , construct an uncertainty-guided attention mechanism. Based on the prediction mean of multiple sub-networks, calculate the uncertainty recognition boundary region, and use depthwise separable convolution and channel attention in this region to enhance the features.

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

[0089] S21. Calculate the entropy of the prediction mean of any two branches as the uncertainty u, and use a dynamically growing threshold to determine the mask M of the high-uncertainty region. The formula is as follows:

[0090]

[0091] Among them, v = 1e -16 is a hyperparameter used to prevent the situation of log(0). t represents the current epoch number, and t max represents the preset maximum training epoch number, represents the predicted probability of the c-th category in the average value of the two predicted category probability maps obtained by processing the input view through the two branches of j and k. u i represents the uncertainty map obtained by calculating the entropy from and applied to the i-th branch module, where i, j, and k represent different branches, and C represents the total number of categories.

[0092] S22. Divide the mask M of the high-uncertainty region into windows of a fixed size of 2×2. As long as any part of the high-uncertainty region is included in this window, that is, as long as M = 1 exists in the window, this window is determined as the boundary window (i.e., the organ boundary window).

[0093] S23. Calculate the features enhanced by depthwise separable convolution and channel attention for the original features, and fuse the features before and after attention enhancement on these boundary windows to represent the sparse enhanced features. The formula is as follows:

[0094]

[0095] Among them, f iRepresents the original feature at the i-th position in the feature map, Represents the attention-enhanced feature at the i-th position in the feature map, f i ′ represents the feature at the i-th position after passing through the uncertainty-guided attention module. DepthWise_Conv() is the depthwise separable convolution operation, and Channel_Attention() is the channel attention operation.

[0096] The calculation formulas for the depthwise separable convolution and the channel attention-enhanced feature are as follows:

[0097] Q = DWConv(Conv 1×1 (x))

[0098] K = DWConv(Conv 1×1 (x))

[0099] V = DWConv(Conv 1×1 (x))

[0100]

[0101] Among them, x represents the input feature matrix, Q, K, and V respectively represent the query, key, and value in the attention mechanism. Conv 1×1 Represents the convolution operation with a 1×1 convolution kernel. DWConv represents the depthwise separable convolution operation. τ represents the learnable temperature coefficient, which is combined with the normalization calculation to perform redistribution on the result of the dot product of Q T and K, so that the result falls in the non-saturated region of the softmax operation, preventing the problem of gradient disappearance caused by the chain multiplication when calculating the gradient.

[0102] S3. Further decode the enhanced feature, and calculate the supervised loss and the cross pseudo-supervised loss at the full resolution respectively.

[0103] Specifically, step S3 specifically includes the following sub-steps:

[0104] S31. Send the enhanced feature f′ into the subsequent module of the decoder, and obtain the predicted class probability via the classification head at the full-resolution feature map For unlabeled data, select the predicted class probability with the largest value in the predicted class probability distribution of each pixel, and compare the predicted class probabilities of any two branches, and select the larger one. The corresponding class is used as the pseudo-label that wins the competition The formula is as follows:

[0105]

[0106] Among them, i, j, and k are different from each other, representing different value combinations of the three branches, and C represents the total number of classes of the current data.

[0107] S32. Calculate the prediction probabilities for the weakly augmented views and strongly augmented views of the labeled data and the cross-entropy loss L between them and the true label Y ce and the Dice loss L dice , which serve as the supervised loss L sup , and are defined as follows:

[0108]

[0109] where λ1 and λ2 are the weighting coefficients of the two different losses, and Y represents the true label.

[0110] S33. Calculate the cross-entropy and Dice losses between the pseudo-labels that win the competition of any two sub-networks and the predictions of the remaining sub-networks for the weakly augmented views and strongly augmented views of the unlabeled data, which serve as the cross-pseudo-supervised loss L cps , and are calculated as follows:

[0111]

[0112] where is the prediction of the unlabeled data for the i-th branch, is the pseudo-label obtained by the prediction competition of the j-th and k-th branches. For different augmented views, the pseudo-labels that win the competition of the weakly augmented views are adopted to impose the consistency constraint between different augmented views.

[0113] S4. Construct an uncertainty-guided CutMix augmentation strategy mechanism. Based on the prediction means of multiple sub-networks, calculate the average uncertainty at the Patch level, and obtain new augmented samples through the uncertainty-guided CutMix strategy. Specifically: Calculate the entropy of the prediction means obtained by the input views through three branches as the uncertainty map, divide each view into multiple 56×56 image patches, then calculate the average uncertainty of each image patch as the uncertainty at the image patch level, and use the low-uncertainty image patches of the strongly augmented views to cover the high-uncertainty image patches of the weakly augmented views to construct new augmented samples.

[0114] Specifically, step S4 specifically includes the following sub-steps:

[0115] S41. Calculate the entropy of the prediction means of the three branches for the weakly augmented views and strongly augmented views respectively as their uncertainties, and the formula is as follows:

[0116]

[0117] where n = 56, and an image with a resolution of 224×224 will thus be divided into 16 image patches. Denote the predicted class probability map obtained by processing any input view through the $i$-th branch. Denote the mean of the three predicted class probability maps of the three branches. Denote the mean of the predicted probability maps obtained by processing the weakly augmented view through the three branches. The predicted probability of the $c$-th class in Denote the mean of the predicted probability maps obtained by processing the weakly augmented view through the three branches. The predicted probability of the $c$-th class in $U$. w Denote the predicted uncertainty map corresponding to the weakly augmented view, $U$. s Denote the predicted uncertainty map corresponding to the strongly augmented view.

[0118] S42. Divide each view into multiple non-overlapping $56\times56$ image patches, and calculate the mean uncertainty of each image patch as the uncertainty at the image patch level. The formula is as follows:

[0119]

[0120] where $n = 56$. Denote the average uncertainty of the $j$-th image patch in the weakly augmented view, $U$. w,i Denote the uncertainty of the $i$-th pixel in the $j$-th image patch of the weakly augmented view. Denote the average uncertainty of the $j$-th image patch in the strongly augmented view, $U$. s,i Denote the uncertainty of the $i$-th pixel in the $j$-th image patch of the strongly augmented view.

[0121] S43. According to the average uncertainty of the image patches in the weakly augmented view and the average uncertainty of the image patches in the strongly augmented view respectively determine the image patch mask $M$ with the highest uncertainty in the weakly augmented view max and the image patch mask $M$ with the lowest uncertainty in the strongly augmented view min , and use the low-uncertainty image patches of the current strongly augmented view to cover the high-uncertainty image patches of the weakly augmented view to construct a new augmented sample $X$. s→w .

[0122]

[0123] where $X$ w and $X$ s are the weakly augmented and strongly augmented samples of the same image, and are the high-uncertainty image patch mask of $X$ w and the low-uncertainty image patch mask of $X$ s respectively, and $\odot$ represents the element-wise multiplication operation.

[0124] S5. Feed the uncertainty-guided CutMix augmented samples into the segmentation model, and calculate the supervised loss and cross pseudo-supervised loss at the low resolution and full resolution respectively.

[0125] Specifically, step S5 includes the following sub-steps:

[0126] S51. Feed the uncertainty-guided CutMix augmented samples into the model. At the low resolution of the penultimate layer of the decoder, for the labeled data, calculate the supervised loss \(L'\) between the prediction and the downsampled CutMix ground truth label; for the unlabeled data, calculate the cross pseudo-supervised loss \(L'\) between the prediction and the winning pseudo-label, and the formula is as follows: d ′ s_sup ; for the unlabeled data, calculate the cross pseudo-supervised loss \(L'\) between the prediction and the winning pseudo-label, d ′ s_cps The formula is as follows:

[0127]

[0128] where \(Y'\) d ′ s is obtained by downsampling the ground truth label \(Y'\) of the corresponding CutMix, represents the winning pseudo-label of the \(i\)-th branch at the low resolution, represents the feature prediction at the low resolution of the \(i\)-th branch of the labeled data, represents the feature prediction before attention enhancement of the \(i\)-th branch of the unlabeled data, and \(\lambda_1\) and \(\lambda_2\) are the weighting coefficients of the two different losses.

[0129] S52. Feed the low-resolution feature map into the uncertainty-guided attention module to enhance the feature representation. The enhanced features are further decoded. At the full resolution, for the labeled data, calculate the supervised loss \(L'\) between the prediction and the CutMix ground truth label; for the unlabeled data, calculate the cross pseudo-supervised loss \(L'\) between the prediction and the winning pseudo-label, and the formula is as follows: s ′ up ; for the unlabeled data, calculate the cross pseudo-supervised loss \(L'\) between the prediction and the winning pseudo-label, c ′ ps The formula is as follows:

[0130]

[0131] where \(Y'\) represents the ground truth label of the corresponding CutMix, represents the winning pseudo-label of the \(i\)-th branch, represents the CutMix augmented view of the labeled data The predicted class probability map obtained via the \(i\)-th branch, represents the CutMix augmented view of the unlabeled data The predicted class probability map obtained via the \(i\)-th branch.

[0132] Sum all losses, calculated as follows:

[0133] L = L sup + L ds_sup + β1(L cps + L ds_cps )

[0134] L′ = L s ′ up + L d ′ s_sup + β1(L c ′ ps + L d ′ s_cps )

[0135] L overall = L + L′

[0136] Wherein, is the preheating equation, t refers to the current epoch number, and t max refers to the total number of epochs. L represents the weighted loss calculated for the weakly augmented views and strongly augmented views of the labeled data and unlabeled data, L′ represents the weighted loss calculated for the CutMix augmented views of the labeled data and unlabeled data, and L overall refers to the sum of all losses, that is, the training objective for each iteration.

[0137] Example 2

[0138] This example also involves the application and verification of a semi-supervised medical image segmentation method based on uncertainty and attention mechanism. Using the segmentation model of Example 1, the segmentation model is updated by the gradient backpropagation method. When the training reaches the maximum number of times, the trained segmentation model is used to make predictions on any test image..

[0139] Table 1 shows the comparison of the segmentation effects of this method and existing algorithms on the Synapse abdominal multi-organ nine-classification dataset when the proportion of labeled data is 15% and 30% respectively (↑ indicates that the larger the value of this index, the better, and ↓ indicates that the smaller the value of this index, the better).

[0140] Table 1

[0141]

[0142]

[0143] From the comparison of various indicators in Table 1, it can be seen that when using only less labeled data, compared with the prior art, the proposed method can effectively identify the fuzzy boundary regions (i.e., difficult segmentation regions) in medical images through the uncertainty-guided attention mechanism. By adopting depthwise separable convolution and channel attention at these sparse positions, it can effectively perceive local information and class distribution information, and at the same time, the feature enhancement by sparse injection reduces the possibility of network overfitting. Meanwhile, the uncertainty-guided CutMix augmentation can, to a certain extent, prompt the model to learn a more comprehensive data distribution, and use high-quality pseudo-labels to force the model to learn new information from difficult samples, ultimately significantly improving the accuracy of medical image segmentation.

[0144] The method of the present invention uses the uncertainty-guided attention mechanism to identify the fuzzy boundary regions according to the uncertainty predicted by different branches, and then uses depthwise separable convolution and channel attention to enhance the feature extraction of the boundary regions, promoting the feature representation ability in a sparse injection manner, and alleviating the problem of insufficient feature extraction ability in the fuzzy boundary regions.

[0145] The method of the present invention uses the uncertainty-guided CutMix augmentation mechanism to dynamically locate the perturbation addition regions according to the uncertainty predicted by different branches, and uses the low-uncertainty image patches in different augmented views to replace the high-uncertainty image patches, thereby using high-quality pseudo-labels to supervise difficult perturbation samples, while further expanding the data perception distribution and stabilizing the model training process, and ultimately improving the overall performance of the segmentation model.

[0146] Embodiment 3

[0147] The electronic device of the present invention includes a central processing unit (CPU), which can execute various appropriate actions and processes according to the computer program instructions stored in the read-only memory (ROM) or the computer program instructions loaded from the storage unit into the random access memory (RAM). In the RAM, various programs and data required for device operation can also be stored. The CPU, ROM, and RAM are connected to each other through a bus. The input / output (I / O) interface is also connected to the bus.

[0148] Multiple components in the device are connected to the I / O interface, including: an input unit, such as a keyboard, a mouse, etc.; an output unit, such as various types of displays, speakers, etc.; a storage unit, such as a disk, an optical disc, etc.; and a communication unit, such as a network card, a modem, a wireless communication transceiver, etc. The communication unit allows the device to exchange information / data with other devices through a computer network such as the Internet and / or various telecommunication networks.

[0149] The processing unit executes the various methods and processes described above. For example, in some embodiments, the method may be implemented as a computer software program tangibly embodied in a machine-readable medium, such as a storage unit. In some embodiments, part or all of the computer program may be loaded and / or installed onto the device via the ROM and / or the communication unit. When the computer program is loaded into the RAM and executed by the CPU, one or more steps of the methods described above may be performed. Alternatively, in other embodiments, the CPU may be configured to execute the method by any other suitable means (e.g., by means of firmware).

[0150] The functions described above herein can be performed at least in part by one or more hardware logic components. By way of example, and without limitation, the types of hardware logic components that may be used include: field programmable gate arrays (FPGAs), application specific integrated circuits (ASICs), application specific standard products (ASSPs), systems on a chip (SOCs), complex programmable logic devices (CPLDs), and the like.

[0151] The program code for implementing the methods of the present invention may be written in any combination of one or more programming languages. These program codes may be provided to a processor or controller of a general purpose computer, a special purpose computer, or other programmable data processing apparatus, such that the program codes, when executed by the processor or controller, cause the functions / operations specified in the flowchart and / or block diagram to be implemented. The program code may be executed entirely on the machine, partly on the machine, as a stand-alone software package partly on the machine and partly on a remote machine, or entirely on the remote machine or server.

[0152] In the context of the present invention, a machine-readable medium may be a tangible medium that can contain or store a program for use by or in connection with an instruction execution system, apparatus, or device. A machine-readable medium may be a machine-readable signal medium or a machine-readable storage medium. A machine-readable medium may include, but is not limited to, electronic, magnetic, optical, electromagnetic, infrared, or semiconductor systems, apparatus, or devices, or any suitable combination of the foregoing. More specific examples of a machine-readable storage medium would include an electrical connection based on one or more wires, a portable computer diskette, a hard disk, a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or Flash memory), an optical fiber, a portable compact disc read-only memory (CD-ROM), an optical storage device, a magnetic storage device, or any suitable combination of the foregoing.

[0153] As described above, it is only the specific implementation manner of the present invention, but the protection scope of the present invention is not limited thereto. Any person skilled in the art within the technical scope disclosed by the present invention can easily think of various equivalent modifications or substitutions, and these modifications or substitutions should be covered within the protection scope of the present invention. Therefore, the protection scope of the present invention shall be subject to the protection scope of the claims.

Claims

1. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism, characterized in that Input the medical image into the trained semi-supervised medical image segmentation model to achieve image segmentation. The segmentation model includes an encoder and three attention-enhanced UNet decoders, and the three decoders respectively correspond to three sub-networks; The training process of the segmentation model includes the following steps: Step 1: Obtain labeled and unlabeled medical image data, and respectively obtain corresponding weakly augmented views and strongly augmented views through weak augmentation operations and strong augmentation operations; Step 2: Input the weakly augmented view and the strongly augmented view into the segmentation model, and calculate the supervised loss and the cross pseudo-supervised loss respectively at the low resolution of the decoder; Step 3: Calculate the uncertainty and identify the boundary region based on the predicted means of the three sub-networks, enhance the features in the boundary region and then decode, and calculate the supervised loss and the cross pseudo-supervised loss respectively at the full resolution of the decoder; Step 4: Calculate the average uncertainty at the patch level based on the predicted means of the three sub-networks, obtain new augmented samples through the uncertainty-guided CutMix strategy and then input them into the segmentation model, and calculate the supervised loss and the cross pseudo-supervised loss respectively at the low resolution and the full resolution of the decoder; Step 5: Sum all the losses in Steps 2-4 as the training objective for each iteration until the maximum number of iterations is reached.

2. The semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, wherein The said Step 2 includes: Compute the predictions at low resolution for the weakly augmented views and strongly augmented views of the labeled data The cross-entropy loss L between the downsampled ground truth labels ce and the Dice loss L dice , as the supervision loss L at low resolution ds_sup , is defined as follows: where λ1 and λ2 are the weighted coefficients of two different losses, and Y ds represents the true label after downsampling; For the weakly augmented view and strongly augmented view of unlabeled data, calculate the cross-entropy \(L\) between the winning pseudo-labels of any two sub-networks and the predictions of the remaining sub-networks ce and the Dice loss \(L\) dice , as the cross pseudo-supervision loss \(L\) at low resolution ds_cps , which is calculated as follows: Among them, is the feature prediction before the enhancement of the i-th branch attention, represents the pseudo-label obtained by winning the prediction competition between the two branches j and k, and λ1 and λ2 are the weighting coefficients of two different losses.

3. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 2, characterized in that, For the weakly augmented view and the strongly augmented view of the unlabeled data, the process of obtaining the competing winning pseudo-labels is specifically as follows: In the predicted class probability distribution of each pixel, select the predicted class probability with the largest value, and compare the predicted class probabilities of any two branches, and select the class corresponding to the larger one as the competing winning pseudo-label.

4. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that, The process of calculating the uncertainty and identifying the boundary region in the said Step 3 includes: Calculate the entropy of the predicted means of any two branches as the uncertainty u, and use a dynamically growing threshold to determine the mask M of the high-uncertainty region. The formula is as follows: Among them, ∈ = 1e -16 is a hyperparameter, t represents the current epoch number, and t max represents the preset maximum number of training epochs. represents the average of the two predicted class probability maps obtained by processing the input view through the two branches j and k. The predicted probability of the c-th class in i is represented by the predicted probability The uncertainty map calculated by entropy and applied to the i-th branch module, where i, j, and k represent different branches, and C represents the total number of classes. Divide the mask M of the high-uncertainty region into windows of a fixed size. If a window contains any part of the high-uncertainty region, then this window is identified as a boundary window.

5. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that The process of enhancing the features in the boundary region in the said Step 3 is specifically as follows: Calculate the features enhanced by depthwise separable convolution and channel attention on the original features, and fuse the features before and after attention enhancement on the boundary window to represent the sparse enhanced features. The formula is as follows: Among them, f i represents the original feature at the i-th position in the feature map, represents the attention-enhanced feature at the i-th position in the feature map, and f i ' represents the feature at the i-th position after passing through the uncertainty-guided attention module. DepthWise_Conv() is the depthwise separable convolution operation, and Channel_Attention() is the channel attention operation.

6. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that In the said Step 3, calculating the supervised loss and the cross pseudo-supervised loss respectively at the full resolution of the decoder includes: Calculate the prediction probabilities for weakly augmented views and strongly augmented views of the labeled data The cross-entropy loss L between the predicted probabilities and the true label Y ce and the Dice loss L dice , as the supervised loss L sup , are defined as follows: Among them, λ1 and λ2 are the weighting coefficients of two different losses, and Y represents the true label; For the weakly augmented view and strongly augmented view of the unlabeled data, calculate the cross-entropy loss \(L\) between the winning pseudo-labels of any two sub-networks and the predictions of the remaining sub-networks ce and the Dice loss \(L\) dice , as the cross pseudo-supervision loss \(L\) cps , which is calculated as follows: Among them, is the prediction of the unlabeled data of the i-th branch, is the pseudo-label obtained by winning the prediction competition between the j-th and k-th branches.

7. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that, In the said Step 4, the process of calculating the average uncertainty at the Patch level includes: Calculate the entropy of the predicted means of the three branches for the weakly augmented view and the strongly augmented view respectively, and use it as their uncertainty. The formula is as follows: Among them, represents the predicted class probability map obtained by processing any input view through the i-th branch, represents the mean of the three predicted class probability maps of the three branches, represents the mean of the predicted probability maps obtained by processing the weakly augmented view through the three branches for the predicted probability of the c-th class in, represents the mean of the predicted probability maps obtained by processing the weakly augmented view through the three branches for the predicted probability of the c-th class in, U w represents the predicted uncertainty map corresponding to the weakly augmented view, U s represents the predicted uncertainty map corresponding to the strongly augmented view, ∈ = 1e -16 is a hyperparameter; Divide each view into multiple non-overlapping patches, and calculate the average uncertainty of each patch as the uncertainty at the patch level. The formula is as follows: where n×n is the number of non-overlapping image patches, represents the average uncertainty of the j-th image patch in the weakly enhanced view, U w,i represents the uncertainty of the i-th pixel in the j-th image patch in the weakly enhanced view, represents the average uncertainty of the j-th image patch in the strongly enhanced view, U s,i represents the uncertainty of the i-th pixel in the j-th image patch in the strongly enhanced view.

8. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that, In the above-mentioned step 4, the process of obtaining new augmented samples through the uncertainty-guided CutMix strategy includes: according to the average uncertainty of the image patches in the weakly augmented view and the average uncertainty of the image patches in the strongly augmented view respectively determine the image patch mask M with the highest uncertainty in the weakly augmented view max and the image patch mask M with the lowest uncertainty in the strongly augmented view min , use the low-uncertainty image patches of the current strongly augmented view to cover the high-uncertainty image patches of the weakly augmented view, and construct a new augmented sample X s→w : Among them, X w and X s are the weakly enhanced and strongly enhanced samples of the same image, and are the high-uncertainty image patch masks of the weakly enhanced sample X w and the low-uncertainty image patch masks of the strongly enhanced sample X s respectively. ⊙ represents the element-wise multiplication operation.

9. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that, In the said Step 4, calculating the supervised loss and the cross pseudo-supervised loss respectively at the low resolution and the full resolution of the decoder includes: Send the uncertainty-guided CutMix augmented samples into the segmentation model. At the low resolution of the second-to-last layer of the decoder, for the labeled data, calculate the supervised loss \(L\) between the prediction and the downsampled CutMix ground truth label. For the unlabeled data, calculate the cross pseudo-supervised loss \(L'\) between the prediction and the winning pseudo label. The formula is as follows: d ′ s_sup ; For the unlabeled data, calculate the cross pseudo-supervised loss \(L'\) between the prediction and the winning pseudo label. The formula is as follows: d ′ s_cps The formula is as follows: Among them, Y d ′ s is downsampled from the true label Y' that performs the corresponding CutMix, represents the winning pseudo-label of the i-th branch at low resolution, represents the feature prediction at the low resolution of the i-th branch of the labeled data, represents the feature prediction before attention enhancement of the i-th branch of the unlabeled data, and λ1 and λ2 are the weighting coefficients of two different losses; At full resolution, for labeled data, calculate the supervised loss $L$ between the prediction and the CutMix ground truth label s ′ up ; for unlabeled data, calculate the cross pseudo-supervised loss $L$ between the prediction and the winning pseudo-label in competition c ′ ps , and the formula is as follows: Among them, Y′ represents the true label for executing the corresponding CutMix strategy, represents the winning pseudo-label of the i-th branch, represents the CutMix augmented view of the labeled data The predicted class probability map obtained via the i-th branch, represents the CutMix augmented view of the unlabeled data The predicted class probability map obtained via the i-th branch.

10. A semi-supervised medical image segmentation method based on uncertainty and attention mechanism according to claim 1, characterized in that the sum All the losses in Steps 2-4 are specifically: L = L sup + L ds_sup + β1(L cps + L ds_cps ) L′ = L s ′ up + L d ′ s_sup + β1(L c ′ ps + L d ′ s_cps ) L overall = L + L' Among them, is the preheating equation; t refers to the current epoch number; t max refers to the total number of epochs; L represents the weighted loss calculated for the weakly augmented views and strongly augmented views of the labeled data and unlabeled data; L ds_sup and L ds_cps are the supervised loss and cross-pseudo-supervised loss at low resolution in step 2 respectively; L sup and L cps are the supervised loss and cross-pseudo-supervised loss at full resolution in step 3 respectively; L′ represents the weighted loss calculated for the CutMix augmented views of the labeled data and unlabeled data; L d ′ s_sup and L d ′ s_cps are the supervised loss and cross-pseudo-supervised loss at low resolution in step 4 respectively; L s ′ up and L c ′ ps are the supervised loss and cross-pseudo-supervised loss at full resolution in step 4 respectively; L overall refers to the sum of all losses, that is, the training objective for each iteration.

Images