A lung bronchus segmentation method based on semi-supervised learning

Abstract

The application discloses a lung bronchus segmentation method based on semi-supervised learning and relates to the technical field of intelligent medical treatment. A large amount of unlabeled data is utilized by adopting a semi-supervised consistency learning method; a teacher model takes the predicted entropy of each voxel in input as a measure of approximate uncertainty, and a student model learns from meaningful and reliable targets through the uncertainty information of the teacher model; the output result of the student model is input into a confidence network, the segmentation result of the unlabeled data and the segmentation result of the labeled data are used for training the confidence network, and a confidence loss is obtained; the confidence loss is contained in a total loss, and a model is trained to minimize the total loss. The application can improve the segmentation result of the model and realize the task of automatically segmenting the lung bronchus by using a small amount of labeled data set.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of intelligent medical treatment, and in particular to a lung bronchus segmentation method based on semi-supervised learning. BACKGROUND

[0002] Lung bronchus segmentation is an important research direction in the field of medical image processing. The purpose of lung bronchus segmentation is to automatically or semi-automatically segment the lung and bronchus in medical images so that doctors can more accurately diagnose and treat lung diseases.

[0003] Traditional threshold-based segmentation methods are simple and easy to use, but are affected by noise and uneven gray scale, resulting in poor segmentation results. Edge detection-based methods use the features of object edges in images for segmentation. Common edge detection algorithms include Canny, Sobel, etc. This method is sensitive to noise and the edge detection results are unstable, resulting in poor segmentation results. Region growing-based methods start from one or more seed points and recursively add adjacent pixel points to generate lung or bronchus regions. This method is suitable for images with uniform gray scale and simple structure, but has poor segmentation results for noisy and complex structure images. In recent years, the gradual maturity of deep learning-based methods has made it possible to segment lung bronchus, but there are relatively few related studies. Labeling medical image datasets often requires experienced doctors to spend a lot of time labeling. Semi-supervised learning only requires part of the labeled data, and adds a large amount of unlabeled data to complete model training, which greatly reduces the manpower and material resources consumed by labeling data.

[0004] Therefore, the skilled in the art is committed to developing a lung bronchus segmentation method based on semi-supervised learning to automatically segment lung bronchus with a small amount of labeled dataset. SUMMARY

[0005] In view of the above defects of the prior art, the technical problem to be solved by the present application is to improve the model prediction result and realize the task of automatically segmenting lung bronchus with a small amount of labeled dataset.

[0006] To achieve the above-mentioned purpose, the present application provides a lung bronchus segmentation method based on semi-supervised learning, comprising the following steps:

[0007] Step 1, collect lung CT image dataset, label part of the dataset to obtain part of the labeled training dataset;

[0008] Step 2, constructing a student-teacher model, the student model and the teacher model are two split models with the same structure, the teacher model parameters are updated by the student model training parameter exponential shift average method; the teacher model estimates the uncertainty of each target prediction, and the student model learns from meaningful and reliable targets by using the uncertainty information of the teacher model; the unlabeled data set is added to the student model and the teacher model after adding different random disturbances for consistency training and obtaining consistency loss;

[0009] Step 3, input the labeled data set after adding random disturbance into the student model for supervised training, and calculate the cross entropy loss, and take the calculated cross entropy loss as the supervised loss;

[0010] Step 4, input the output result of the student model into the confidence network, and train the confidence network with the segmentation result of the unlabeled data input into the student model and the segmentation result of the labeled data, and obtain the confidence loss;

[0011] Step 5, the total loss is the sum of the consistency loss obtained in step 2, the cross entropy loss obtained in step 3 and the confidence loss obtained in step 4 after multiplying the weight coefficient, and the trained model is obtained when the total loss converges.

[0012] Further, in step 1, data preprocessing is performed on the data set: first, the original data set and the label are cropped to retain the lung bronchus region; then data augmentation is performed on the data set.

[0013] Further, the segmentation model includes 3DU-Net, V-Net, nnU-Net.

[0014] Further, in step 2, the teacher model calculates the prediction entropy of each voxel in the input as a measure of approximate uncertainty.

[0015] Further, in step 2, the uncertainty of each voxel predicted by the teacher model is obtained, and after screening, the student model learns from the voxel points with low uncertainty.

[0016] Further, the random disturbance is adding Gaussian random noise to the data set.

[0017] Further, the confidence network includes an input layer, a hidden layer and an output layer.

[0018] Further, the input layer receives the segmentation result output by the student network.

[0019] Further, the hidden layer includes a convolution layer, an activation function and a pooling layer, the convolution layer extracts the features of the input data through convolution operation, and further extracts and compresses the features through the activation function and the pooling layer.

[0020] Further, the output layer receives the reconstructed data from the hidden layer and outputs the prediction result of the model.

[0021] In the preferred embodiment of the present application, since in the traditional student-teacher model, the student model does not make uncertainty estimation on the teacher model and directly learns, it will affect the final learning result of the student model. The teacher model of the present application calculates the prediction entropy of each voxel in the input as a measure of approximate uncertainty. The student model gradually learns from meaningful and reliable targets through the uncertainty information of the teacher model. The specific way of the teacher model to estimate uncertainty is that for each voxel in the input, the prediction entropy is selected as a measure of approximate uncertainty, and the calculation formula is as follows:

[0022]

[0023] In the case of random input Gaussian noise, the teacher model is randomly forward propagated T times, represents the probability of the t-th forward propagation prediction being the c-th class. The uncertainty of the teacher model prediction for each voxel is obtained. After screening, the student model learns to the voxel with low uncertainty.

[0024] The traditional semi-supervised face model does not consider the difference between the labeled data and the unlabeled data after passing through the model output result, ignores the loss of judging whether the output result is labeled data or unlabeled data, and obtains a poor model segmentation result. The output result of the student model of the present application is input into the confidence network, the segmentation result of the unlabeled data input into the student model and the segmentation result of the labeled data are trained on the confidence network, and the confidence loss is obtained. The total loss contains the confidence loss, and the training model makes the total loss loss minimum, which can improve the model segmentation result. The confidence network is composed of an input layer, a hidden layer and an output layer. The input layer receives the segmentation result output by the student network, the hidden layer includes a convolution layer, an activation function and a pooling layer, the convolution layer extracts the features of the input data through convolution operation, and further extracts and compresses the features through the activation function and the pooling layer. The hidden layer is the core part of the confidence network, which learns to represent the complex structure and abstract features of the input data. The output layer receives the reconstructed data from the hidden layer and outputs the prediction result of the model.

[0025] The confidence loss is as follows:

[0026]

[0027] y i represents the segmentation result of the i-th voxel of the labeled data input into the student model, y~(i) is the segmentation result of the i-th voxel of the labeled data input into the student model. i represents the segmentation result of the i-th voxel of the unlabeled data input into the student model.

[0028] The medical image data set is often labeled by experienced doctors with a lot of time. The present application uses a semi-supervised consistency learning method to utilize a large amount of unlabeled data. A student-teacher model is constructed, the student model and the teacher model are two segmentation models with the same structure, and the teacher model parameters are updated by the student model training parameter exponential translation average method. The teacher model estimates the uncertainty of each target prediction. The student model gradually learns from meaningful and reliable targets by using the uncertainty information of the teacher model. Different random disturbances are added to the unlabeled data set, and consistency training is carried out by adding the student model and the teacher model to obtain the consistency loss.

[0029] Compared with the prior art, the present application has the following obvious substantial features and significant advantages:

[0030] 1. The present application does not need all labeled data, and uses unlabeled data to improve the training model result.

[0031] 2. The final prediction result of the model trained by the present application is better than that of the model trained without using the present application.

[0032] The concept, specific structure and technical effects of the present application will be further described below with reference to the accompanying drawings, so as to fully understand the purpose, features and effects of the present application. BRIEF DESCRIPTION OF DRAWINGS

[0033] Figure 1 is a semi-supervised learning lung bronchus segmentation algorithm system framework of a preferred embodiment of the present application;

[0034] Figure 2 is the input and output of a belief network of a preferred embodiment of the present application;

[0035] Figure 3 is the structure of a belief network of a preferred embodiment of the present application;

[0036] Figure 4 is a segmentation result using the model obtained by the present application. DETAILED DESCRIPTION

[0037] The present application can be embodied in many different forms, and the scope of protection of the present application is not limited to the embodiments described herein.

[0038] In the drawings, like reference numerals refer to like parts throughout the various views. The sizes and the thicknesses of the components shown in the drawings are arbitrarily shown for making the drawings easy to understand, and the present application is not limited to the sizes and the thicknesses of the components. The thicknesses of the components are exaggerated in some places for the purpose of making the drawings clear.

[0039] The experimental operation is run in a Linux system. The code of the entire experimental design is written in Python, and the Python version is 3.7.

[0040] The experimental data is from ATM22. ATM22 is a multi-center collection of 500 CT scans. The chest CT scans are collected from the public LIDC-IDRI dataset and Shanghai Chest Hospital. Each chest CT scan is first pre-processed by some powerful deep learning models and integrated strategies to obtain the preliminary segmentation result, and then carefully outlined and double-checked by three radiologists with more than five years of professional experience to obtain the final refined airway tree structure. The lung CT image size in the ATM-22 dataset is 512*512*L (where L is the number of slices in the CT volume data). The dataset used in this paper is 60 data in ATM22.

[0041] The implementation process of the embodiment is as shown in Figure 1 , and specifically includes the following steps:

[0042] A lung bronchus segmentation method based on semi-supervised learning mainly includes the following steps:

[0043] Step 1, collect lung CT image dataset, label part of the dataset to obtain part of the labeled training dataset. In this example, 16 labeled lung CT images are selected, and 44 unlabeled lung CT images are added as original data sets, and data preprocessing is performed on the data set;

[0044] Further, the data preprocessing in step 1 includes the following steps: first, the original data set and the label are cropped to retain the lung bronchus region. In the data set, all data set labels are marked as 1 in the lung bronchus region, and the other parts are marked as 0. The label is cropped according to the maximum range of the three-dimensional coordinates x, y, z of the marked 1 region, and the original data set is cropped according to the corresponding label 1 part, so that the lung bronchus part can be obtained. For the unlabeled data set, the maximum range of the three-dimensional coordinates x, y, z of all image labels in the labeled data set is used as the unlabeled data set and the candidate region. Then the data set is data enhanced.

[0045] Optionally, the data enhancement can adopt random cropping, random rotation, and translation. In this example, the data enhancement adopted is random cropping and random rotation.

[0046] Step 2, constructing a student-teacher model, the student model and the teacher model are two split models with the same structure, and the teacher model parameters are updated by the student model training parameter exponential moving average method. The exponential moving average method can be represented by the following formula:

[0047] θ t ' = αθ t-1 ' + (1-α)θ t

[0048] Where θ t is the model parameter of the student model at the training period t; θ t-1 ' and θ t ' are the time parameters of the teacher model at the training period t-1 and t, respectively, and α is the weight of the original parameter retained in the updating process. The teacher model calculates the prediction entropy of each voxel in the input as a measure of approximate uncertainty. The student model gradually learns from meaningful and reliable targets through the uncertainty information of the teacher model.

[0049] The specific way of the teacher model estimating uncertainty is that for each voxel in the input, the prediction entropy is selected as the measure of approximate uncertainty, and the calculation formula is as follows:

[0050]

[0051] In the case of random input Gaussian noise, the teacher model is randomly forward propagated T times, The probability of the c-th class predicted by the t-th forward propagation is represented. The uncertainty of the teacher model prediction for each voxel is obtained.

[0052] When calculating the consistency loss of adding different random disturbances to the unlabeled data set and adding the student model and the teacher model segmentation results respectively, only the data with a credibility greater than 0.75 is calculated for the teacher model to calculate the consistency loss, and the consistency loss is represented by the following formula:

[0053]

[0054] Where f v is the prediction result of the student model, f′ v is the prediction result of the teacher model at the v-th voxel, μ v is the uncertainty calculated by the teacher model at the v-th voxel.

[0055] Optionally, the segmentation model in the student model and the teacher model includes but is not limited to any one of 3DU-Net, V-Net (V-Net, V-shaped network), and nnU-Net. Different neural network segmentation models can be applied according to the specific task requirements, and V-Net is selected as the segmentation network in the present example.

[0056] Step 3, the labeled data set is added to the random disturbance and input into the student model for supervised training, and the cross-entropy loss is calculated, and the calculated cross-entropy loss is used as the supervised loss;

[0057] Optionally, the random disturbance is to add Gaussian random noise to the data set.

[0058] Further, the cross-entropy loss calculation formula of step 3 is as follows:

[0059]

[0060] Wherein, y i represents the true value of the ith voxel, represents the prediction value of the student model at the ith voxel.

[0061] Step 4, the output result of the above student model is input into the confidence network, the segmentation result of the student model input into the unlabeled data and the segmentation result of the labeled data are input into the confidence network for training, and the confidence loss is obtained, the input and output of the confidence network are as shown in Figure 2 , and the confidence network structure is as shown in Figure 3 .

[0062] Further, the confidence network of step 4 is composed of an input layer, a hidden layer and an output layer. The input layer receives the segmentation result output by the student network, the hidden layer includes a convolution layer, an activation function and a pooling layer, the convolution layer extracts the features of the input data through convolution operation, and further extracts and compresses the features through the activation function and the pooling layer and the like. The hidden layer is the core part in the confidence network, which learns to represent the complex structure and abstract features of the input data. The output layer receives the reconstructed data from the hidden layer and outputs the prediction result of the model.

[0063] Further, the confidence loss of step 4 is as follows:

[0064]

[0065] y i represents the segmentation result of the labeled data input into the student model, represents the segmentation result of the unlabeled data input into the student model.

[0066] Step 5, the total loss is the sum of the consistency loss obtained in step 2, the cross-entropy loss obtained in step 3 and the confidence loss obtained in step 4 after being multiplied by the weight coefficient, and the trained model is obtained when the total loss converges. Figure 4 A test result of the trained model obtained in the present example, the test input image is a lung CT image.

[0067] Further, the total loss in step 5 is as follows:

[0068] L = μ1L c + μ2L s + μ3L D

[0069] μ1, μ2, μ3 are weight coefficients of L c , L s , L D , respectively. L c is a consistency loss, L s is a cross-entropy loss, and L D is a confidence loss.

[0070] The preferred embodiments of the present application are described in detail above. It should be understood that modifications and variations can be made by those of ordinary skill in the art without departing from the spirit and scope of the present application. Accordingly, it is intended that all such possible modifications and variations be included within the scope of the present application as defined by the following claims.

Claims

1. A lung bronchus segmentation method based on semi-supervised learning, characterized in that, Includes the following steps: Step 1: Collect a dataset of lung CT images, and annotate a portion of the dataset to obtain a partially annotated training dataset; Step 2: Construct student and teacher models. The student and teacher models are two segmentation models with identical structures. The teacher model parameters are updated using the exponential shift average method based on the student model training parameters. The exponential shift average method is expressed by the following formula: in These are the model parameters for the student model during training period t. and These are the time parameters of the teacher model at training periods t-1 and t, respectively. The weights that retain the original parameters during the update process; the teacher model uses the prediction entropy of each voxel in the input as a measure of approximate uncertainty; the student model learns gradually from meaningful and reliable targets by using the uncertainty information from the teacher model. The teacher model estimates uncertainty by using prediction entropy as a measure of approximate uncertainty for each voxel in the input, calculated as follows: The teacher model is subjected to T random forward propagations with randomly input Gaussian noise. Let represent the probability that the forward propagation prediction is of class c in the t-th time, and obtain the uncertainty of the teacher model's prediction for each voxel. When calculating the consistency loss of the segmentation results of the student model and the teacher model after adding different random perturbations to the unlabeled dataset, the consistency loss is only calculated for data with a confidence level greater than 0.75 in the teacher model. The consistency loss is expressed by the following formula: in, For the prediction results of the student model, The teacher model's prediction at the v-th ... The uncertainty calculated by the teacher model at the v-th genus; Step 3: Add random perturbation to the labeled dataset and input it into the student model for supervised training, and calculate the cross-entropy loss. Use the calculated cross-entropy loss as the supervised loss. Step 4: Input the output of the student model into the confidence network. Use the unlabeled data and the segmentation results of the student model and the labeled data to train the confidence network and obtain the confidence loss. Step 5: The total loss is the sum of the consistency loss obtained in Step 2, the cross-entropy loss obtained in Step 3, and the confidence loss obtained in Step 4, multiplied by the weight coefficients respectively. The trained model is obtained when the total training loss converges. The confidence network includes an input layer, a hidden layer, and an output layer; The input layer receives the segmentation results output by the student network; The hidden layer includes a convolutional layer, an activation function, and a pooling layer. The convolutional layer extracts features from the input data through convolution operations, and the activation function and pooling layer further extract and compress the features.

2. The lung bronchus segmentation method based on semi-supervised learning as described in claim 1, characterized in that, Step 1 involves preprocessing the dataset: first, the original dataset and labels are cropped to retain the region containing the lungs and bronchi; then, data augmentation is performed on the dataset.

3. The lung bronchus segmentation method based on semi-supervised learning as described in claim 1, characterized in that, The segmentation models include 3DU-Net, V-Net, and nnU-Net.

4. The lung bronchus segmentation method based on semi-supervised learning as described in claim 1, characterized in that, In step 2, the teacher model calculates the prediction entropy of each voxel in the input as a measure of approximate uncertainty.

5. The lung bronchus segmentation method based on semi-supervised learning as described in claim 1, characterized in that, In step 2, the uncertainty of the teacher model's prediction for each voxel is obtained. After screening, the student model learns from voxel points with low uncertainty.

6. The lung bronchus segmentation method based on semi-supervised learning as described in claim 1, characterized in that, The random perturbation is the addition of Gaussian random noise to the dataset.

7. The lung bronchus segmentation method based on semi-supervised learning as described in claim 1, characterized in that, The output layer receives the data reconstructed from the hidden layer and outputs the model's prediction results.

Images