Semi-supervised medical image segmentation method for generating confidence mask based on soft pseudo label

Abstract

The invention provides a semi-supervised medical image segmentation method for generating a confidence mask based on a soft pseudo tag, and the method comprises the steps: collecting and preprocessing different types of medical image data sets, and constructing a training set and a test set; selecting a semi-supervised medical image segmentation network; carrying out basic segmentation model selection on the segmentation network according to different medical image data set characteristics; inputting the data of the training set into a basic segmentation model, generating a mask by using a soft pseudo tag result between different branch networks, and training the basic segmentation model; data of a test set is input into a trained basic segmentation model to obtain a reasoning result, a segmentation index is calculated, the basic segmentation model is evaluated, under the condition that less data is used, the accuracy of image segmentation is guaranteed, meanwhile, the training time cost is reduced, a semi-supervised method is separated from model selection, the method can be suitable for different scenes, and the efficiency of image segmentation is improved. Popularization and application are facilitated.

Description

Technical Field

[0001] The present invention belongs to the technical field of deep learning segmentation models, and in particular to a semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels. Background Art

[0002] In the existing medical field, semi-supervised medical image segmentation is mainly used to solve the problem of difficulty in obtaining high-precision labels in clinical scenarios. It uses a small amount of labeled data and a large amount of unlabeled data to achieve the same effect as fully supervised image segmentation.

[0003] In this field, semi-supervised learning methods mainly include pseudo-labeling technology and consistency regularization technology; the segmentation model is a teacher-student network model based on the traditional U-shaped architecture; among them:

[0004] ① Pseudo-label technology: using a model trained with labeled data to predict unlabeled data, the result is directly used as the true label, and then put back into the model for retraining; the above method is limited by the accuracy of the model trained with labeled data;

[0005] ② Consistency regularization technology: mainly refers to the model encouraging the consistency of the segmentation results of the perturbed input and the original input, or the consistency of the segmentation results of different segmentation models for the same input.

[0006] The above scheme is mainly embodied in the prior art as follows:

[0007] 1. Similar solutions using soft fake label technology: UPRC;

[0008] This technical abstract records: Gross target volume (GTV) segmentation plays an irreplaceable role in radiotherapy planning for pediatric pharyngeal carcinoma (NPC);

[0009] Although convolutional neural networks (CNNs) have achieved good performance in this task, they rely on a large number of labeled images for training, which is expensive and time-consuming to acquire; we propose a new framework with uncertainty-rectified pyramid consistency (URPC) regularization for semi-supervised NPC GTV segmentation; specifically, we extend a backbone segmentation network to produce pyramid predictions at different scales; the pyramid prediction network (PPNet) is supervised by the ground truth of labeled images and a multi-scale consistency loss for unlabeled images, motivated by the fact that predictions for the same input at different scales should be similar and consistent; however, since these predictions are at different resolutions, encouraging them to be directly consistent at each pixel is less robust and may lose some fine details.

[0010] To address this issue, an uncertainty correction module was designed to enable the framework to gradually learn from meaningful and reliable consensus regions at different scales; experimental results on a dataset with 258 NPC MR images showed that when only 10% or 20% of the images were labeled, the above method greatly improved the segmentation performance by utilizing unlabeled images; in addition, when 50% of the images were labeled, the average Dice score of URPC was 82.74%, which was close to fully supervised learning.

[0011] The outstanding feature of the above paper is the semi-supervised segmentation model designed for a specific medical imaging dataset, which is very different from the present invention. Soft pseudo labels mainly refer to that when the model predicts, it does not directly obtain the "prediction label" based on the prediction probability, but sharpens the prediction probability (understood as sharpening with temperature hyperparameters). The above scheme is based on the result after sharpening to ensure the consistency of model prediction, without considering the mask and the overall architecture of the model design.

[0012] 2. Using a mask-based consistency regularization scheme: MC-Net;

[0013] The technical abstract describes: A new mutual consistency network (MC-Net+) to effectively utilize unlabeled data for semi-supervised medical image segmentation; The MC-Net+ model is motivated by the observation that deep models trained with limited annotations tend to output highly uncertain and easily misclassified predictions in fuzzy areas of medical image segmentation (e.g., glue edges or thin branches). Utilizing these challenging samples can make semi-supervised segmentation model training more effective;

[0014] Therefore, the MC-Net+ model consists of two new designs;

[0015] First, the model contains a shared encoder and multiple slightly different decoders (i.e., using different upsampling strategies); the statistical differences of the outputs of multiple decoders are calculated to represent the uncertainty of the model, which indicates the difficult areas for the model to segment;

[0016] Second, we apply a novel mutual consistency constraint between the probability output of one decoder and the soft pseudo labels of other decoders. In this way, we minimize the discrepancy between multiple outputs (i.e., model uncertainty) during training and force the model to produce invariant results in such challenging regions, aiming to regularize model training.

[0017] MCNet does not consider the concept of mask and directly calculates the cross MSE loss of soft false labels and output; therefore, it does not reproduce well in the comparative experiment of the present invention; the number and depth of branches of the model are different, and the present application is completely different from the cross-consistency working structure recorded in the abstract.

[0018] In summary, the existing technologies are unreasonably designed, cumbersome to operate, and inefficient. When using the same data, the segmentation of image edges is not accurate enough, and most of the technologies are coupled with the model and cannot be generalized. Summary of the invention

[0019] In order to solve the above technical problems, the present invention provides a semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels. While using less data, it ensures the accuracy of image segmentation and reduces the training time cost. It also distinguishes the semi-supervised method from model selection and can be applied to different scenarios.

[0020] A semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels, including:

[0021] Step 1: Collect and preprocess different types of medical imaging data sets to construct training sets and test sets;

[0022] As an example, the types of medical image data sets include: lesion recognition + organ segmentation.

[0023] As an example, the preprocessing is: random cropping, random rotation and noise enhancement.

[0024] Step 2: Selecting a semi-supervised medical image segmentation network;

[0025] The segmentation network adopts: a branch network with a single encoder and two decoders;

[0026] As an example, the branch network of the single encoder and dual decoders is transposed convolution + linear interpolation, where different decoders correspond to different upsampling methods.

[0027] Step 3: According to the characteristics of different medical image data sets, a basic segmentation model is selected for the segmentation network;

[0028] As an example, the selection result is: Unet and / or Vnet network model.

[0029] As a preferred example, in the brain tumor segmentation operation, the Unet network model is selected; in the left atrium segmentation operation, the Vnet network model is used.

[0030] Step 4: Input the data of the training set into the basic segmentation model, ensure the ratio of labeled data to unlabeled data in each batch, set hyperparameters, generate a mask using the soft pseudo-label results between different branch networks, and train the basic segmentation model;

[0031] As an example, the hyperparameters are: temperature coefficient and number of training times.

[0032] As an example, the operation of generating a mask using the soft false label result is obtained by a sharpen function with a temperature coefficient;

[0033] Specifically, the temperature is adjusted with the number of training rounds; the generation of the mask depends on the probability after sharpening, and the confidence regions are calculated on different classification results and then the union is taken to finally obtain the confidence mask.

[0034] Step 5: Input the test set data into the trained basic segmentation model to obtain the inference results, calculate the segmentation index, and evaluate the basic segmentation model.

[0035] Beneficial effects of the present invention:

[0036] The present invention overcomes the problem that in the semi-supervised medical image segmentation scenario, the deep learning segmentation model cannot fully utilize the unlabeled image data, resulting in a decrease in segmentation accuracy, especially in the edge area of ​​the segmentation. It achieves the decoupling of the model and the method, and can be effectively migrated to different segmentation tasks.

[0037] Masks are generated based on soft pseudo labels to guide the dual branches of the model to promote consistency in high uncertainty areas when segmenting unlabeled images.

[0038] The image processing method of the present invention can effectively reduce the annotation cost of 3D medical images and assist doctors in clinical diagnosis while ensuring high segmentation accuracy. BRIEF DESCRIPTION OF THE DRAWINGS

[0039] Figure 1 Schematic diagram of the overall segmentation model of the semi-supervised medical image segmentation method based on soft pseudo labels to generate confidence masks in the present invention.

[0040] Figure 2 The figure is a schematic diagram of calculation of generating mask of the semi-supervised medical image segmentation method based on soft pseudo labels to generate confidence mask of the present invention.

[0041] Figure 3 The present invention provides a comparison of the segmentation result indicators of the semi-supervised medical image segmentation method based on soft pseudo labels to generate confidence masks on brain tumors using 10% or 20% labeled data.

[0042] Figure 4This is a visualized 3D and slice image of the left atrium segmentation of the semi-supervised medical image segmentation method based on soft pseudo labels to generate confidence masks.

[0043] Figure 5 This is a comparison chart of the time required for a single round of training using a brain tumor dataset on a Nvidia 3090 series graphics card using different prior art models of the semi-supervised medical image segmentation method for generating confidence masks based on soft pseudo labels and the model of the present invention (set as PLCL). DETAILED DESCRIPTION

[0044] Below, reference Figures 1 to 5 As shown, a semi-supervised medical image segmentation method for generating confidence masks based on soft pseudo labels includes:

[0045] Step 1: Collect and preprocess different types of medical imaging data sets to construct training sets and test sets;

[0046] As an example, the types of medical image data sets include: lesion recognition + organ segmentation.

[0047] As an example, the preprocessing is: random cropping, random rotation and noise enhancement.

[0048] As an example, when constructing the training set, the labeled data is divided into 10% or 20% ratios for experiments of different ratios; at the same time, the model size is input: brain tumor: 96*96*96, left atrium: 112*112*80.

[0049] Step 2: Selecting a semi-supervised medical image segmentation network;

[0050] The segmentation network adopts: a branch network with a single encoder and two decoders;

[0051] As an example, the branch network of the single encoder and dual decoders is transposed convolution + linear interpolation, where different decoders correspond to different upsampling methods.

[0052] Step 3: According to the characteristics of different medical image data sets, a basic segmentation model is selected for the segmentation network;

[0053] As an example, the selection result is: Unet and / or Vnet network model.

[0054] As a preferred example, in the brain tumor segmentation operation, the Unet network model is selected; in the left atrium segmentation operation, the Vnet network model is used.

[0055] As an example, in the Unet network model, the encoder part includes a Resnet convolution layer; transposed convolution is used in all three layers of downsampling, and the convolution layers of the dual decoders remain consistent.

[0056] As an example, in the Unet network model, the upsampling layers are: transposed convolution and bilinear interpolation, and the corresponding end from the decoder to the encoder includes a basic skip connection.

[0057] Step 4: Input the data of the training set into the basic segmentation model, ensure the ratio of labeled data to unlabeled data in each batch, set hyperparameters, generate a mask using the soft pseudo-label results between different branch networks, and train the basic segmentation model;

[0058] As an example, the hyperparameters are: temperature coefficient and number of training times.

[0059] As an example, the operation of generating a mask using the soft false label result is obtained by a sharpen function with a temperature coefficient;

[0060] Specifically, the temperature is adjusted with the number of training rounds; the generation of the mask depends on the probability after sharpening, and the confidence regions are calculated on different classification results and then the union is taken to finally believe the mask.

[0061] Furthermore, the training set is used as the input of the basic segmentation model to construct the loss function;

[0062] The loss function is composed of the supervised learning loss obtained by labeled data and the consistency loss obtained by unlabeled data;

[0063] As an example, the loss function formula is designed as follows:

[0064]

[0065] in, t max is the maximum number of iterations, λ = 0.5; L seg is the cross entropy loss and Dice loss function composed of the model prediction output and the true label of the labeled data; L in With L mc They refer to the internal consistency constraints obtained from unlabeled data and the mutual consistency constraints of the outputs between branches; L in It is a consistency loss with adaptive weights, which is used to assign smaller weights to image pixels with high uncertainty, making the training smoother and less prone to over-guidance.

[0066] The sharpen function is designed as follows:

[0067]

[0068] Where: p(y|x; θ i ) refers to the input x in the model with parameter θ i The probability output when a 3D probability map is given, p represents the predicted probability of each point; p * (y|x;θ i ) with an asterisk is the output of this probability map after sharpening.

[0069] As an example, when the probability of predicting that an input is a lesion is 0.7, then after being processed by the sharpen function, the sharpened output result is 0.9.

[0070] At this time, the L in and L mc The following results are obtained:

[0071]

[0072] Among them, the E function is a consistent mask calculation method, which is expressed as the sum of the products of the same classification probabilities is greater than the hyperparameter ε. Refers to the output result of each pixel point, and the KL divergence is calculated, w i is the inverse of the divergence index of each pixel, D kl is the sum of the scatter of all pixels; The numerator is the mask multiplied by the mean square error loss of the two branch predictions for the pixel, and the denominator is the sum of each value of the mask separately;

[0073] The hyperparameter ε is set to 0.5 in the present invention, and the mean square error loss in the two regions can be calculated based on the mask; the formula calculation expression of the E function is designed as follows:

[0074]

[0075] Among them, the indicator function in the middle is equivalent to a value of 1 if the internal condition holds, and a value of 0 if it does not hold; so E i It means that the mask value corresponding to each pixel can only be 0 or 1.

[0076] Step 5: Input the test set data into the trained basic segmentation model to obtain the inference results, calculate the segmentation index, and evaluate the basic segmentation model.

[0077] In order to better illustrate the principle design of the present invention, the following is further illustrated by specific experimental data:

[0078] Embodiment 1:

[0079] Reference Figure 1 and Figure 2 As shown, the entire model finally relies on SGD and stochastic gradient descent to back-propagate and update parameters;

[0080] Furthermore, the present invention compares experimental data with the currently popular semi-supervised medical image segmentation method; the evaluation criteria are: Dice, Jaccard, 95HD, ASD, see Figure 3 shown.

[0081] The test data is cropped according to the sliding window to keep it consistent with the training data, input into the model for prediction, and the final probability map is spliced ​​to obtain the final segmentation result;

[0082] In the experiment, using 20% ​​of the labeled data on the brain tumor dataset, we can achieve a dice score of 85.87%, which is higher than other existing methods and close to the result of using all the data for training.

[0083] Visual analysis: On the left atrium data, the experiment shows the final 3D segmentation result and the effect of 2D slicing. It can be found that in the edge area, the analysis effect of the present invention is more delicate and close to GT (true label); the present invention uses the upsampling method of transposed convolution, and the experimental results show that the final results of the two branches are not much different.

[0084] The experimental results further confirm that the combination of soft pseudo label technology and the dual-branch model can effectively improve the semi-supervised medical image segmentation effect; the unique mask generation method based on soft pseudo labels effectively shortens the training time and improves the accuracy of edge segmentation; the present invention can achieve high-accuracy image segmentation while paying more attention to the details of the segmentation boundary, and the proposed confidence mask generation method reduces the training time compared with the latest model.

[0085] Embodiment 2:

[0086] The present invention also includes an electronic device for a semi-supervised medical image segmentation method based on soft pseudo labels to generate confidence masks. The electronic device includes a memory and a processor. The memory stores a computer program. When the processor executes the computer program, the semi-supervised medical image segmentation method based on soft pseudo labels to generate confidence masks can be implemented.

[0087] Embodiment 3:

[0088] The present invention also includes a computer-readable storage medium having a computer program stored thereon. When the computer program is executed by a processor, a semi-supervised medical image segmentation method for generating a confidence mask based on soft pseudo labels can be implemented.

[0089] Through the description of the above implementation methods, technicians in this field can clearly understand that the above-mentioned implementation methods can be implemented by means of software plus a necessary general hardware platform, and of course can also be implemented by hardware, but in many cases the former is a better implementation method; based on such an understanding, the technical solution of the present application is essentially or the part that contributes to the prior art can be embodied in the form of a software product, and the computer software product is stored in a storage medium (such as ROM / RAM, disk, CD), including a number of instructions for enabling a terminal (which can be a mobile phone, computer, server, air conditioner, or network device, etc.) to execute the methods described in each embodiment of the present application.

[0090] The embodiments of the present application are described above in conjunction with the accompanying drawings, but the present application is not limited to the above-mentioned specific implementation methods. The above-mentioned specific implementation methods are merely illustrative and not restrictive. Under the guidance of the present application, ordinary technicians in this field can also make many forms without departing from the purpose of the present application and the scope of protection of the claims, all of which are within the protection of the present application.

[0091] Those of ordinary skill in the art will appreciate that the units and algorithm steps of each example described in conjunction with the embodiments disclosed in the present application can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professional and technical personnel can use different methods to implement the described functions for each specific application, but such implementation should not be considered to be beyond the scope of this application.

[0092] Those skilled in the art can clearly understand that, for the convenience and brevity of description, the specific working processes of the systems, devices and units described above can refer to the corresponding processes in the aforementioned method embodiments and will not be repeated here.

[0093] In the embodiments provided in the present application, it should be understood that the disclosed devices and methods can be implemented in other ways; for example, the device embodiments described above are only schematic, and the division of the units is only a logical function division. There may be other division methods in actual implementation, such as multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of devices or units, which can be electrical, mechanical or other forms.

[0094] In addition, each functional unit in each embodiment of the present application may be integrated into one processing unit, or each unit may exist physically separately, or two or more units may be integrated into one unit.

[0095] If the functions are implemented in the form of software functional units and sold or used as independent products, they can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present application, or the part that contributes to the prior art or the part of the technical solution, can be embodied in the form of a software product, which is stored in a storage medium and includes several instructions for a computer device (which can be a personal computer, server, or network device, etc.) to perform all or part of the steps of the methods described in each embodiment of the present application. The aforementioned storage medium includes: various media that can store program codes, such as USB flash drives, mobile hard drives, ROM, RAM, magnetic disks, or optical disks.

[0096] The above are only preferred embodiments of the present invention. It should be understood that the description of the above embodiments is only used to help understand the method and core ideas of the present invention, and is not used to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, etc. made within the ideas and principles of the present invention should be included in the scope of protection of the present invention.

Claims

1. A semi-supervised medical image segmentation method based on generating confidence masks from soft pseudo labels, characterized in that: include: Step 1: Collect and preprocess different types of medical imaging data sets to construct training sets and test sets; Step 2: Selecting a semi-supervised medical image segmentation network; The segmentation network adopts: a branch network with a single encoder and two decoders; Step 3: According to the characteristics of different medical image data sets, a basic segmentation model is selected for the segmentation network; Step 4: Input the data of the training set into the basic segmentation model, ensure the ratio of labeled data to unlabeled data in each batch, set hyperparameters, generate masks using the soft pseudo-label results between different branch networks, and train the basic segmentation model; Step 5: Input the test set data into the trained basic segmentation model to obtain the inference results, calculate the segmentation index, and evaluate the basic segmentation model.

2. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 1, characterized in that: The types of medical image data sets include: lesion recognition + organ segmentation.

3. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 1, characterized in that: The preprocessing includes: random cropping, random rotation and noise enhancement.

4. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 1, characterized in that: When constructing the training set, the labeled data is divided into 10% or 20% ratios for experiments with different ratios; at the same time, the model size is input: brain tumor: 96*96*96, left atrium: 112*112*80.

5. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 1, characterized in that: The branch network of the single encoder and dual decoders is transposed convolution + linear interpolation, wherein different decoders correspond to different upsampling methods.

6. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 1, characterized in that: The selection results are: Unet and / or Vnet network models; In the brain tumor segmentation operation, the Unet network model is selected; in the left atrium segmentation operation, the Vnet network model is used.

7. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 6, characterized in that: In the Unet network model, the encoder part includes a Resnet convolution layer; when encountering 3 layers of downsampling, transposed convolution is used, and the convolution layer of the dual decoder remains consistent; In the Unet network model, the upsampling layers are: transposed convolution and bilinear interpolation, and the corresponding end from the decoder to the encoder includes a basic skip connection.

8. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 1, characterized in that: The hyper parameters are: temperature coefficient and number of training times; the operation of generating a mask using the soft pseudo label result is obtained by a sharpen function with a temperature coefficient; Specifically, the temperature is adjusted with the number of training rounds; the generation of the mask depends on the probability after sharpening, and the confidence region is calculated on different classification results and then the union is taken to finally obtain the confidence mask; Use the training set as the input of the basic segmentation model and construct the loss function; The loss function is composed of the supervised learning loss obtained by labeled data and the consistency loss obtained by unlabeled data.

9. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 8, characterized in that: The loss function formula is designed as follows: in, t max is the maximum number of iterations, λ = 0.5; L seg is the cross entropy loss and Dice loss function composed of the model prediction output and the true label of the labeled data; L in With L mc They refer to the internal consistency constraints obtained from unlabeled data and the mutual consistency constraints of the outputs between branches; L in It is a consistency loss with adaptive weights, which is used to assign smaller weights to image pixels with high uncertainty, making the training smoother and less prone to over-guidance. Design the sharpen function as follows: Where: p(y|x; θ i ) refers to the input x in the model with parameter θ i The probability output when a 3D probability map is given, p represents the predicted probability of each point; p * (y|x;θ i ) with an asterisk is the output of this probability map after sharpening; At this time, the L in and L mc The following results are obtained: Among them, the E function is a consistent mask calculation method, which is expressed as the sum of the products of the same classification probabilities is greater than the hyperparameter ε; Refers to the output result of each pixel point, and the KL divergence is calculated, w i is the inverse of the divergence index of each pixel, D kl is the sum of the scatter of all pixels; The numerator is the mask multiplied by the mean square error loss of the two branch predictions for the pixel, and the denominator is the sum of each value of the mask separately.

10. The semi-supervised medical image segmentation method based on generating confidence masks based on soft pseudo labels according to claim 9, characterized in that: The confidence hyperparameter ε is set, which is set to 0.5 in the present invention. Based on the mask, the mean square error loss in the two regions can be calculated; the formula calculation expression of the E function is designed as follows: Among them, the indicator function in the middle is equivalent to a value of 1 if the internal condition holds, and a value of 0 if it does not hold; so E i It means that the mask value corresponding to each pixel can only be 0 or 1.

Images