Semi-supervised medical image segmentation method based on extra supervision

Abstract

The invention belongs to the field of medical image processing, and particularly relates to a semi-supervised medical image segmentation method based on extra supervision. The method comprises the following steps: acquiring a medical image data set and preprocessing the medical image data set to obtain a preprocessed medical image; processing the medical image by adopting a two-branch medical image segmentation network and calculating an expected total loss; carrying out pre-training according to the previous step, judging whether the image is in a pre-training round or not, directly carrying out the next round of pre-training if the image is in the pre-training round, and otherwise, constructing a new image pair; a two-branch medical image segmentation network is adopted to carry out additional supervision learning on the new image pair, and additional supervision loss is calculated; calculating the total loss of the model according to the expected total loss and the additional supervision loss; adjusting model parameters according to the total loss of the model to obtain a trained medical image segmentation model; according to the method, related knowledge can be better extracted from the label-free image, the influence caused by noise is avoided, and the accuracy of image segmentation is improved.

Description

Technical Field

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

[0002] Medical images play a vital role in the medical diagnosis and treatment process for medical staff. They can reflect the internal structure of the human body, provide detailed information about lesions and injuries, and help doctors diagnose and analyze diseases. Generally speaking, a large amount of labeled data and existing image segmentation algorithms can achieve good segmentation results; but in the field of medical image processing, since labeling data is usually very expensive and time-consuming, semi-supervised medical image segmentation using only a small amount of labeled data and a large amount of unlabeled data has important practical significance and application value. It reduces the demand for manually labeled data by introducing domain knowledge, and uses unlabeled data for learning to improve the generalization ability and accuracy of the model. At the same time, other technologies such as data enhancement and transfer learning can be combined to further improve the performance of the model. Semi-supervised medical image segmentation has broad application prospects in the field of medical image processing, and can be applied to disease diagnosis, surgical navigation, medical research and many other aspects.

[0003] The current mainstream semi-supervised medical image segmentation mainly uses deep learning models to learn labeled and unlabeled images. The purpose is to extract as much useful information as possible from unlabeled images, while avoiding the interference caused by the model's incorrect predictions. Currently, a series of technologies such as pseudo-labeling, consistency learning, and transfer learning are often used to further improve the generalization ability and accuracy of the model. However, the existing medical image segmentation models still have the problem that the generalization ability of semi-supervised methods is not strong enough, and it is difficult to avoid the noise problem in unlabeled image prediction. Summary of the invention

[0004] In view of the shortcomings of the prior art, the present invention proposes a semi-supervised medical image segmentation method based on additional supervision, the method comprising: obtaining a medical image to be predicted and inputting it into a trained medical image segmentation model to obtain a medical image segmentation result;

[0005] The training process of the medical image segmentation model includes:

[0006] Step S1: obtaining a medical image dataset and preprocessing it to obtain preprocessed medical images; wherein the medical images include labeled images and unlabeled images;

[0007] Step S2: using a two-branch medical image segmentation network to process the medical image and calculate the expected total loss;

[0008] Step S3: Perform pre-training according to step S2, and determine whether it is in the pre-training round. If it is in the pre-training round, directly proceed to the next round of pre-training; otherwise, construct a new image pair;

[0009] Step S4: Use a two-branch medical image segmentation network to perform additional supervision learning on the new image pair and calculate the additional supervision loss;

[0010] Step S5: Calculate the total loss of the model according to the expected total loss and the additional supervision loss; adjust the model parameters according to the total loss of the model to obtain a trained medical image segmentation model.

[0011] Preferably, step S2 comprises:

[0012] Use a two-branch medical image segmentation network to perform supervised learning and consistency learning on labeled images, and calculate the supervised learning and consistency loss of labeled images;

[0013] A two-branch medical image segmentation network is used to perform consistency learning on unlabeled images and calculate the consistency loss of unlabeled images;

[0014] The expected total loss is calculated based on supervised learning for labeled images, consistency loss for labeled images, and consistency loss for unlabeled images.

[0015] Furthermore, the formula for the expected total loss is:

[0016]

[0017] Among them, L t represents the expected total loss, represents the total loss of labeled images for the i-th network branch, represents the total loss of unlabeled images for the ith network branch, represents the supervised loss of labeled images for the i-th network branch, L css represents the consistency loss of the two network branches for labeled images, λ j represents the consistency loss weight during the j-th round of training, L csu represents the consistency loss of the two network branches on unlabeled images.

[0018] Furthermore, the consistency loss weight during the j-th round of training is:

[0019]

[0020] Among them, epoch represents the total number of training rounds.

[0021] Preferably, in step S3, the process of constructing a new image pair includes:

[0022] S31: Generate a mask according to the prediction result of the unlabeled image in this round and the prediction result of the unlabeled image in the previous round;

[0023] S32: Calculate a new image based on the mask, the unlabeled image of this round, and the labeled image of this round;

[0024] S33: Calculate a new pseudo label based on the mask, the prediction result of the unlabeled image in this round, and the label of the labeled image;

[0025] S34: Take the new image and the new pseudo label as a new image pair.

[0026] Furthermore, the formula for calculating the new image and the new pseudo label is:

[0027] I new =I u (x,y)*(1-M(x,y))+I l (M(x,y))

[0028] l new = l u (x,y)*(1-M(x,y))+l l (M(x,y))

[0029] Among them, I new represents the new image, l new represents the new pseudo label, I u represents an unlabeled image, M(x,y) represents a mask, and (x,y) represents the pixel coordinates in the image; I l Indicates a labeled image, l u Represents the prediction result of the unlabeled image, l l Represents the label of the labeled image.

[0030] Preferably, the formula for calculating the additional supervision loss is:

[0031]

[0032]

[0033] Among them, L e represents the additional supervision loss, represents the additional supervision loss of the i-th network branch, Table 2 New images of the second network branch, represents the new pseudo-label for the first network branch, A new image representing the first network branch, represents the new pseudo-label of the second network branch, and f(·) represents the output after passing through the network branch.

[0034] Preferably, the total model loss is the sum of the expected total loss and the additional supervision loss.

[0035] The beneficial effects of the present invention are:

[0036] The present invention designs an additionally supervised semi-supervised medical image segmentation model to learn labeled image pairs, and can also use unlabeled data to perform additional supervision on the learning process of labeled data to prevent overfitting. When using new image pairs for training, additional supervision is performed using labeled images to effectively prevent overfitting and prevent the model from learning the noise in the pseudo-labels. At the same time, a parameter λ is designed. j To control the model's over-learning of new image pairs and prevent the model from learning the noise in the pseudo-labels. The present invention can better extract relevant knowledge from unlabeled images, avoid the influence of noise, and improve the accuracy of image segmentation. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] Figure 1 is a training flow chart of the medical image segmentation model of the present invention;

[0038] Figure 2 is a training structure diagram of the medical image segmentation model in the present invention;

[0039] Figure 3 Schematic diagram of the new image pair construction process in the present invention. DETAILED DESCRIPTION

[0040] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0041] This paper proposes a semi-supervised medical image segmentation method based on additional supervision, such as Figure 1 , Figure 2 As shown, the method includes the following contents:

[0042] Obtain the medical image to be predicted and input it into the trained medical image segmentation model to obtain the medical image segmentation result.

[0043] The training process of the medical image segmentation model includes:

[0044] Step S1: Obtain a medical image dataset and preprocess it to obtain preprocessed medical images; wherein the medical images include labeled images and unlabeled images.

[0045] Obtain a medical image dataset, which includes labeled images and unlabeled images. Preprocess the images, specifically:

[0046] The labeled image is resized to 256×256×3, and the corresponding label size is resized to 256×256×1. The unlabeled image is resized to 256×256×3.

[0047] Then, two different image enhancement strategies were used to process the images. The first one was to randomly select two of the four preprocessing methods, namely brightness adjustment, contrast adjustment, Gaussian blur, and random rotation, which would not reduce the image information or have little impact on the existing information. This process is called weak enhancement. The second one was to randomly select one of the three methods, namely random cropping, random elimination, and random exchange, which might reduce the information or greatly change the image information, and then extract one from the weak processing enhancement method for combination processing. This process is called strong enhancement. The images were processed according to the selected preprocessing method to obtain the weakly enhanced image and the strongly enhanced image corresponding to each image.

[0048] Step S2: Use a two-branch medical image segmentation network to process the medical image and calculate the expected total loss.

[0049] The present invention designs a semi-supervised medical image segmentation model (two-branch medical image segmentation network) with multiple loss functions that can simultaneously learn labeled and unlabeled images. The model shares two network branches with the same structure but different parameter initializations. Preferably, the branch structure often uses a U-Net network, which can also be replaced by other deep learning networks.

[0050] For labeled images:

[0051] The weakly enhanced image and the strongly enhanced image corresponding to the labeled image are respectively input into the two network branches of the medical image segmentation network for supervised learning and consistency learning. The total loss of network branch i in the jth round of training is:

[0052]

[0053] in, represents the total loss of labeled images for the i-th network branch; represents the supervised loss of the labeled image of the i-th network branch, which represents the difference between the predicted result and the label of the labeled image; L css represents the consistency loss of the two network branches for labeled images, which represents the difference between the images output by the two network branches; λ j Represents the consistency loss weight during the j-th round of training.

[0054] For unlabeled images:

[0055] The weakly enhanced image and the strongly enhanced image corresponding to the unlabeled image are respectively input into the two network branches of the medical image segmentation network for consistency learning. The total loss of network branch i in the jth round of training is:

[0056]

[0057] in, represents the total loss of unlabeled images for the i-th network branch, L csu represents the consistency loss of the two network branches on unlabeled images.

[0058] Parameter λ j The model can be controlled to over-learn new image pairs, preventing the model from learning the noise in the pseudo-labels. The consistency loss weight designed by the present invention during the j-th round of training is:

[0059]

[0060] Among them, epoch represents the total number of training rounds.

[0061] The expected total loss is calculated based on supervised learning of labeled images, consistency loss of labeled images, and consistency loss of unlabeled images:

[0062]

[0063] Step S3: Perform pre-training according to step S2, and determine whether it is in the pre-training round. If it is in the pre-training round, directly proceed to the next round of pre-training; otherwise, construct a new image pair.

[0064] Preset the pre-training rounds and perform pre-training according to step S2, and perform formal training after the pre-training according to step S2 is completed. During formal training, after completing step S2, construct new image pairs to continue the subsequent training process, such as Figure 3 As shown, specifically:

[0065] S31: Generate a mask according to the prediction result of the unlabeled image in this round and the prediction result of the unlabeled image in the previous round.

[0066] The prediction results of the unlabeled images obtained in the previous round of training stored in the memory bank and the prediction results of the unlabeled images in this round are used to determine the unstable area, and a mask is generated to represent the unstable area. The formula is as follows:

[0067] M(x,y)={(x,y)|P1(x,y)≠P2(x,y)}x∈X,y∈Y

[0068] Among them, M(x,y) is the mask of the unstable area, P1(x,y) and P2(x,y) represent the prediction results of the unlabeled image in the previous round and the prediction results of the unlabeled image in this round extracted from the memory bank, respectively, X and Y represent the width and height of the image, respectively, and (x,y) is the pixel coordinate of the image.

[0069] S32: Calculate a new image based on the mask, the unlabeled image of this round, and the labeled image of this round.

[0070] Use the corresponding position image in the labeled image to replace the unstable area in the unlabeled image to reconstruct the new image, which is expressed as:

[0071] I new =I u (x,y)*(1-M(x,y))+I l *(M(x,y))

[0072] Among them, I new represents the new image, I u represents an unlabeled image, I l Indicates a labeled image.

[0073] S33: A new pseudo label is calculated based on the mask, the prediction result of the unlabeled image in this round, and the label of the labeled image.

[0074] Use the part of the label corresponding to the labeled image to replace the corresponding area in the unlabeled image prediction result to reconstruct the new pseudo label, which is expressed as:

[0075] l new = l u (x,y)*(1-M(x,y))+l l *(M(x,y))

[0076] Among them, l new represents the new pseudo label, l n Represents the prediction result of the unlabeled image, l l Represents the label of the labeled image.

[0077] S34: Take the new image and the new pseudo-label as a new image pair P new .

[0078] P new =(I new ,l new )

[0079] Step S4: Use a two-branch medical image segmentation network to perform additional supervised learning on the new image pair and calculate the additional supervision loss.

[0080] Both network branches construct their own new image pairs, namely The new image pairs are exchanged according to the network branches for additional supervision training. The additional supervision loss of each network branch is as follows:

[0081]

[0082] in, represents the additional supervision loss of the first network branch, represents the additional supervision loss of the second network branch; Table 2 New images of the second network branch, represents the new pseudo-label for the first network branch, A new image representing the first network branch, represents the new pseudo-label of the second network branch, and f(·) represents the output after passing through the network branch.

[0083] The additional supervision loss of the two-branch medical image segmentation network is:

[0084]

[0085] Among them, L e represents the additional supervision loss, represents the additional supervision loss of the i-th network branch.

[0086] Step S5: Calculate the total loss of the model according to the expected total loss and the additional supervision loss; adjust the model parameters according to the total loss of the model to obtain a trained medical image segmentation model.

[0087] Preferably, the total model loss L is the sum of the expected total loss and the additional supervision loss:

[0088] L=L t +L e

[0089] The model parameters are adjusted according to the total model loss. When the preset total training rounds are reached or the total model loss converges, the training is stopped, the model parameters are saved, and a trained medical image segmentation model is obtained.

[0090] After obtaining the trained model, the image to be predicted is input into the two network branches in the trained medical image segmentation model, and the prediction results of the two network branches are obtained respectively. The parts with the same prediction results are taken as the final prediction results, and the parts with different prediction results are selected as the final prediction results with higher prediction confidence.

[0091] In summary, the present invention uses data obtained through different preprocessing methods during pre-training to enable the sub-network to learn the features in the labeled image from different angles. At the same time, an additionally supervised semi-supervised medical image segmentation model is designed to learn the labeled image pairs while also using unlabeled data to perform additional supervision on the learning process of the labeled data to prevent overfitting. In the process of constructing new image pairs, the parts of the unlabeled images with different results in the current round and the previous round are regarded as unstable areas. These unstable areas often contain a lot of noise, so these parts are directly replaced with the corresponding positions in the labeled images to construct new image pairs as pseudo-labels, which not only removes the influence of high-noise areas on the network learning effect, but also increases the diversity of the data set in this way. When using new image pairs for training, labeled images are used for additional supervision to prevent overfitting and prevent the model from learning the noise in the pseudo-labels. At the same time, a parameter λ is designed. j To control the model's over-learning of new image pairs and prevent the model from learning the noise in the pseudo-labels. The present invention can better extract relevant knowledge from unlabeled images, avoid the influence of noise, and improve the accuracy of image segmentation.

[0092] The above embodiments further illustrate the purpose, technical solutions and advantages of the present invention in detail. It should be understood that the above embodiments are only preferred implementation modes of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made to the present invention within the spirit and principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A semi-supervised medical image segmentation method based on additional supervision, characterized in that: include: Obtain a medical image to be predicted and input it into a trained medical image segmentation model to obtain a medical image segmentation result; The training process of the medical image segmentation model includes: Step S1: obtaining a medical image dataset and preprocessing it to obtain preprocessed medical images; wherein the medical images include labeled images and unlabeled images; Step S2: using a two-branch medical image segmentation network to process the medical image and calculate the expected total loss; Step S3: Perform pre-training according to step S2, and determine whether it is in the pre-training round. If it is in the pre-training round, directly proceed to the next round of pre-training; otherwise, construct a new image pair; Step S4: Use a two-branch medical image segmentation network to perform additional supervision learning on the new image pair and calculate the additional supervision loss; Step S5: Calculate the total loss of the model according to the expected total loss and the additional supervision loss; adjust the model parameters according to the total loss of the model to obtain a trained medical image segmentation model.

2. The method for semi-supervised medical image segmentation based on additional supervision according to claim 1, characterized in that: The step S2 comprises: Use a two-branch medical image segmentation network to perform supervised learning and consistency learning on labeled images, and calculate the supervised learning and consistency loss of labeled images; A two-branch medical image segmentation network is used to perform consistency learning on unlabeled images and calculate the consistency loss of unlabeled images; The expected total loss is calculated based on supervised learning for labeled images, consistency loss for labeled images, and consistency loss for unlabeled images.

3. The method for semi-supervised medical image segmentation based on additional supervision according to claim 2, characterized in that: The formula for the expected total loss is: Among them, L t represents the expected total loss, represents the total loss of labeled images for the i-th network branch, represents the total loss of unlabeled images for the ith network branch, represents the supervised loss of labeled images for the i-th network branch, L css represents the consistency loss of the two network branches for labeled images, λ j represents the consistency loss weight during the j-th round of training, L csu represents the consistency loss of the two network branches on unlabeled images.

4. The method for semi-supervised medical image segmentation based on additional supervision according to claim 3, characterized in that: The consistency loss weight during the jth round of training is: Among them, epoch represents the total number of training rounds.

5. The method for semi-supervised medical image segmentation based on additional supervision according to claim 1, characterized in that: In step S3, the process of constructing a new image pair includes: S31: Generate a mask according to the prediction result of the unlabeled image in this round and the prediction result of the unlabeled image in the previous round; S32: Calculate a new image based on the mask, the unlabeled image of this round, and the labeled image of this round; S33: Calculate a new pseudo label based on the mask, the prediction result of the unlabeled image in this round, and the label of the labeled image; S34: Take the new image and the new pseudo label as a new image pair.

6. The method for semi-supervised medical image segmentation based on additional supervision according to claim 5, characterized in that: The formula for calculating the new image and the new pseudo label is: I new =I u (x,y)*(1-M(x,y))+I l (M(x,y)) l new =l u (x,y)*(1-M(x,y))+l l (M(x,y)) Among them, I new represents the new image, l new represents the new pseudo label, I u represents an unlabeled image, M(x,y) represents a mask, and (x,y) represents the pixel coordinates in the image; I l Indicates a labeled image, l u Represents the prediction result of the unlabeled image, l l Represents the label of the labeled image.

7. The method for semi-supervised medical image segmentation based on additional supervision according to claim 1, characterized in that: The formula for calculating the additional supervision loss is: Among them, L e represents the additional supervision loss, represents the additional supervision loss of the i-th network branch, Table 2 New images of the second network branch, represents the new pseudo-label for the first network branch, A new image representing the first network branch, represents the new pseudo-label of the second network branch, and f(·) represents the output after passing through the network branch.

8. The method for semi-supervised medical image segmentation based on additional supervision according to claim 1, characterized in that: The total model loss is the sum of the expected total loss and the additional supervision loss.

Images