A boundary-aware semi-supervised medical image segmentation method

Abstract

The present invention discloses a boundary-aware semi-supervised medical image segmentation method, which constructs a semi-supervised medical image segmentation model for training. The semi-supervised medical image segmentation model for training includes two sub-models with identical structures, each of which includes an encoder, a foreground decoder, and a boundary decoder. A training data set is used as input of the semi-supervised medical image segmentation model, and a loss function is constructed. Training is performed using a stochastic gradient descent optimizer until a preset number of training rounds is reached. The encoder and foreground decoder trained in one of the sub-models are used as the final medical image segmentation model. The present invention designs methods such as boundary segmentation auxiliary tasks, cross-task consistency constraints, and cross-model consistency constraints, so that the model can achieve high-accuracy image segmentation effects, especially the segmentation accuracy of boundary areas, when only a small number of labeled samples are present.

Description

Technical Field

[0001] The present invention relates to a medical image segmentation method, which is particularly useful for obtaining precise segmentation boundaries in semi-supervised medical image segmentation when labeled image data is insufficient but unlabeled data is sufficient. Background Art

[0002] Medical image segmentation applies semantic segmentation technology to medical images. Its purpose is to isolate areas of medical images with specific meanings, such as labeling tissues and organs in CT or MRI images. This provides a reliable basis for clinical diagnosis and pathological research, helping doctors make more accurate diagnoses. Automatic identification of organ and tissue regions in medical images is of great clinical significance.

[0003] Currently, deep learning-based methods have become the mainstream approach for medical image segmentation. However, training deep learning network models typically requires sufficient data manually annotated by professional doctors, which consumes a lot of time and manpower, making it difficult to implement in actual clinical practice, hindering the application of deep learning in real-world clinical practice.

[0004] Semi-supervised learning can learn from both labeled and unlabeled data, thus reducing the cost of labeling. In recent years, researchers have proposed a series of semi-supervised deep learning methods for segmenting various tissues and organs. However, existing semi-supervised medical image segmentation methods often perform poorly in the boundary regions of segmented targets and fail to fully exploit discriminative features in the image. Therefore, it is necessary to develop a new segmentation method that can reduce labeling costs while improving segmentation accuracy, especially in the boundary regions of the target. Summary of the Invention

[0005] The purpose of the present invention is to reduce the burden of annotation and provide a semi-supervised medical image segmentation method that can obtain accurate segmentation boundaries.

[0006] The objectives of the present invention are achieved through the following technical solutions:

[0007] A semi-supervised medical image segmentation method based on boundary perception is specifically as follows:

[0008] For the medical image to be segmented, Medicine to be divided Image input to the trained medical image segmentation model middle , obtain the segmentation result;

[0009] The medical image segmentation model is trained by the following method:

[0010] Medical image data is collected. Professionals perform pixel-level annotation on the foreground of the target area of some medical images as foreground labels, and then calculate boundary labels based on the foreground labels. The original images, foreground labels, and boundary labels together constitute the training dataset.

[0011] A semi-supervised medical image segmentation model for training is constructed, wherein the semi-supervised medical image segmentation model for training includes two sub-models with identical structures, each sub-model including an encoder, a foreground decoder, and a boundary decoder; for each sub-model, the encoder is used to extract a high-level semantic feature map from an input medical image, and the foreground decoder and the boundary decoder respectively output a foreground segmentation probability map and a boundary segmentation probability map based on the feature map;

[0012] The training dataset is used as input for a semi-supervised medical image segmentation model, and a loss function is constructed. The model is trained using a stochastic gradient descent optimizer until a preset number of training rounds is reached. The encoder and foreground decoder trained in one of the sub-models are used as the final medical image segmentation model.

[0013] The above loss function The specific definitions are as follows:

[0014]

[0015] in, is the supervised loss for labeled data, is the unsupervised loss for unlabeled data, and λ(t) is the weight that changes with the number of training rounds t;

[0016] For labeled samples in the training set, a supervised loss function is constructed using the true foreground labels and true boundary labels to optimize the model parameters.

[0017]

[0018]

[0019]

[0020] is the foreground segmentation loss, that is, the segmentation loss between the foreground segmentation results p1 and p2 of the two sub-models and the true foreground label y; is the boundary segmentation loss, that is, the segmentation loss between the boundary segmentation results b1 and b2 of the two sub-models and the true boundary label B; α is the weight to balance the two loss terms;

[0021] It can be one or a combination of common segmentation loss functions such as dice loss function, cross entropy loss function, etc.

[0022] For unlabeled samples in the training set, we use consistency constraints between tasks and sub-models to construct an unsupervised loss function to optimize model parameters.

[0023]

[0024]

[0025]

[0026] in, is the cross-task consistency loss function, that is, the segmentation loss between the foreground segmentation result p1(p2) and the boundary segmentation result b1(b2) in a single sub-model; is the cross-model consistency loss function, that is, the segmentation loss between the two foreground segmentation results p1 and p2 of the two sub-models and the segmentation loss between the two boundary segmentation results b1 and b2; β is the weight to balance the two loss terms; Represents the transformation from foreground segmentation probability map to boundary segmentation probability map.

[0027] Furthermore, during the training process, the training data set also includes data obtained by data augmentation based on the original data;

[0028] Furthermore, t represents the current training round number, t max Indicates the maximum number of training rounds.

[0029] Furthermore, for The present invention proposes to use pooling operations to implement the conversion process, so that the forward propagation calculation process of the entire model is differentiable, so that the two decoders of each sub-model can be collaboratively optimized by the gradient descent method; the conversion process is specifically implemented as follows:

[0030]

[0031] Furthermore, in the inference stage, the sliding window method is first used to crop the original medical image into several image blocks with the same size as the training image, and then input into the trained model to obtain several corresponding segmentation results, which are then spliced together through a process opposite to the sliding window to obtain the final image segmentation result.

[0032] The advantages of the present invention are: In response to the problems faced by automatic medical image segmentation in actual clinical applications, namely high labeling costs and low accuracy of foreground boundary areas of segmentation results, the present invention designs a novel semi-supervised deep learning model, and designs methods such as boundary segmentation auxiliary tasks, cross-task consistency constraints, and cross-model consistency constraints. This allows the model to achieve high-accuracy image segmentation effects, especially the accuracy of boundary area segmentation, when only a small number of labeled samples are available. On multiple mainstream public datasets, the accuracy of this method exceeds the current state-of-the-art methods. In addition, the model has a simple structure and strong versatility, making it suitable for different medical image data and tasks. BRIEF DESCRIPTION OF THE DRAWINGS

[0033] Figure 1 Here is an example of a medical image.

[0034] Figure 2 This is the overall structure diagram of the segmentation model in the example.

[0035] Figure 3 It is a flowchart of the training phase of the medical image segmentation model proposed in the present invention.

[0036] Figure 4 It is a flowchart of the inference stage of the medical image segmentation model proposed in this invention.

[0037] Figure 5 It is a visual display of the image segmentation results in the example. DETAILED DESCRIPTION

[0038] The present invention will be further described in detail below with reference to the accompanying drawings and specific examples.

[0039] The following uses the 2018 MICCAI left atrium segmentation challenge public dataset as an example to illustrate the specific implementation of the present invention.

[0040] First, if Figure 3 As shown, follow these steps to train a medical image segmentation model:

[0041] Step 1: Collect medical image data and annotate a certain percentage of the images to identify the target foreground area. Obtain foreground labels, and then calculate boundary labels based on the foreground labels. The original images, foreground labels, and boundary labels together constitute the training dataset. Figure 1 This is a visualization of an example of data in the dataset. From left to right, it shows the 2D display of the original image, the foreground label, the boundary label, and the 3D display of the foreground label.

[0042] The dataset of this embodiment contains 100 3D MRI images, 80 of which are randomly divided as a training set, and the remaining 20 images are used as a test set. In the training set, a certain proportion of images are randomly selected as labeled images, that is, their labels are used during the training process, while the labels of the remaining images are not used. In this embodiment, we randomly select 10% and 20% of the images as labeled images, respectively. The image labels have the same resolution as the original image, with the foreground pixel value being 1 and the background pixel value being 0. This embodiment is a binary classification task, but the method of the present invention is also applicable to multi-classification tasks.

[0043] Step 2: Perform data processing and data enhancement on the original dataset.

[0044] First, for the original dataset images, the labeled region rectangle is randomly expanded in all directions to a certain size, and the original image and label are cropped. Second, the cropped images are normalized so that each image has a mean of 0 and a variance of 1. During training, two labeled data and two unlabeled data are randomly selected from the cropped images, and image patches of size 112×112×80 are cropped. Data augmentation is performed using random rotation and random flipping; the rotation angles are 90°, 180°, or 270°, and the flipping method is along the x-axis or y-axis. The processed images are used as input for subsequent training models.

[0045] Step 3: Build a boundary-aware semi-supervised medical image segmentation model.

[0046] The specific structure of the boundary-aware semi-supervised medical image segmentation model used for training is as follows Figure 1 The entire model consists of two sub-models, each of which contains an encoder, a foreground decoder, and a boundary decoder. In each sub-model, the encoder extracts high-level semantic feature maps from the input medical image, and the foreground decoder and boundary decoder output foreground segmentation probability maps and boundary segmentation probability maps based on these feature maps, respectively.

[0047] The encoder and decoder in this model do not rely on a specific network structure, and mainstream networks such as 3D U-Net and V-Net can be used as the backbone network of this model. In this embodiment, the encoder and decoder of the sub-model use V-Net as the backbone network. Among them, an encoder is composed of four convolution modules and four downsampling modules, each convolution module is composed of a convolution layer, a batch normalization (BatchNorm) layer and a ReLU activation function, and each downsampling module is obtained by combining a convolution layer and a ReLU activation function. The decoder is composed of four convolution modules, four upsampling modules and an output module. The convolution module structure is consistent with the convolution module in the encoder. The upsampling module is obtained by combining a transposed convolution and a ReLU activation function. The output layer is a convolution layer, which maps the high-dimensional feature map to a two-channel segmentation probability map. The output layer of the foreground decoder outputs a foreground segmentation probability map, and the output layer of the boundary decoder outputs a boundary segmentation probability map.

[0048] Step 4: Model training. Each batch of four data, two labeled data and two unlabeled data, is used to optimize the model, which is counted as one training round. Use the stochastic gradient descent optimizer with a learning rate of 0.01 and train for 15,000 rounds. As the loss function of the entire model.

[0049]

[0050] in, is the supervised loss for labeled data, is the unsupervised loss of unlabeled data, λ(t) is the weight that changes with the number of training rounds t, The specific calculation formula is shown in (2). The specific calculation formula of is shown in (5), and the calculation formula of λ is shown in (9).

[0051]

[0052] Supervised loss uses the true labels to optimize the model parameters and is only calculated on labeled samples. It includes two parts: foreground segmentation loss and boundary segmentation loss. is the foreground segmentation loss, that is, the segmentation loss between the foreground segmentation results p1 and p2 of the two sub-models and the true foreground label y; is the boundary segmentation loss, that is, the segmentation loss between the boundary segmentation results b1 and b2 of the two sub-models and the true boundary label B; α is the weight for balancing the two loss terms, which needs to be adjusted according to the specific task. In this embodiment, α=1; and The specific calculation formulas are shown in (3) and (4) respectively.

[0053]

[0054]

[0055] In formulas (3) and (4), A common loss function can be used. This embodiment uses a combination of the dice loss function and the cross entropy loss function, that is,

[0056] For unlabeled samples in the training set, an unsupervised loss function is constructed using consistency constraints between tasks and sub-models. This optimizes the model parameters. The specific calculation formula is shown in formula (5).

[0057]

[0058] in, is the cross-task consistency loss function, that is, the segmentation loss between the foreground segmentation result p1(p2) and the boundary segmentation result b1(b2) in a single sub-model. The specific calculation formula is shown in formula (6); is the cross-model consistency loss function, that is, the segmentation loss between the two foreground segmentation results p1 and p2 of the two sub-models and the segmentation loss between the two boundary segmentation results b1 and b2. The specific calculation formula is shown in formula (7); β is the weight for balancing the two loss terms, which needs to be adjusted according to the specific task. In this embodiment, β = 1.

[0059]

[0060]

[0061] In formula (6), Represents the transformation from the foreground segmentation probability map to the boundary segmentation probability map. The method proposes to use pooling operation to realize the conversion process, so that the forward propagation calculation process of the entire model is differentiable, so that the two decoders of each sub-model can be collaboratively optimized by gradient descent method; the specific implementation of the conversion process is shown in formula (8):

[0062]

[0063] Furthermore, considering that the model accuracy is low in the early stages of training and cross-task consistency and cross-model consistency are unreliable, we adopt the method of controlling λ to avoid the interference of the unsupervised loss term on early training. Specifically, the calculation formula of λ changing with argument t is shown in formula (9).

[0064]

[0065] Among them, t max is the maximum number of training rounds. In this embodiment, t max is 15,000.

[0066] The model is trained according to the above settings until the preset number of training rounds is reached. After training is complete, the trained set of encoders and foreground decoders is used as the medical image segmentation model for subsequent inference.

[0067] Inference stage, such as Figure 4 As shown in the figure, the sliding window method is first used to crop the original image to be tested into several image blocks with the same size as the training image, and then input them into the medical image segmentation model obtained through the above training process to obtain several image segmentation results, which are then spliced together to obtain the final segmentation result.

[0068] In this embodiment, the test image is first standardized, and then a sliding window with a step size of 18×18×4 is used to crop the test image into several image blocks of 112×112×80. The trained segmentation model is input to obtain several corresponding foreground segmentation probability maps, and finally the obtained segmentation probability maps are spliced to obtain the final segmentation result.

[0069] Table 1 Comparison of the segmentation effects of the present invention and the existing algorithm when the label ratio is 10% and 20% respectively (↑ indicates that the larger the value of the indicator, the better; ↓ indicates that the smaller the value of the indicator, the better)

[0070]

[0071] As can be seen from the comparison of various indicators in Table 1, even with only a small number of annotations, the present invention improves the accuracy of medical image segmentation compared to existing technologies, and the segmentation results have more accurate boundaries and better integrity. By designing auxiliary tasks and leveraging cross-task and cross-model consistency, the present invention enables the model to pay more attention to boundary areas and learn more discriminative features during training, thereby improving segmentation accuracy.

[0072] Figure 5 This is a visualization of test data and its segmentation results. As you can see, the segmentation results obtained by this method are closer to the true labels. In particular, the results of this method are more accurate and complete in the boundary areas marked by circles.

[0073] So far, the specific implementation process of the present invention has been described.

[0074] The above embodiments are merely examples for the purpose of clarity of explanation and are not intended to limit the implementation methods. Based on the above description, the present invention may be modified in other forms according to different practical tasks.

Claims

1. A boundary-aware semi-supervised medical image segmentation method, characterized in that: The method is specifically: For the medical image to be segmented, the medical image to be segmented is input into a trained medical image segmentation model to obtain a segmentation result; The medical image segmentation model is trained by the following method: Medical image data is collected. Professionals perform pixel-level annotation on the foreground of the target area of some medical images as foreground labels, and calculate boundary labels based on the foreground labels. The original images, foreground labels, and boundary labels together constitute the training dataset. A semi-supervised medical image segmentation model for training is constructed, wherein the semi-supervised medical image segmentation model for training includes two sub-models with identical structures, each sub-model including an encoder, a foreground decoder, and a boundary decoder; for each sub-model, the encoder is used to extract a high-level semantic feature map from an input medical image, and the foreground decoder and the boundary decoder respectively output a foreground segmentation probability map and a boundary segmentation probability map based on the feature map; Use the training dataset as input to the semi-supervised medical image segmentation model and construct the loss function , using stochastic gradient descent optimizer for training until the preset number of training rounds is reached; the trained encoder and foreground decoder in one of the sub-models are used as the final medical image segmentation model; For unlabeled samples in the training set, we use consistency constraints between tasks and sub-models to construct an unsupervised loss function to optimize model parameters. in, is the cross-task consistency loss function, that is, the segmentation loss between the foreground segmentation result p1(p2) and the boundary segmentation result b1(b2) in a single sub-model; is the cross-model consistency loss function, that is, the segmentation loss between the two foreground segmentation results p1 and p2 of the two sub-models and the segmentation loss between the two boundary segmentation results b1 and b2; β is the weight to balance the two loss terms; Represents the conversion of foreground segmentation probability map to boundary segmentation probability map; for The calculation process of is implemented by using pooling operation, so that the forward propagation calculation process of the entire model is differentiable, so that the two decoders of each sub-model are collaboratively optimized through the gradient descent method; the specific implementation of the conversion process is as follows:

2. The boundary-aware semi-supervised medical image segmentation method according to claim 1, characterized in that: The above constructed loss function The specific definitions are as follows: in, is the supervised loss for labeled data, is the unsupervised loss for unlabeled data, and λ(t) is the weight that changes with the number of training rounds t; For labeled samples in the training set, a supervised loss function is constructed using the true foreground labels and true boundary labels to optimize the model parameters. is the foreground segmentation loss, that is, the segmentation loss between the foreground segmentation results p1 and p2 of the two sub-models and the true foreground label y; is the boundary segmentation loss, that is, the segmentation loss between the boundary segmentation results b1 and b2 of the two sub-models and the true boundary label B; α is the weight to balance the two loss terms; is one or a combination of common segmentation loss functions such as dice loss function, cross entropy loss function, etc.

3. The boundary-aware semi-supervised medical image segmentation method according to claim 1, characterized in that: During the training process, the training data set also includes data obtained by data enhancement based on the original data.

4. The boundary-aware semi-supervised medical image segmentation method according to claim 1, characterized in that: t represents the current training round number, t max Indicates the maximum number of training rounds.

5. The boundary-aware semi-supervised medical image segmentation method according to claim 1, characterized in that: In the inference stage, the sliding window method is first used to crop the original medical image into several image blocks with the same size as the training image, and then input them into the trained model to obtain several corresponding segmentation results. Then, the image is spliced together through the opposite process of the sliding window to obtain the final image segmentation result.

Images