Attention-guided nonlinear perturbation consistency semi-supervised medical image segmentation method

Abstract

The present invention belongs to the field of medical image processing technology, specifically a semi-supervised medical image segmentation method with attention-guided nonlinear perturbation consistency. The present invention uses a small amount of labeled data and unlabeled data to train a model for medical image segmentation, adopts nonlinear transformation for data enhancement, and constrains it through consistency of output level and feature level; nonlinear transformation realizes contrast change by randomly generating Bezier curves for grayscale mapping; consistency constraint of output level encourages generation of similar prediction labels under different enhancements through cross-pseudo-label supervision; consistency constraint of feature level encourages network learning of robust representation that is invariant to nonlinear changes through multi-scale attention-aware feature alignment. The present invention effectively utilizes unlabeled data, improves generalization through data enhancement, and improves the accuracy of medical image segmentation through information mining of unlabeled data. The present invention is applicable to all grayscale-based medical image segmentation tasks.

Description

Technical Field

[0001] The present invention belongs to the technical field of medical image processing, and in particular relates to a medical image segmentation method. Background Art

[0002] Medical image segmentation is a crucial task in medical image analysis and clinical applications. With the rapid development of artificial intelligence, computer-assisted automatic segmentation algorithms based on deep learning methods are constantly emerging. However, many outstanding achievements are inseparable from large amounts of labeled data. In recent years, the use of semi-supervised learning methods to alleviate the data scarcity problem faced by deep learning scenarios has been widely recognized. These methods only require a small amount of labeled data supplemented by a large amount of unlabeled data for learning. Consistency regularization is a mainstream algorithm in semi-supervised learning. It improves the generalization and robustness of segmentation algorithms by encouraging consistency in results under different perturbations at the data, feature, or model levels.

[0003] The key to the consistency regularization algorithm lies in the construction of effective perturbations and the calculation of stable consistency losses. In terms of perturbation construction, existing methods fail to perform powerful transformations based on the characteristics of the medical image itself; and in the construction of consistency losses, existing methods only focus on the construction of perturbations at the output level. However, the pseudo-labels generated by unlabeled data may introduce bias and fail to fully exploit the inherent information of unlabeled data. In addition, in medical business scenarios, there are also situations where medical data sources are rich in sources and data from different sources differ. For example, data from different devices and centers will also have certain differences in image quality. Existing methods are difficult to meet the needs of limited annotations and large amounts of unlabeled data scenarios in medical image segmentation. Therefore, it is of great significance to study accurate semi-supervised medical segmentation methods that can adapt to different medical image data with high generalization and stability. Summary of the Invention

[0004] The purpose of the present invention is to provide an efficient semi-supervised medical image segmentation method that can utilize a large amount of unlabeled data and limited labeled data to achieve higher segmentation accuracy and generalization while reducing the need for labeling.

[0005] The semi-supervised medical image segmentation method proposed in this paper is an attention-guided nonlinear perturbation consistency method. The segmentation framework includes a nonlinear transformation module and an attention-guided multi-scale feature consistency module. The nonlinear transformation module uses randomly generated Bezier curves to generate multiple sets of grayscale mapping functions. For each input image, the module uses monotonically decreasing and increasing mapping functions to generate strongly enhanced and weakly enhanced images, respectively. This module adjusts the image contrast to increase the sample information in the training set, helping the model learn a diverse grayscale distribution and reducing the risk of model overfitting. In the attention-guided multi-scale feature consistency module, a distance map is calculated for the output of each decoder layer. Operations at multiple scales can integrate rich information from different scales to help improve segmentation results. In addition, to make the network's representation learning more robust to different enhancement transformations, the present invention uses a spatial attention module to increase attention to the target area. Spatial attention is used to reweight distances to help the network learn more valuable and stable feature representations. With the improvement of generalization and network representation capabilities achieved by these two modules, combined with label supervision and pseudo-label cross-supervision at the output level, a more accurate segmentation network can be obtained.

[0006] Specifically, the attention-guided nonlinear transformation consistency semi-supervised medical image segmentation method of the present invention has the following process: Figure 1 The specific steps are as follows.

[0007] Step 1, image preprocessing stage, this step is the normal default operation and is not Figure 1 Draw in

[0008] For two-dimensional images, slice-level intensity normalization preprocessing is performed, and input size normalization preprocessing is performed through resizing. For three-dimensional images, cubic spline interpolation is used for resampling and the image is changed to a specified size. When the image is large, if necessary, the region of interest can be cropped using prior knowledge (expert outlining, pre-trained models). Simple data augmentation is performed during the data reading process, mainly including random transformations such as random rotation, shearing, scaling, and flipping operations.

[0009] Step 2, perform nonlinear transformation on the input;

[0010] Data is read through PyTorch's built-in Dataloader, and Sampler is used to set the sampling within a batch, with half of the batch being labeled data and half being unlabeled data, which serve as the input for the subsequent two branches. The data is then subjected to two different enhancements, strong and weak, through a nonlinear transformation module. The Bezier curves corresponding to strong and weak enhancements are randomly generated monotonically decreasing and monotonically increasing mapping functions, respectively. Different Bezier curves achieve smooth and variable curves by randomly generating two control points to produce intensity nonlinear transformation functions with diverse local contrast, which can simulate the differences between different imaging images, making this method adaptable to a variety of medical data.

[0011] A Bezier curve can be generated by two endpoints P0 and P3 and two control points P1 and P2. The generation formula is as follows:

[0012]

[0013] Where t is a fractional value of the length along the line; the relationship between the original image and the transformed image obtained using different Bezier curves is as follows Figure 2 As shown in .

[0014] Step 3: Perform attention-guided multi-scale feature alignment;

[0015] In order to constrain the strong perturbation in step 2, the weakly enhanced and strongly enhanced images generated in step 2 are respectively input into two parallel and identical branch U-Net networks. Using feature alignment at the decoding layer can encourage the network to learn valuable high-level semantic features that are invariant under nonlinear transformations; at the decoding layer, the features output by each layer of the network are built consistent on the two branch networks:

[0016] First, the L1 distance is used to calculate the feature distance map, and then the spatial attention map of each layer is constructed using the branch corresponding to the weakly enhanced input. The feature distance map is reweighted and averaged through the attention map to calculate the consistency loss between the two branches in each layer. The calculation process is as follows: Figure 3 As shown in Figure 2, averaging the feature consistency losses at different scales can obtain the attention-guided multi-scale feature consistency loss, which is used as the loss sub-item of the gradient backpropagation during network learning. The formula for defining the attention-guided multi-scale feature consistency loss is as follows:

[0017]

[0018] in, and are the feature outputs of the kth layer of the weak enhancement branch and the strong enhancement branch, H and W are the height and width of the feature of the layer, a ij is the spatial attention map, and S is the total number of decoding layers.

[0019] Step 4: Perform output cross-pseudo-label learning;

[0020] The decoder of the U-Net network can obtain a predicted output mask of the same size as the input; for labeled data, the cross entropy loss and Dice loss are calculated using its true annotation:

[0021]

[0022] For unlabeled data, the segmentation results predicted by the network can be used as pseudo labels for unlabeled data to help learning unlabeled data. The predicted labels generated by the two branches are used as each other's pseudo labels to calculate the Dice loss, which is calculated as follows:

[0023]

[0024] in, represents the prediction of the weakly enhanced image, Represents the segmentation prediction of the strongly enhanced image; the segmentation mask can be generated as a pseudo label through argmax.

[0025] Step 5: Calculate the total training loss;

[0026] For the loss calculation of a single iteration, a batch loss function contains supervised loss terms and unsupervised loss terms, and the definition formula is as follows:

[0027]

[0028] Among them, b l and b u They represent the labeled and unlabeled data subsets in the batch respectively.

[0029] Step 6, iterative optimization;

[0030] Repeat steps 2, 3, 4, and 5, and use the SGD optimization algorithm for iterative optimization. As the loss value decreases, the network weights are updated to train the segmentation model. The segmentation model parameters are saved based on the optimal results of the validation set for subsequent inference.

[0031] Step 7, test image;

[0032] After simple preprocessing without data enhancement, the test image is input into the pre-trained model saved by the weak enhancement branch to obtain the complete segmentation result.

[0033] Compared with the existing technology, the attention-guided nonlinear transformation consistency semi-supervised medical image segmentation of the present invention has the following advantages:

[0034] (1) Through nonlinear transformation, powerful perturbations can be constructed to fully learn image grayscale information and improve generalization;

[0035] (2) An attention-based multi-scale feature consistency method promotes the network to learn robust features that are invariant to transformations and achieves efficient use of unsupervised data;

[0036] (3) This paper achieves the best semi-supervised medical image segmentation results on the public ACDC dataset. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] Figure 1 This is the flowchart of the present invention, which includes a nonlinear transformation module and attention-guided multi-scale feature consistency.

[0038] Figure 2 Schematic diagram of different Bezier curves and their corresponding transformation results in the nonlinear transformation module of the present invention.

[0039] Figure 3 This is the calculation process of the single-layer attention-guided feature consistency loss in the attention-guided multi-scale feature consistency module in the present invention.

[0040] Figure 4 The visualization results of the proposed method for the MR image cardiac segmentation task on the ACDC dataset at different ratios of labeled data to unlabeled data are shown in the bottom row. DETAILED DESCRIPTION

[0041] Example 1 The method of the present invention is used to perform ACDC cardiac MR multi-region segmentation.

[0042] The semi-supervised medical image segmentation method based on attention-guided nonlinear transformation consistency provided by the present invention is end-to-end. The specific implementation process of the embodiment is as follows:

[0043] Step 1: The ACDC dataset contains a total of 100 data examples, including 70 training set data, 10 validation set data, and 20 test set data. First, the image intensity is normalized to [0, 1] and the image is resized to 256×256. Random transformations including random rotation, shearing, scaling, and flipping operations are used in each training iteration to enhance the data.

[0044] Step 2: Input the image into the nonlinear transformation module to obtain strongly enhanced and weakly enhanced images, which are then input into two identical U-Net networks respectively; the network encoder and decoder contain 4 layers of downsampling and upsampling layers respectively, and the network structure is U-shaped.

[0045] Step 3: Calculate the attention-guided multi-scale feature consistency loss for the four-layer network output of the two-branch decoder.

[0046] Step 4: Get the prediction result of the network output and calculate the total loss of a single iteration.

[0047] Step 5: Iterate and optimize to obtain the segmentation network. The network uses the SGD optimizer with the parameters weight decay = 0.0001 and momentum = 0.9. The initial learning rate is 0.01, the batch size is 16, and the number of iterations is 30k. As the loss value decreases, the highest value of the validation set prediction result is used to obtain the trained network parameters.

[0048] In step 6, the test set is input into the weak enhancement branch and the segmentation results are predicted without data enhancement. The average results of multiple segmentation regions with different annotation ratios are shown in Table 1. Among them, LS represents training using only available annotated data, and Seg represents training using To do supervision, Pseu stands for using cross pseudo label supervision NLT stands for the use of nonlinear transformation modules, and AMFC stands for the use of attention-guided multi-scale consistency modules. The results show the effectiveness of each design of the present invention, and Dice values ​​of 82.81 and 87.19 are obtained when only 5% and 10% of the training data are used as labeled data.

[0049] Table 1

[0050]

[0051] The experimental results of this embodiment show that the method of the present invention has higher segmentation accuracy and better generalization; the nonlinear transformation and attention-guided multi-scale feature consistency proposed in the present invention can help the network to extract features more effectively and make full use of unlabeled data to help improve network performance; and the method is simple to implement and can be easily integrated into various dual-branch semi-supervised segmentation frameworks.

[0052] References:

[0053] [1]XIA Y, LIU F, YANG D, et al. 3D Semi-Supervised Learning withUncertainty-Aware Multi-View Co-Training[C / OL] / / 2020IEEE Winter Conference onApplications of Computer Vision(WACV).Snowmass Village,CO,USA:IEEE,2020:3635-3644[2022-09-05].

[0054] https: / / ieeexplore.ieee.org / document / 9093608 / .

[0055] [2]CAO X, CHEN H, LI Y, et al. Uncertainty Aware Temporal-Ensembling Modelfor Semi-Supervised ABUS Mass Segmentation[J / OL]. IEEE Transactions on MedicalImaging, 2021, 40(1): 431-443. https: / / doi.org / 10.1109 / TMI.2020.3029161.

[0056] [3]LUO X, WANG G, LIAO W, et al. Semi-supervised medical image segmentationvia uncertainty rectified pyramid consistency[J / OL]. Medical Image Analysis, 2022, 80: 102517.

[0057] https: / / doi.org / 10.1016 / j.media.2022.102517.

[0058] [4]LI S, ZHANG C, HE X. Shape-aware Semi-supervised 3D SemanticSegmentation for Medical Images: Volume 12261[M / OL]. 2020: 552-561[2022-07-14]. http: / / arxiv.org / abs / 2007.10732.

[0059] [5]LITJENS G, KOOI T, BEJNORDI B E, et al. A survey on deep learning inmedical image analysis[J / OL]. Medical Image Analysis, 2017, 42: 60-88.

[0060] https: / / doi.org / 10.1016 / j.media.2017.07.005.

[0061] [6] BASAK H, BHATTACHARYA R, HUSSAIN R, et al. An Embarrassingly Simple Consistency Regularization Method for Semi-Supervised Medical Image Segmentation[M / OL]. arXiv, 2022[2022-07-06]. http: / / arxiv.org / abs / 2202.00677.

[0062] [7] LUO X, CHEN J, SONG T, et al. Semi-supervised Medical Image Segmentation through Dual-task Consistency[M / OL]. arXiv, 2021[2022-06-25]. http: / / arxiv.org / abs / 2009.04448.

[0063] [8] LUO X, HU M, SONG T, et al. Semi-Supervised Medical Image Segmentation via Cross Teaching between CNN and Transformer[J]. arXiv preprint arXiv:2112.04894, 2021.

[0064] [9] LI X, YU L, CHEN H, et al. Semi-supervised Skin Lesion Segmentation via Transformation Consistent Self-ensembling Model[M / OL]. arXiv, 2018[2022-06-18].

[0065] http: / / arxiv.org / abs / 1808.03887.

[0066]

[10] LI X, YU L, CHEN H, et al. Transformation Consistent Self-ensembling Model for Semi-supervised Medical Image Segmentation[M / OL]. arXiv, 2020[2022-06-18].

[0067] http: / / arxiv.org / abs / 1903.00348.

[0068]

[11] XIE Q, DAI Z, HOVY E, etc. Unsupervised data augmentation for consistency training[J]. Advances in Neural Information Processing Systems, 2020, 33: 6256-6268.

Claims

1. An attention-guided nonlinear perturbation consistency semi-supervised medical image segmentation method, characterized by: A nonlinear transformation module and an attention-guided multi-scale feature consistency module are used in segmentation. The nonlinear transformation module uses randomly generated Bezier curves to generate multiple sets of grayscale mapping functions, and uses monotonically decreasing and increasing mapping functions to generate strongly enhanced and weakly enhanced images for each input image. The nonlinear transformation module increases the sample information of the training set by adjusting the contrast of the image, helping the model learn a diverse grayscale distribution to reduce the risk of model overfitting. In the attention-guided multi-scale feature consistency module, a distance map is calculated for the output of each layer of the decoder. Operations at multiple scales can integrate rich information at different scales to help improve the segmentation effect. In order to make the network's representation learning more robust in different enhancement transformations, a spatial attention module is used to increase attention to the target area, and spatial attention is used to reweight the distance to help the network learn more valuable stable feature representations. With the improvement of generalization and network representation capabilities by these two modules, coupled with label supervision and pseudo-label cross-supervision at the output level, a more accurate segmentation network is obtained.

2. The medical image segmentation method according to claim 1, characterized in that: The specific steps are: Step 1, image preprocessing, includes: For two-dimensional images, perform slice-level intensity normalization preprocessing and input size normalization preprocessing by resizing; For three-dimensional images, cubic spline interpolation is used to resample the image and change the image to the specified size; When the image is large, the region of interest is cropped using prior knowledge; Perform data augmentation during data reading, including random transformations such as random rotation, shearing, scaling, and flipping operations; Step 2, perform nonlinear transformation on the input; Data is read through PyTorch's built-in Dataloader, and Sampler is used to set the sampling within a batch, with half of the batch being labeled data and half being unlabeled data, which serves as the input for the subsequent two branches. The data is then subjected to two different enhancements, strong and weak, through a nonlinear transformation module. The Bezier curves corresponding to strong and weak enhancements are randomly generated monotonically decreasing and monotonically increasing mapping functions, respectively. Different Bezier curves achieve smooth and variable curves by randomly generating two control points, thereby producing intensity nonlinear transformation functions with diverse local contrasts, simulating the differences between different imaging. The Bezier curve is generated by two end points P0 and P3 and two control points P1 and P2. The generation formula is as follows: Where t is a fractional value along the length of the line; Step 3: Perform attention-guided multi-scale feature alignment; The weakly enhanced and strongly enhanced images generated in step 2 are fed into two parallel, identical branch U-Net networks. Using feature alignment at the decoding layer can encourage the network to learn valuable high-level semantic features that are invariant under nonlinear transformations. At the decoding layer, the features output by each layer of the network are built into the consistency of the two branch networks: first, the L1 distance is used to calculate the feature distance map, and then the spatial attention map of each layer is constructed using the branch corresponding to the weakly enhanced input. The feature distance map is reweighted and averaged through the attention map, and the consistency loss between the features of each layer of the two branches is calculated; the feature consistency loss at different scales is averaged to obtain the attention-guided multi-scale feature consistency loss, which is used as the loss sub-item of the gradient return during network learning; the formula for the attention-guided multi-scale feature consistency loss is defined as follows: in, and are the feature outputs of the kth layer of the weak enhancement branch and the strong enhancement branch, H and W are the height and width of the feature of the layer, a ij is the spatial attention map, S is the total number of decoding layers; Step 4: Perform output cross-pseudo-label learning; The decoder of the U-Net network can obtain a predicted output mask of the same size as the input; For labeled data, use its true annotation to calculate the cross entropy loss and Dice loss: For unlabeled data, the segmentation results predicted by the network are used as pseudo labels for the unlabeled data to help the learning of unlabeled data. The predicted labels generated by the two branches are used as each other's pseudo labels to calculate the Dice loss, which is calculated as follows: in, represents the prediction of the weakly enhanced image, Represents the segmentation prediction of the strongly enhanced image; generates the segmentation mask as a pseudo label through argmax; Step 5: Calculate the total training loss; For the loss calculation of a single iteration, a batch loss function contains supervised loss terms and unsupervised loss terms, and the definition formula is as follows: Among them, b l and b u Represent the labeled and unlabeled data subsets in the batch respectively; Step 6, iterative optimization; Repeat steps 2, 3, 4, and 5, using the SGD optimization algorithm for iterative optimization. Update the network weights as the loss value decreases, train the segmentation model, and save the segmentation model parameters based on the optimal result of the validation set for subsequent inference. Step 7, test image; After preprocessing without data enhancement, the test image is input into the pre-trained model saved by the weak enhancement branch to obtain the complete segmentation result.

Images