A semi-supervised medical image segmentation method based on staggered structure

Abstract

The application discloses a kind of semi-supervised medical image segmentation methods based on staggered structure, the method will be segmentation model U-Net as backbone network, semi-supervised segmentation architecture Mean TeacherU-Net is constructed, the architecture is constrained to the prediction of supervised part and label, and the consistency of the prediction of student model and the prediction of teacher model is strengthened;The performance of the model is improved using staggered structure, staggered structure includes staggered structure in time sequence and staggered structure in direction, the two structures further constrain the feature map of intermediate layer in student model and the feature map of intermediate layer in teacher model.The application strengthens the correlation of feature map and the consistency of student model and teacher model in network coding and decoding process, so as to improve segmentation effect, with higher practical value.

Description

TECHNICAL FIELD

[0001] The present application belongs to the field of image segmentation, and particularly relates to a semi-supervised medical image segmentation method based on staggered structure. BACKGROUND

[0002] Labeling data is mostly a labor-intensive work and requires professional skills, so the cost of labeling all training data is high, and it is usually difficult to obtain all labeled data sets for full supervision learning.

[0003] In contrast, the cost of obtaining unlabeled data is low, and semi-supervised learning, which uses unlabeled data together with labeled data as training data, can improve the performance of the model. In addition, most semi-supervised learning strategies are to obtain additional information specific to other learning strategies by extending full supervision learning or unsupervised learning, so this method also has theoretical significance in machine learning and as a model of human learning. At present, semi-supervised learning has become a hot direction in the field of machine learning.

[0004] Under this background, in order to utilize unlabeled data in medical image segmentation tasks, the present application uses a U-Net model as a backbone network to construct a semi-supervised architecture Mean TeacherU-Net, and extracts the middle layers of the student model and the teacher model in the architecture to form a staggered structure in time sequence and direction. SUMMARY

[0005] The present application proposes a semi-supervised medical image segmentation method based on staggered structure to utilize unlabeled data in medical image segmentation tasks, which uses labeled and unlabeled images to train the segmentation network together, and at the same time, through the construction of staggered structure in time sequence and direction, the consistency between the teacher and student networks is enhanced.

[0006] Technical scheme: The semi-supervised medical image segmentation method based on staggered structure comprises the following steps:

[0007] (1) Preprocess the medical images that need to be segmented to obtain labeled and unlabeled data;

[0008] (2) Construct a semi-supervised segmentation model based on staggered structure: use a U-Net model as a backbone network to construct a semi-supervised architecture Mean TeacherU-Net, and extract the middle layers of the student model and the teacher model in the architecture to form a staggered structure in time sequence and direction;

[0009] (3) Use labeled and unlabeled data to train the segmentation model;

[0010] (4) using the trained model to test on the target image to obtain the final medical image segmentation result.

[0011] Further, the pre-processing of step (1) is implemented as follows:

[0012] The target organ region is intercepted, the CT scan file is converted into an image, the image size is converted to 256x256, and the image is subjected to augmentation operations including standardization, normalization, random rotation and cropping.

[0013] Further, the network of step (2) includes a U-Net network U t as a teacher model, including an encoder E t and a decoder D t , and a U-Net network U s as a student model, including an encoder E s and a decoder D s .

[0014] Further, the U-Net uses Dice loss L Dice and cross-entropy loss L CE as the loss function L sup of the supervised part:

[0015]

[0016]

[0017] L sup =L Dice +L CE

[0018] where y l is the label, is the predicted value, and the weight w of L CE is a hyperparameter.

[0019] When using the U-Net model as the backbone network to build the semi-supervised architecture Mean Teacher U-Net, the teacher model uses the exponential moving average (EMA) of the student model to update its own parameters:

[0020] θ t，i =αθ t，i-1 +(1-α)θ s，i

[0021] where α is a hyperparameter, θ t，i is the parameter of the teacher model at the i-th iteration, and θ s，i is the parameter of the student model at the i-th iteration.

[0022] Using MSE loss as the consistency loss L consis Strengthen the consistency between the predictions of the student model and the teacher model:

[0023] λ consis =δω(e, ε)

[0024]

[0025] Where δ and ε are hyperparameters, e is the number of training epochs, and ω is the sigmoid rampup method used to adjust the weights λ of the loss function. consis The number of training epochs increases to reduce the impact of models with lower performance in the early stages of training on the training process. It is the teacher model's prediction of labeled data. It is the teacher model's prediction of unlabeled data. It is the student model's prediction of labeled data. It is the student model's prediction of unlabeled data.

[0026] Furthermore, the implementation process of extracting the intermediate layer of the student model and teacher model in step (2) to form an interleaved structure in terms of timing and direction is as follows:

[0027] The temporally interleaved structure uses convolution to process the feature maps f of the intermediate layers of the student model encoder. s，en Feature maps f of the intermediate layer of the decoder s，de The feature map f of the intermediate layer of the entire model s，mid , and the feature map f of the intermediate layer of the teacher model encoder t，en Feature maps f of the intermediate layer of the decoder t，de The feature map f of the intermediate layer of the entire model t，mid Extract the data and apply the loss function to f. s，en and f t，de f s，mid and f t，mid f s，de and f t，en Apply constraints, loss function L tc for:

[0028]

[0029]

[0030]

[0031]

[0032] Interlacing structures in the direction of ft，en with f s，en , f t，de with f s，de , loss function L dc is:

[0033]

[0034]

[0035]

[0036] wherein Φ t，en is the part of f t，en that overlaps with f s，en , Φ t，de is the part of f t，de that overlaps with f s，de , Φ s，en is the part of f s，en that overlaps with f t，en , Φ s，de is the part of f s，de that overlaps with f t，de .

[0037] Loss function L semisup of the semi-supervised part and total loss function L total is:

[0038] L semisup = L consis + λ tc L tc + λ dc L dc

[0039] L total = λ sup L sup + L semisup = λ sup L sup + L consis + λ tc L tc + λ dc L dc

[0040] wherein hyperparameter λ sup is the weight of the supervised part loss function, λ tc is the weight of the loss function of the staggered structure in time sequence, λ dc is the weight of the loss function of the staggered structure in direction.

[0041] Further, the step (3) comprises the following steps:

[0042] (31) Configure server environment, upload model code, select GPU;

[0043] (32) Set hyperparameters, including number of iterations, weights of loss function, learning rate, etc.

[0044] (33) Initialize model parameters.

[0045] (34) Divide data set, i.e. randomly select 70% of data as training set, and the remaining 30% as test set, and run code.

[0046] (35) Train model on training set and update parameters.

[0047] (36) Save model and visualize results after a certain number of iterations.

[0048] (37) Select the best saved model as the final model.

[0049] (38) Test the effect of the final model on the test set and output segmentation results.

[0050] Advantages: Compared with the prior art, the advantages of the present application are: 1. The semi-supervised segmentation network proposed by the present application utilizes unlabeled data. Since the label annotation process of medical image segmentation problems requires professional skills and has high labor and time cost, the present application not only reduces the burden of professional personnel on medical image annotation, but also has good segmentation performance, and has practical application value. 2. The interleaved structure proposed by the present application enhances the consistency between the teacher and student models, so that the model converges faster and improves the segmentation performance. BRIEF DESCRIPTION OF DRAWINGS

[0051] Figure 1 is a flowchart of the present application;

[0052] Figure 2 is a schematic diagram of the semi-supervised segmentation network based on interleaved structure constructed by the present application;

[0053] Figure 3 is a flowchart of the present application for training and testing the model. DETAILED DESCRIPTION

[0054] The present application will be further described in detail below in conjunction with the accompanying drawings.

[0055] As shown in Figure 1 , the present application discloses a semi-supervised medical image segmentation method based on interleaved structure, which specifically comprises the following steps:

[0056] Step 1: Medical image preprocessing.

[0057] First, the original datasets of medical images typically contain not only the target region but also some non-target regions, so it is necessary to first extract the target organ region. Second, the original medical images are usually presented as CT scan files, while this invention is applicable to image segmentation, requiring the CT scan files to be converted into image files. To facilitate image processing by the model, the image size is modified to 256×256, and the image pixel values ​​are standardized, i.e., the mean is subtracted and the corresponding variance is divided; and a normalization operation is performed to scale the image pixel values ​​to the range of [-1, 1]. Since the number of medical images is relatively small, to avoid overfitting of the model, this invention uses image augmentation operations such as random cropping and random rotation.

[0058] Step 2: Construct a semi-supervised medical image segmentation network based on an interleaved structure.

[0059] The application scenario of this invention is semi-supervised medical image segmentation, that is, using both unlabeled and labeled data simultaneously to obtain good segmentation results on the target image while reducing the annotation burden on professionals.

[0060] The semi-supervised medical image segmentation network constructed in this invention makes good use of unlabeled data. This network includes a U-Net network U acting as the teacher model. t Including encoder E t With decoder D t A U-Net network U serving as a student model s Including encoder E s With decoder D s This structure can utilize unlabeled images and accelerate model convergence and improve segmentation results by constraining the consistency between teacher and student models, unlike conventional image segmentation methods.

[0061] Given labeled data (x) l ), tag (y l ) and unlabeled data (x u The present invention aims to use these images to train a model in order to obtain better segmentation performance on test images.

[0062] like Figure 2 The model structure is shown in the figure: the preprocessed image is simultaneously input into the teacher and student U-Net respectively, and after being downsampled by the encoder, it is then input into the decoder for upsampling. At the same time, the model adopts a skip-connection operation to directly forward the low-level feature map from the encoder to the corresponding layer of the decoder, thereby reducing the gap between the two.

[0063] The student model updates parameters using gradient descent; the loss function L in the supervised part of this process is... sup (Dice lost L) Diceand cross-entropy loss L CE is as follows:

[0064]

[0065]

[0066] L sup = L Dice + L CE

[0067] where y l is the label, is the predicted value, and L CE is the weight of the loss function. The weight w is a hyperparameter.

[0068] On the other hand, the teacher model uses the exponential moving average (EMA) of the student model to update its own parameters, as follows:

[0069] θ t，i = αθ t，i-1 + (1-α)θ s，i

[0070] where α is a hyperparameter, θ t，i is the parameter of the teacher model at the i-th iteration, and θ s，i is the parameter of the student model at the i-th iteration. Meanwhile, the input of the teacher model is added with noise, which ensures that there is a difference between the two models, facilitating the parameter update when the models are trained in the direction of keeping consistency.

[0071] Using the MSE loss as the consistency loss L consis to strengthen the consistency between the prediction of the student model and the prediction of the teacher model, as follows:

[0072] λ consis = δω(e,ε)

[0073]

[0074] where δ, ε are hyperparameters, e is the training epoch, and ω() is the sigmoid rampup method proposed by Laine et al. in Temporal ensembling for semi-supervised learning, which is used to increase the weight λ consis of the loss function with the increase of the training epoch, so as to reduce the influence of the model with lower performance in the early stage of training on the training process, is the prediction of the teacher model on the labeled data, is the prediction of the teacher model on the unlabeled data, is the prediction of the student model on the labeled data, is the prediction of the student model on the unlabeled data.

[0075] In addition, the model extracts the intermediate layers of the student model and the teacher model in the architecture to form the interleaving structure in the time sequence and the direction. The interleaving structure in the time sequence uses convolution to extract the feature maps f s，en , f s，de , f s，mid and f t，en , f t，de , f t，mid of the intermediate layers of the encoder and the decoder of the student model and the teacher model, and uses the loss function to constrain f s，en and f t，de , f s，mid and f t，mid , f s，de and f t，en , which makes the feature maps in the encoder better fused with the feature maps in the decoder, and strengthens the correlation of the feature maps in the network encoding and decoding process and the consistency of the student model and the teacher model. The loss function L tc is as follows:

[0076]

[0077]

[0078]

[0079]

[0080] The interleaving structure in the direction constrains f t，en and f s，en , f t，de and f s，de , which makes the connection between the feature maps in multiple directions emphasized, further strengthening the consistency of the student and the teacher model. The loss function L dc is as follows:

[0081]

[0082]

[0083]

[0084] wherein, Φ t，en is f t，en in f s，enOverlapping part, Φ t，de is f t，de Overlapping part, Φ s，de is f s，en Overlapping part, Φ s，en is f t，en Overlapping part, Φ s，de is f s，de Overlapping part, Φ t，de .

[0085] The loss function of the semi-supervised part L semisup is as follows:

[0086] L semisup = L consis + λ tc L tc + λ dc L ac

[0087] The total loss function L total is as follows:

[0088] L total = λ sup L sup + L semisup = λ sup L sup + L consis + λ tc L tc + λ dc L ac

[0089] wherein the hyperparameter λ sup is the weight of the supervised part loss function, λ tc is the weight of the loss function of the staggered structure in time sequence, and λ dc is the weight of the loss function of the staggered structure in direction.

[0090] Step 3: training the model, specifically as shown in Figure 3 :

[0091] 1) configuring the server environment, uploading the model code, and selecting the GPU;

[0092] 2) setting the hyperparameters, such as initially setting the total number of iterations to 20000, saving the model and visualizing the results when the initial iteration number reaches 1000, setting the weight of the loss function to 1, and setting the learning rate to 0.0002;

[0093] 3) initializing the model parameters;

[0094] 4) Divide the dataset, that is, randomly select 70% of the data as the training set, and the remaining 30% as the test set, and run the code;

[0095] 5) Train the model on the training set and update the parameters;

[0096] 6) Save the model and visualize the results after a certain number of iterations;

[0097] 7) Select the best saved model as the final model;

[0098] 8) Test the final model on the test set and output the segmentation results.

[0099] Step 4: Model prediction segmentation results.

[0100] Input the test image into the model saved in step 3, output the segmentation test indicators and display the final image segmentation results.

[0101] The above describes in detail a semi-supervised medical image segmentation method based on an interleaved structure provided by the present application. It is worth noting that there are many methods and ways to implement this technical solution, and the above description is only the preferred embodiment of the present application, which is used to help understand the method and core idea of the present application; at the same time, for those skilled in the art, modifications and adjustments made on the basis of the core idea of the present application will be considered as the protection scope of the present application. In summary, the content of the present application should not be understood as a limitation of the present application, and the protection scope of the present application should be defined by the appended claims.

Claims

1. A semi-supervised medical image segmentation method based on an interleaved structure, characterized in that, The method comprises the following steps: (1) preprocessing the medical image to be segmented to obtain labeled and unlabeled data; (2) Constructing a semi-supervised segmentation model based on staggered structure: using U-Net model as the backbone network to construct semi-supervised architecture Mean TeacherU-Net, extracting the intermediate layers of student model and teacher model in the architecture to constitute staggered structure in time sequence and direction, the implementation process is as follows: the staggered structure in time sequence uses convolution to extract the feature maps of the intermediate layers of student model encoder , the feature maps of the intermediate layers of decoder , the feature maps of the intermediate layers of the whole model , and the feature maps of the intermediate layers of teacher model encoder , decoder , the feature maps of the intermediate layers of the whole model , and using loss function to constrain and , and , and , is: Interleaved structure in the direction of the pair of With , With Constraints, loss function : wherein is the portion of overlapping with is the portion of overlapping with is the portion of overlapping with is the portion of overlapping with Loss function for semi-supervised part and total loss function is: wherein the hyperparameters are weights for the supervised part loss function, are weights for the loss function of the interleaving structure in time, are weights for the loss function of the interleaving structure in direction, is the loss function of the supervised part, is the loss function using MSE loss as consistency. (3) training the segmentation model using labeled and unlabeled data; (4) using the trained model to test on the target image to obtain the segmentation result.

2. The semi-supervised medical image segmentation method based on the interleaved structure according to claim 1, characterized in that, The preprocessing process of step (1) is as follows: Cutting the target organ region, converting the CT scan file into an image, converting the image size to 256x256, and performing augmentation operations including standardization, normalization, random rotation, and cropping on the image.

3. The semi-supervised medical image segmentation method based on the interleaved structure according to claim 1, characterized in that, The network of step (2) comprises a U-Net network as a teacher model , comprising an encoder and a decoder , a U-Net network as a student model , comprising an encoder and a decoder .

4. The semi-supervised medical image segmentation method based on the interleaved structure according to claim 3, characterized in that, The U-Net uses a Dice loss and a cross-entropy loss as loss function for the supervised part : )) wherein, is a label, is a predicted value, a weight is a hyperparameter; When using the U-Net model as the backbone network to construct the semi-supervised architecture Mean TeacherU-Net, the teacher model uses the exponential moving average of the student model to update its own parameters: wherein, is a hyperparameter, is the parameter of the teacher model at the th iteration, is the parameter of the student model at the th iteration, is the parameter of the teacher model at the th iteration; Using MSE loss as consistency loss Strengthen the consistency between the prediction of the student model and the prediction of the teacher model: wherein, , is a hyperparameter, is a training epoch, is a sigmoid rampup method for letting the weight of the loss function increase with increasing training epochs to reduce the influence of models with lower performance in the early training phase on the training process, is the prediction of the teacher model on the labeled data, is the prediction of the teacher model on the unlabeled data, is the prediction of the student model on the labeled data, is the prediction of the student model on the unlabeled data.

5. The semi-supervised medical image segmentation method based on the interleaved structure according to claim 1, characterized in that, The step (3) comprises the following steps: (31) configure the server environment, upload the model code, and select the GPU; (32) set the hyperparameters, including the number of iterations, the weight of the loss function, and the learning rate; (33) initialize the model parameters; (34) divide the data set, that is, randomly select 70% of the data as the training set and the remaining 30% as the test set, and run the code; (35) train the model on the training set and update the parameters; (36) save the model and visualize the results after a certain number of iterations; (37) select the best model from the saved models as the final model; (38) test the effect of the final model on the test set and output the segmentation result.

Images