A semi-supervised medical image segmentation method based on perturbation co-training

Abstract

The present invention discloses a semi-supervised medical image segmentation method based on perturbation collaborative training, which uses a small amount of labeled medical images while mining the information in a large amount of unlabeled medical images to improve the performance of the segmentation model. For labeled data, the labels are used to supervise the prediction results of the two sub-networks; for unlabeled data, the two sub-networks generate pseudo-labels based on their respective prediction results, and then use the pseudo-labels to supervise the prediction results of the other sub-network. In the process of generating pseudo-labels, the sub-network perturbs the features extracted by the encoder based on the feature information of the other sub-network, thereby increasing the diversity of pseudo-labels and thus improving the effect of collaborative training. The method described in the present invention can be applied to medical image segmentation tasks of various parts and has good universality and versatility.

Description

Technical Field

[0001] The present invention relates to medical image analysis, and in particular to a semi-supervised medical image segmentation method based on perturbation collaborative training. Background Art

[0002] Medical image analysis is a key application scenario for semantic segmentation. Segmenting target organs or lesion locations from medical images can help doctors quickly identify targets and develop treatment plans. Training accurate semantic segmentation models typically requires a large number of medical images with pixel-level or voxel-level annotations, and labeling medical images requires a significant amount of time and effort by professional doctors. This high labor cost is prohibitive in real-world applications and often becomes a significant factor limiting the performance of semantic segmentation models. In contrast, hospital databases contain a large number of unlabeled medical images. Therefore, semi-supervised learning can be used to improve the performance of semantic segmentation models by leveraging a small amount of labeled data while simultaneously mining information from a large amount of unlabeled data.

[0003] Most semi-supervised semantic segmentation methods currently used in medical image analysis rely on consistency constraints, which ensure that the model produces similar predictions for perturbed data. One approach, known as collaborative training, uses two or more sub-networks to make predictions on the same data and uses their predictions to supervise each other, thereby achieving consistency across different features of the same data.

[0004] Traditional co-training relies on the crucial assumption of conditional independence, which states that different views of the same data should be conditionally independent. In the era of deep learning, this translates to diversity in different sub-networks. However, most current semi-supervised medical image segmentation algorithms based on co-training fail to enhance sub-network diversity, potentially leading to similarities between sub-networks and thus impacting co-training model performance. Summary of the Invention

[0005] Purpose of the invention: The purpose of the present invention is to provide a semi-supervised medical image segmentation method based on perturbation collaborative training, so that in the process of generating pseudo labels, the feature information of another sub-network is introduced to perturb the features of this sub-network, thereby increasing the diversity of pseudo labels.

[0006] Technical solution: The semi-supervised medical image segmentation method based on perturbation collaborative training described in the present invention includes the following steps:

[0007] (1) Inputting a training dataset, wherein the training dataset contains a large amount of unlabeled data and a small amount of labeled data; randomly initializing the parameters of two encoder-decoder segmentation networks with the same structure;

[0008] (2) Sample labeled medical images and their corresponding labels (Xl ,Y l ) and unlabeled medical images X u , where X u , H, W, D represent the height, width and depth of the image respectively, Y l The value range of the elements is {0,1}, 0 represents background voxels and 1 represents foreground voxels;

[0009] (3) preprocessing and data augmentation of the input data, wherein the preprocessing includes normalizing the input image to the interval [0, 1] and randomly cropping the input image and label;

[0010] (4) For labeled data X l , input it into the two sub-networks, and obtain the prediction results P of the two sub-networks l,1 ,P l,2 ∈[0,1] H×W×D , based on the prediction result P l,1 ,P l,2 and label Y l Calculate Dice loss and binary cross entropy loss and sum them as the training loss on labeled data

[0011] (5) The unlabeled data X u Input into the two sub-networks, and during the calculation process, the encoder features are perturbed based on the feature information between the sub-networks to generate pseudo labels

[0012] (6) Again, the unlabeled data X u It is input into the two sub-networks to obtain the prediction results P of the two sub-networks. u,1 ,P u,2 ∈[0,1] H×W×D , based on the prediction result P u,1 ,P u,2 and pseudo labels Calculate Dice loss and binary cross entropy loss and sum them as the training loss on unlabeled data

[0013] (7) Training loss based on labeled data and the training loss on unlabeled data Calculate the overall training loss Update the segmentation network;

[0014] (8) Determine whether the maximum number of iterations has been reached, otherwise return to step (2);

[0015] (9) Stop training and use the trained model to make predictions for the input image X.

[0016] The training loss on the labeled data described in step (4) is defined as:

[0017]

[0018] in represents the binary cross entropy loss function, Represents the Dice loss function.

[0019] The unlabeled loss in step (7) is defined as follows:

[0020]

[0021] The overall training loss described in step (7) The definition is as follows:

[0022]

[0023] where β t Control the balance between labeled loss and unlabeled loss; gradually increase β in the form of an exponential function in the iterative round t t =β0exp(-5*(1-t / T max ) 2 ), where β0 is the control and Balanced hyperparameter, T max is the maximum number of iterations.

[0024] The step (5) is specifically as follows:

[0025] (5.1) For unlabeled data X u , the encoders of the two sub-networks each extract features Where j represents the feature extracted by the encoder of the jth layer;

[0026] (5.2) Randomly sample j from {1,2,…M} and randomly sample λ from Beta(α,α), where Beta represents the Beta distribution and α is a hyperparameter used to control the shape of the Beta distribution;

[0027] (5.3) Based on the following formula and Make a perturbation:

[0028] (5.4) The perturbed feature F u,1 and F u,2 Send it to the decoder to get the prediction result of the sub-network And according to and Generate pseudo labels and

[0029] The step (9) is specifically as follows:

[0030] (9.1) Preprocessing the input image, wherein the preprocessing includes resampling and registering the input image and normalizing the input image to the interval [0, 1] to obtain an optimized input image;

[0031] (9.2) using a sliding window strategy to cut out a prediction sub-block of the same size as the training image from the optimized input image;

[0032] (9.3) Input the predicted sub-blocks into the trained model and obtain the predicted segmentation results of each predicted sub-block through forward calculation. Here, the prediction results of a single sub-network or the integration of the prediction results of two sub-networks are used;

[0033] (9.4) Use the Gaussian weighting strategy to fuse the model prediction results of each prediction sub-block to obtain the prediction result of the input image X.

[0034] A computer storage medium stores a computer program, which, when executed by a processor, implements the aforementioned semi-supervised medical image segmentation method based on perturbation collaborative training.

[0035] A computer device includes a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the semi-supervised medical image segmentation method based on perturbation collaborative training is implemented.

[0036] Beneficial effects: Compared with the prior art, the present invention has the following advantages:

[0037] 1. The method of the present invention can improve the performance of the model by utilizing the labeled data while mining the information in the unlabeled data based on the idea of ​​collaborative training when there is less labeled data and more unlabeled data.

[0038] 2. The method of the present invention utilizes the feature information between sub-networks to perturb the features during the collaborative training process to generate pseudo labels, thereby enhancing the diversity of pseudo labels and further improving the segmentation accuracy of the collaborative training model;

[0039] 3. The method described in the present invention can be applied to medical image segmentation tasks of various parts and has good universality and versatility. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] Figure 1 A schematic diagram of the network structure of the method of the present invention;

[0041] Figure 2 A flowchart of the steps of the method of the present invention;

[0042] Figure 3 This is a flowchart of the pseudo-label generation process of the method of the present invention;

[0043] Figure 4 The figure is a flowchart of the prediction process of the method of the present invention. DETAILED DESCRIPTION

[0044] The technical solution of the present invention will be further described below with reference to the accompanying drawings.

[0045] like Figure 1 、 2 As shown, a semi-supervised medical image segmentation method based on perturbation collaborative training includes the following steps:

[0046] (1) Input a training dataset containing a large number of unlabeled 3D medical images and a small number of labeled 3D medical images, and randomly initialize the parameters of two encoder-decoder segmentation sub-networks with the same structure;

[0047] (2) Randomly sample a small batch of 3D medical images from the training set, with a batch size of 8, which includes 4 labeled data and 4 unlabeled data;

[0048] (3) Preprocessing and data augmentation of the input data, including normalizing the input image to the interval [0, 1] and using random cropping to crop fixed-size sub-blocks from the normalized image for training;

[0049] (4) Calculate the prediction results of the two sub-networks for the labeled data, and calculate the training loss on the labeled data based on the true labels It is defined as follows:

[0050]

[0051] Among them, P l,i represents the prediction result of the i-th sub-network for the labeled data, Y l Indicates the label of the labeled data, l BCE and l Dice Represent binary cross entropy loss and Dice loss respectively;

[0052] (5) Figure 3As shown in the figure, for unlabeled data, the sub-network uses the feature information of the other sub-network encoder to perturb the features of its own encoder during the calculation process, obtains the prediction results and generates pseudo labels, thereby increasing the diversity of pseudo labels and better meeting the diversity requirements of collaborative training;

[0053] (5.1) The encoders of the two sub-networks each extract features from the unlabeled data Where M represents the encoder level of the feature. The deeper the feature, the richer the semantic information, but the smaller the resolution of the feature.

[0054] (5.2) Randomly sample j uniformly from {1,2,…,M} and sample λ according to Beta(α,α), where Beta is a probability distribution and α controls the shape of its probability density function, which is set to 1.

[0055] (5.3) According to the formula and Make a perturbation:

[0056] (5.4) The perturbed feature F u,2 ,F u,2 The two sub-networks are fed into the decoder to obtain the prediction results and generate pseudo labels.

[0057] (6) Calculate the prediction results of the two sub-networks on the unlabeled data again, without applying feature perturbations, and calculate the training loss on the unlabeled data based on the pseudo labels generated by the other sub-network It is defined as follows:

[0058]

[0059] Among them, P u,i represents the prediction result of the i-th sub-network for unlabeled data, represents the pseudo label generated by the i-th sub-network, l BCE and l Dice Represent binary cross entropy loss and Dice loss respectively;

[0060] (7) Calculate the overall loss based on the labeled loss and the unlabeled loss Stochastic gradient descent with Nesterov momentum is used as the optimization algorithm to update the parameters of the two sub-networks, and the learning rate is adjusted using the cosine decay strategy. It is expressed as follows:

[0061]

[0062] where β tControl the balance between labeled loss and unlabeled loss, and gradually increase β in the form of an exponential function in the iterative round t t =β0 exp(-5*(1-t / T max ) 2 ), where β0 is the control and The balanced hyperparameter is set to 1, T max is the maximum number of iterations, which is 6000;

[0063] (8) Then determine whether the maximum number of training iterations, 6000, has been reached. If not, return to step (2);

[0064] (9) Figure 4 As shown in the figure, if the maximum number of iterations is reached, the training stops and enters the prediction process. For a given input image, the network's prediction result is given.

[0065] (9.1) Preprocess the input image, including resampling, registration, and normalization to the interval [0, 1]. Registration here refers to resampling the input image and label data to the spatial spacing of the three-dimensional medical images in the training set using the Spline interpolation algorithm and the nearest neighbor interpolation algorithm.

[0066] (9.2) A sliding window strategy is used to crop a prediction sub-block with the same size as the training image from the optimized input image.

[0067] (9.3) The prediction sub-blocks are input into the two sub-networks, and the prediction results of each prediction sub-block are obtained through forward calculation. Here, the prediction results of a single sub-network can be used, or the integration of the prediction results of the two sub-networks can be used.

[0068] (9.4) The Gaussian weighting strategy is used to fuse the model prediction results of each prediction sub-block to obtain the prediction result of the input image. This can alleviate the problem of inaccurate prediction of the edge of the prediction sub-block and make the final prediction result more accurate and smooth.

[0069] This example conducts experiments on the NIH pancreas dataset, which contains 80 abdominal CT enhanced scans. The labels in the dataset are provided by senior radiologists. Only some of the labels are used, and the other labels are discarded to simulate a semi-supervised scenario. The experiment here only uses pancreas segmentation as an example to verify the effectiveness of the present invention, but the application scope of the present invention is not limited to pancreas segmentation. It can also be used for the segmentation of other organ regions or lesions. The experiment compares the best existing method, which is the "Semi-supervised Semantic Segmentation Method Using Cross Pseudo-Supervision" proposed in a paper published at CVPR, the top conference in the field of computer vision in 2020. The Dice similarity coefficient is used as the evaluation metric in the experiment to compare the segmentation model performance of the method proposed by the present invention and the best existing method in the 10% labeled data and 20% labeled data scenarios. The experimental results show that the present invention can effectively improve the generalization performance of the model in the semi-supervised scenario. The comparison with the best method also illustrates the effectiveness of the present invention in perturbing features during the pseudo-label generation process of the sub-network.

Claims

1. A semi-supervised medical image segmentation method based on perturbation collaborative training, characterized in that: The following steps are involved: (1) Inputting a training dataset, wherein the training dataset contains a large amount of unlabeled data and a small amount of labeled data; randomly initializing the parameters of two encoder-decoder segmentation networks with the same structure; (2) Sample labeled medical images and their corresponding labels (X l ,Y l ) and unlabeled medical images X u ,in Y l ∈{0,1} H×W×D , H, W, D represent the height, width and depth of the image respectively, Y l The value range of the elements is {0,1}, 0 represents background voxels and 1 represents foreground voxels; (3) preprocessing and data augmentation of the input data, wherein the preprocessing includes normalizing the input image to the interval [0, 1] and randomly cropping the input image and label; (4) For labeled data X l , input it into the two sub-networks, and obtain the prediction results P of the two sub-networks l,1 ,P l,2 ∈[0,1] H×W×D , based on the prediction result P l,1 ,P l,2 and label Y l Calculate Dice loss and binary cross entropy loss and sum them as the training loss on labeled data (5) The unlabeled data X u Input into the two sub-networks, and during the calculation process, the encoder features are perturbed based on the feature information between the sub-networks to generate pseudo labels (6) Again, the unlabeled data X u It is input into the two sub-networks to obtain the prediction results P of the two sub-networks. u,1 ,P u,2 ∈[0,1] H×W×D , based on the prediction result P u,1 ,P u,2 and pseudo labels Calculate Dice loss and binary cross entropy loss and sum them as the training loss on unlabeled data (7) Training loss based on labeled data and the training loss on unlabeled data Calculate the overall training loss Update the segmentation network; (8) Determine whether the maximum number of iterations has been reached, otherwise return to step (2); (9) Stop training and use the trained model to make predictions for the input image X.

2. The semi-supervised medical image segmentation method based on perturbation collaborative training according to claim 1, characterized in that: The training loss on the labeled data described in step (4) is defined as: Among them, P l,i represents the prediction result of the i-th sub-network for the labeled data, Y l represents the label of the labeled data, represents the binary cross entropy loss function, Represents the Dice loss function.

3. The semi-supervised medical image segmentation method based on perturbation collaborative training according to claim 2, characterized in that: The loss on the unlabeled data described in step (7) is defined as follows: Among them, P u,i represents the prediction result of the i-th sub-network for unlabeled data, represents the pseudo label generated by the 3-i sub-network; The overall training loss described in step (7) The definition is as follows: where β t Control the balance between labeled loss and unlabeled loss; gradually increase β in the form of an exponential function in the iterative round t t =β0exp(-5*(1-t / T max ) 2 ), where β0 is the control and Balanced hyperparameter, T max is the maximum number of iterations.

4. The semi-supervised medical image segmentation method based on perturbation collaborative training according to claim 1, characterized in that: The step (5) is specifically as follows: (5.1) For unlabeled data X u , the encoders of the two sub-networks each extract features Where j represents the feature extracted by the encoder of the jth layer; (5.2) Randomly sample j from {1,2,…M} and randomly sample λ from Beta(α,α), where Beta represents the Beta distribution and α is a hyperparameter used to control the shape of the Beta distribution; (5.3) Based on the following formula and Make a perturbation: (5.4) The perturbed feature F u,1 and F u,2 Send it to the decoder to get the prediction result of the sub-network And according to and Generate pseudo labels and 5. The semi-supervised medical image segmentation method based on perturbation collaborative training according to claim 1, characterized in that: The step (9) is specifically as follows: (9.1) Preprocessing the input image, wherein the preprocessing includes resampling and registering the input image and normalizing the input image to the interval [0, 1] to obtain an optimized input image; (9.2) using a sliding window strategy to cut out a prediction sub-block of the same size as the training image from the optimized input image; (9.3) Input the predicted sub-blocks into the trained model and obtain the predicted segmentation results of each predicted sub-block through forward calculation. Here, the prediction results of a single sub-network or the integration of the prediction results of two sub-networks are used; (9.4) Use the Gaussian weighting strategy to fuse the model prediction results of each prediction sub-block to obtain the prediction result of the input image X.

6. A computer storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the computer program implements a semi-supervised medical image segmentation method based on perturbation collaborative training as described in any one of claims 1 to 5.

7. A computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: When the processor executes the computer program, it implements a semi-supervised medical image segmentation method based on perturbation collaborative training according to any one of claims 1 to 5.

Images