Calculus CT image segmentation method and device based on deep learning and mixed loss

Abstract

The invention relates to a stone CT image segmentation method and device based on deep learning and mixed loss, and belongs to the field of medical image analysis. Preprocessed data are sent to an encoder of a 3D Unet of a segmentation model, step-by-step convolution and maximum pooling calculation are carried out on an input kidney stone CT image, and the input of each step of calculation is the output of the previous step; the preprocessed data are partitioned and flattened into sequence vectors, the sequence vectors are sent into a ViT encoder of a segmentation model for self-attention calculation, and global features of a calculus CT image are obtained; the obtained global features are sent to a global feature decoder of a segmentation model decoder, deconvolution decoding calculation is carried out step by step, and a plurality of feature maps with different channel numbers are obtained; a feature map obtained through the last convolution and pooling calculation is sent to a 3D Unet decoder of a segmentation model decoder, and deconvolution calculation is carried out successively; according to the method, the position and form of the calculus in the calculus CT image can be accurately identified, and the segmentation effect is improved.

Description

Technical Field

[0001] The present invention relates to a stone CT image segmentation method and device based on deep learning and mixed loss, belonging to the technical field of medical image analysis, and can be applied to the diagnosis of stone patients by medical institutions and the processing of stone CT images. Background Art

[0002] Automatic segmentation of medical images can make the changes of anatomical or pathological structures in images clearer. It greatly improves the efficiency and accuracy of patient diagnosis. With the development of deep learning technology, deep learning technology based on convolutional neural networks has begun to be widely used in the segmentation task of medical images. Due to the less interference from noise and low contrast in medical images, convolutional neural networks have achieved better results in segmentation tasks. Medical image segmentation technology based on convolutional neural networks can extract features layer by layer by stacking multiple convolutional layers, and can extract features at a higher dimension.

[0003] In order to solve the problem of medical image segmentation, Ronneberger et al. proposed Unet. The network structure of Unet is similar to the English letter U, so it is called Unet. Since Unet performs well in the problem of medical image segmentation, most researchers will first use the experimental results of Unet as a baseline. Some medical images such as MRI and CT are not 2D, but 3D. For this problem, et al. proposed 3D Unet. Compared with Unet, the input and output of 3D Unet are both 3D data. 3D Unet also adds a Bacth Norm module after the ReLU activation function to make the model converge faster.

[0004] In the process of diagnosing patients through CT images, it is necessary to mark the stone area of ​​the stone CT image. The manual marking method is time-consuming and laborious, and the marking results of doctors with different experience are different. Therefore, an automatic segmentation algorithm for stone CT images is studied, which can efficiently mark the stone area of ​​stone CT images and save manpower and time costs. Summary of the invention

[0005] The present invention provides a stone CT image segmentation method and device based on deep learning and hybrid loss, which are used to automatically segment and annotate the stone area in the CT image, which can effectively reduce the time cost of manual annotation and solve the problem of inconsistent annotation results. It can automatically and accurately identify the stone position and morphology in the stone CT image and improve the segmentation effect.

[0006] The technical solution of the present invention is: in the first aspect, the present invention provides a method for segmenting CT images of stones based on deep learning and hybrid loss, and the specific steps of the method are as follows:

[0007] Step 1: Preprocess the data: Crop the 3D CT images and the corresponding segmentation marker data in the data set simultaneously to keep the ROI area consistent;

[0008] Step 2: Send the preprocessed data to the encoder of the 3D Unet of the segmentation model, and perform step-by-step convolution and maximum pooling calculations on the input kidney stone CT image. The input of each step of calculation is the output of the previous step.

[0009] Step 3: Divide the preprocessed data into blocks and flatten them into sequence vectors, and then send them to the ViT encoder of the segmentation model for self-attention calculation to obtain the global features of the stone CT image;

[0010] Step 4: Send the global features obtained in Step 3 to the global feature decoder of the segmentation model decoder, and gradually perform deconvolution decoding calculations to obtain multiple feature maps with different numbers of channels;

[0011] Step 5, the feature map obtained by the last convolution and pooling calculation in Step 2 is sent to the 3D Unet decoder of the segmentation model decoder, and the deconvolution calculation is performed successively. The input of each step of calculation is the output of the last calculation, so as to gradually restore the resolution of the feature map and the original number of channels; in each step of deconvolution calculation, the feature map obtained by the deconvolution calculation is spliced ​​with the feature map obtained after each convolution and pooling calculation in Step 2, that is, jump connection, and the convolution calculation is performed again after splicing, and then the calculation results of each step are spliced ​​with the multiple feature maps obtained by the global feature decoder in Step 4 to obtain the fused feature as the input of the next step of deconvolution calculation;

[0012] Step 6. Construct a hybrid loss function by combining Dice Loss and Focal Loss as the loss function for segmentation model training. The hybrid loss function is shown as follows:

[0013] Loss=α*Dice Loss+β*Focal Loss

[0014] Among them, α and β are adjustment coefficients, which are used to balance the proportion of the two losses.

[0015] Furthermore, the Step 1 includes: cropping the CT image into 3D data of a uniform size of 96×96×96.

[0016] Furthermore, the Step 3 includes: a total of 4 layers of ViT encoders are stacked, wherein the number of heads of the multi-head self-attention mechanism of each encoder is 2; after passing through the ViT encoder, a total of 768 channels of global feature maps of stone CT images are obtained; the global feature map obtained after encoding is subjected to torch.concat(·) splicing operation with the features after each deconvolution operation in the segmentation model decoder 3D Unet decoder step by step using deconvolution. Specifically, a total of three deconvolutions are used, and the feature map obtained by each deconvolution is subjected to torch.concat(·) splicing operation with the feature map obtained by each deconvolution operation in the segmentation model decoder 3DUnet decoder step by step from bottom to top.

[0017] The present invention also provides a stone CT image segmentation device based on deep learning and mixed loss, including an encoder, a decoder and a training module;

[0018] The encoder includes a 3D Unet encoder and a ViT encoder, which are used to extract local semantic features and global features of the stone CT image respectively;

[0019] The decoder includes a 3D Unet decoder and a global feature decoder, wherein the 3D Unet decoder is used to gradually restore the feature map to the original resolution and the original number of channels, and the global feature decoder is used to adjust the global feature to a suitable size for fusion with the features output at each step in the 3D Unet decoder;

[0020] The training module is used to construct a hybrid loss function by combining Dice Loss and Focal Loss, which is used as the loss function for segmentation model training. The hybrid loss function is as follows:

[0021] Loss=α*Dice Loss+β*Focal Loss

[0022] Among them, α and β are adjustment coefficients, which are used to balance the proportion of the two losses.

[0023] The segmentation model of the present invention first encodes the input CT image of the stone to extract semantic features, and then decodes it for segmentation. The model is composed of an encoder and a decoder. The encoder of the model is composed of a 3D Unet encoder and a ViT encoder (Vision Transformer Encoder, ViT encoder), which are respectively used to extract local semantic features and global features. In this process, the resolution of the feature map is reduced and the number of channels is increased.

[0024] The decoder of the segmentation model includes a 3D Unet decoder and a global feature decoder. The 3D Unet decoder is used to gradually restore the feature map to its original resolution and number of channels, and the global feature decoder is used to adjust the global feature to a suitable size for fusion with the features output at each step in the 3D Unet decoder.

[0025] In addition to the jump connection jump between the encoder and decoder in the original 3D UNet, the model of the present invention can also decode the results of the visual Transformer encoding step by step and add them to the decoder, which can fuse feature maps that are far apart, reduce the information lost during the convolution and pooling of the feature maps, better capture global information, better restore resolution, and obtain ideal segmentation results. The following will introduce the visual Transformer encoding and step-by-step decoding modules in the model and the hybrid loss function used when training the model.

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

[0027] 1. The present invention is based on the modification of the existing 3D Unet image segmentation network and constructs a segmentation model for stone CT images based on 3D Unet. The model incorporates the ViT encoder and the global feature decoder to provide additional global feature information of the stone CT images. The fusion of the global feature information with the original local semantic feature information can improve the diagnostic accuracy and efficiency of medical institutions for stone patients.

[0028] 2. Use a combination of Dice Loss and Focal Loss as the loss function during model training to improve the segmentation effect of the model. The model has better robustness and can also solve the problems of sample imbalance and difficulty to a certain extent.

[0029] 3. It can effectively reduce the time cost of manual labeling and solve the problem of inconsistent labeling results. It can automatically and accurately identify the location and morphology of stones in CT images of stones. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 It is the overall architecture diagram of the present invention;

[0031] Figure 2 This is a schematic diagram of data preprocessing of the present invention;

[0032] Figure 3 A schematic diagram of A and B of the present invention;

[0033] Figure 4 The figure shows the comparison of segmentation results of different methods in the present invention. DETAILED DESCRIPTION

[0034] Example 1: Figure 1-Figure 4 As shown, a stone CT image segmentation method based on deep learning and mixed loss, the specific steps of the method are as follows:

[0035] Step 1, preprocess the data: The shape and position of the stones in the CT image of the present invention are uncertain, and the stone part only occupies a relatively small part of the entire CT image. Therefore, before inputting the data into the model for segmentation training, it is necessary to preprocess the data to alleviate the impact of the non-stone area being too large on the model results. The stone CT image used in the present invention is composed of multiple 2D slices, and the size of the data is H×W×D, where H and W are the length and height of the CT slice, which are fixed values, both of which are 512 pixels. And D is the number of slices, with a numerical range of 89 to 846. The present invention crops the CT image into 3D data of a uniform size of 96×96×96. It is worth mentioning that in the data preprocessing stage, the three-dimensional CT image and the corresponding segmentation mark data in the data set need to be cropped simultaneously to keep the ROI area consistent.

[0036] Step 2: Send the preprocessed data to the encoder of the 3D Unet of the segmentation model, and perform step-by-step convolution and maximum pooling calculations on the input kidney stone CT image. The input of each step of calculation is the output of the previous step.

[0037] Step 3: Divide the preprocessed data into blocks and flatten them into sequence vectors, and then send them to the ViT encoder of the segmentation model for self-attention calculation to obtain the global features of the stone CT image;

[0038] Step 4: Send the global features obtained in Step 3 to the global feature decoder of the segmentation model decoder, and gradually perform deconvolution decoding calculations to obtain multiple feature maps with different numbers of channels;

[0039] Step 5, the feature map obtained by the last convolution and pooling calculation in Step 2 is sent to the 3D Unet decoder of the segmentation model decoder, and the deconvolution calculation is performed successively. The input of each step of calculation is the output of the last calculation, so as to gradually restore the resolution of the feature map and the original number of channels; in each step of deconvolution calculation, the feature map obtained by the deconvolution calculation is spliced ​​with the feature map obtained after each convolution and pooling calculation in Step 2, that is, jump connection, and the convolution calculation is performed again after splicing, and then the calculation results of each step are spliced ​​with the multiple feature maps obtained by the global feature decoder in Step 4 to obtain the fused feature as the input of the next step of deconvolution calculation;

[0040] Step 6, Focal Loss pays more attention to misclassified samples, emphasizing the weight of correcting difficult samples. DiceLoss focuses on the similarity of two types of samples. Although Focal Loss and Dice Loss are two loss functions used to solve class imbalance and optimize classification accuracy problems, these two loss functions focus on two different problems: sample quantity distribution and classification accuracy. Therefore, if they are combined, the sample imbalance and difficulty in the segmentation task can be comprehensively considered to obtain better model prediction results. The present invention constructs a hybrid loss function by combining Dice Loss with Focal Loss, which is used as the loss function for segmentation model training; the hybrid loss function is shown as follows:

[0041] Loss=α*Dice Loss+β*Focal Loss

[0042] Among them, α and β are adjustment coefficients, which are used to balance the proportion of the two losses.

[0043] Furthermore, the Step 1 includes: cropping the CT image into 3D data of a uniform size of 96×96×96.

[0044] Furthermore, the Step 3 includes: a total of 4 layers of ViT encoders are stacked, wherein the number of heads of the multi-head self-attention mechanism of each encoder is 2; after passing through the ViT encoder, a total of 768 channels of global feature maps of stone CT images are obtained; the global feature map obtained after encoding is subjected to torch.concat(·) splicing operation with the features after each deconvolution operation in the segmentation model decoder 3D Unet decoder step by step using deconvolution. Specifically, a total of three deconvolutions are used, and the feature map obtained by each deconvolution is subjected to torch.concat(·) splicing operation with the feature map obtained by each deconvolution operation in the segmentation model decoder 3DUnet decoder step by step from bottom to top.

[0045] The present invention also provides a stone CT image segmentation device based on deep learning and mixed loss, including an encoder, a decoder and a training module;

[0046] The encoder includes a 3D Unet encoder and a ViT encoder, which are used to extract local semantic features and global features of the stone CT image respectively;

[0047] The decoder includes a 3D Unet decoder and a global feature decoder, wherein the 3D Unet decoder is used to gradually restore the feature map to the original resolution and the original number of channels, and the global feature decoder is used to adjust the global feature to a suitable size for fusion with the features output at each step in the 3D Unet decoder;

[0048] The training module is used to construct a hybrid loss function by combining Dice Loss and Focal Loss, which is used as the loss function for segmentation model training. The hybrid loss function is as follows:

[0049] Loss=α*Dice Loss+β*Focal Loss

[0050] Among them, α and β are adjustment coefficients, which are used to balance the proportion of the two losses.

[0051] In order to illustrate the effect of the present invention, the present invention has done the following experiments for verification;

[0052] 1. For model training;

[0053] The present invention uses the Pytorch deep learning framework to build and train the model. During the model training process, a total of 115 iterations were performed, the batch size was set to 4, the learning rate was set to 0.0001, and the Adam gradient optimizer was used. In addition, the graphics card used in the experiment of the present invention is an NVIDIA A100. The ratio of the training set to the test set is 8:2, that is, 411 groups of data are used as the training set and 103 groups of data are used as the test set.

[0054] 2. Evaluation indicators;

[0055] The present invention uses the Dice coefficient as an evaluation index to measure the quality of the experimental results. The Dice coefficient is one of the most commonly used evaluation indicators in the field of medical image segmentation. The Dice coefficient is a function used to measure the similarity of two sets, and the calculation method is as follows:

[0056]

[0057] The meanings of A and B are as follows: Figure 3 As shown in:

[0058] 3. Comparative experiment

[0059] In order to prove the effectiveness of the method proposed in the present invention, the method proposed in the present invention is compared with other segmentation methods on the data set used in the present invention, and the results are shown in the following table. It can be seen from the experimental results that the Dice coefficient of the method proposed in the present invention is the highest in each comparative experiment, and the Dice coefficient is improved by 0.35%-0.68% compared with other segmentation methods. It can be seen from the experimental results that the segmentation model proposed in the present invention can effectively mark the location and morphology of kidney stones in CT images, thereby improving the efficiency and accuracy of diagnosing patients with kidney stones.

[0060] Table 1 Comparative experiment

[0061]

[0062] In addition, some segmentation results of different methods are shown in Figure 4 As shown in the figure, it can be seen that compared with other methods, the segmentation model proposed in the present invention can more completely segment the stone area of ​​the kidney stone CT image. And the segmentation model proposed in the present invention has fewer errors in segmenting the non-stone area into the stone area.

[0063] 4. Ablation experiment

[0064] In order to prove the effectiveness of the method proposed in the present invention, the present invention conducted an ablation experiment on the proposed framework to verify the experimental effects of ViT encoding followed by step-by-step decoding and mixed loss. The present invention conducted two groups of ablation experiments, namely, 3DUnet was trained using a mixed loss method and the segmentation model proposed in the present invention was trained using Dice loss. The results are shown in the following table. As can be seen from the table, the Dice coefficient of the method proposed in the present invention reaches 96.58%, which is 0.21%-0.27% higher than the other two groups of experiments. This shows that the segmentation model proposed in the present invention can effectively segment kidney stone CT images.

[0065] Table 2 Ablation experiment

[0066]

[0067] The specific implementation modes of the present invention are described in detail above in conjunction with the accompanying drawings, but the present invention is not limited to the above implementation modes, and various changes can be made within the knowledge scope of ordinary technicians in this field without departing from the purpose of the present invention.

Claims

1. A stone CT image segmentation method based on deep learning and hybrid loss, characterized by: The specific steps of the method are as follows: Step 1: Preprocess the data: Crop the 3D CT images and the corresponding segmentation marker data in the data set simultaneously to keep the ROI area consistent; Step 2: Send the preprocessed data to the encoder of the 3D Unet of the segmentation model, and perform step-by-step convolution and maximum pooling calculations on the input kidney stone CT image. The input of each step of calculation is the output of the previous step. Step 3: Divide the preprocessed data into blocks and flatten them into sequence vectors, and then send them to the ViT encoder of the segmentation model for self-attention calculation to obtain the global features of the stone CT image; Step 4: Send the global features obtained in Step 3 to the global feature decoder of the segmentation model decoder, and gradually perform deconvolution decoding calculations to obtain multiple feature maps with different numbers of channels; Step 5. Send the feature map obtained by the last convolution and pooling calculation in Step 2 to the 3DUnet decoder of the segmentation model decoder, and perform deconvolution calculations one by one. The input of each step is the output of the previous calculation, so as to gradually restore the feature map resolution and the original number of channels. In each deconvolution calculation, the feature map obtained by deconvolution calculation is spliced ​​with the feature map obtained after each convolution and pooling calculation in Step 2, that is, jump connection. After splicing, convolution calculation is performed again, and then the calculation results of each step are spliced ​​with the multiple feature maps obtained by the global feature decoder in Step 4 to obtain the fused feature as the input of the next deconvolution calculation.

2. The method for stone CT image segmentation based on deep learning and hybrid loss according to claim 1, characterized in that: It also includes Step 6: constructing a hybrid loss function by combining Dice Loss and Focal Loss as the loss function for segmentation model training; the hybrid loss function is shown as follows: Loss=α*Dice Loss+β*Focal Loss Among them, α and β are adjustment coefficients, which are used to balance the proportion of the two losses.

3. The method for stone CT image segmentation based on deep learning and hybrid loss according to claim 1, characterized in that: The Step 1 includes: cutting the CT image into 3D data of a uniform size of 96×96×96.

4. The method for stone CT image segmentation based on deep learning and hybrid loss according to claim 1, characterized in that: The Step 3 includes: stacking 4 layers of ViT encoders in total, wherein the number of heads of the multi-head self-attention mechanism of each encoder is 2; obtaining a total of 768-channel global feature maps of stone CT images after passing through the ViT encoder; using deconvolution to perform torch.concat(·) splicing operations on the global feature maps obtained after encoding and the features after each deconvolution operation in the segmentation model decoder 3D Unet decoder step by step. Specifically, a total of three deconvolutions are used, and the feature maps obtained by each deconvolution are torch.concat(·) splicing operations are performed on the feature maps obtained by each deconvolution operation in the segmentation model decoder 3D Unet decoder step by step from bottom to top.

5. A stone CT image segmentation device based on deep learning and mixed loss, characterized in that: Includes encoder, decoder and training modules; The encoder includes a 3D Unet encoder and a ViT encoder, which are used to extract local semantic features and global features of the stone CT image respectively; The decoder includes a 3D Unet decoder and a global feature decoder, wherein the 3D Unet decoder is used to gradually restore the feature map to the original resolution and the original number of channels, and the global feature decoder is used to adjust the global feature to a suitable size for fusion with the features output at each step in the 3D Unet decoder.

6. The device for stone CT image segmentation based on deep learning and hybrid loss according to claim 5, characterized in that: It also includes a training module, which is used to construct a hybrid loss function by combining Dice Loss and Focal Loss to serve as the loss function for segmentation model training; the hybrid loss function is as follows: Loss=α*Dice Loss+β*Focal Loss Among them, α and β are adjustment coefficients, which are used to balance the proportion of the two losses.

Images