A Medical Image Segmentation Method Based on Residual Axial Attention

Abstract

The present invention provides a medical image segmentation method based on residual axial attention. First, to address the problems of blurred boundaries and complex variations in the lesion parts of medical images, a channel-spatial transformer based on residual axial attention is designed. This module fuses multi-scale features from multiple stages in the encoder, effectively enhancing the captured deep features. Then, to address the semantic gap between the deep features obtained by the transformer and the upsampled features, a semantic alignment module based on residual axial attention is designed, using the shallow features extracted by the decoder to guide the semantic information alignment of the deep features obtained by the transformer. Finally, a loss function combined with edge prediction is designed, taking the regional information as the geometric constraint for the image segmentation task to encourage the model to generate clear lesion boundaries. The present invention utilizes the transformer based on residual axial attention and the semantic alignment module, combined with the edge prediction loss function, to achieve fine segmentation of medical images.

Description

Technical Field

[0001] The present invention relates to image processing technology, and more particularly to a medical image segmentation method based on residual axial attention. Technical Background

[0002] Medical image segmentation is a crucial task in medical image processing, which aims to accurately extract the structures or lesion regions of interest from medical images. The rapid development of deep learning technology has brought a revolutionary change to medical image segmentation, making it more efficient and accurate to extract meaningful information from complex medical images. Image segmentation methods based on deep learning can not only automatically learn the feature representations in images, but also gradually extract the abstract features in images through multi-layer convolution and pooling operations. At the same time, according to the label information in the training data, learn the feature representations of different structures and tissues in the image, so as to achieve accurate segmentation of medical images.

[0003] U-Net is the most widely used encoder-decoder network architecture in the field of medical image segmentation, which can learn from fewer training images and produce more accurate segmentation results. On this basis, networks such as UNet++, AttUNet and V-Net are improved by introducing attention mechanisms and multi-scale feature fusion, etc., to segment images of different medical imaging modalities. Although these methods perform well, due to the locality of convolutional operations, it is difficult for pure CNN architectures to fully capture global context relationships.

[0004] In this context, attention mechanisms can help the model better focus on important regions in the image, thereby improving the accuracy and efficiency of segmentation, providing new possibilities for medical image analysis. The Vision Transformer (ViT) based on self-attention mechanism can effectively capture the long-range dependencies between pixels. However, ViT requires a large amount of training data to obtain good performance, and has a high computational complexity, especially for high-resolution input images. Swin Transformer reduces the computational complexity by introducing a hierarchical mechanism, but due to the limited information interaction between windows, the receptive field is small. Most Transformer-based methods need to be pre-trained on large datasets to achieve satisfactory results, and perform poorly on small sample or imbalanced datasets, suffering from the problem of insufficient generalization ability. In the field of medical image segmentation, it is very expensive to obtain large datasets and labels with high-quality images. In addition, compared with CNN, pure Transformer performs poorly in capturing local representations.

[0005] Therefore, to address the problems existing in the above-mentioned medical image segmentation, some methods attempt to combine CNN and Transformer to enhance the network's capabilities. However, introducing the attention mechanism increases the computational complexity of the model, and there is still a balance problem between performance and computational cost. Therefore, designing an efficient and accurate medical image segmentation model is of great research significance. Summary of the Invention

[0006] The object of the present invention is to solve the problem of the semantic gap between deep features and shallow features in medical image segmentation, and provide a medical image segmentation method based on residual axial attention. Through a deep learning network, medical image features are effectively extracted and fused to achieve precise segmentation of the lesion part of the medical image.

[0007] To achieve the above object, the present invention provides a medical image segmentation method based on residual axial attention, which mainly includes the following five parts: The first part is to preprocess the medical image dataset; the second part is to extract shallow features from the medical image; the third part is to extract deep features from the medical image after extracting shallow features; the fourth part is to perform semantic alignment on the shallow features and deep features extracted from the medical image; the fifth part is the training and testing of the medical image segmentation network model based on residual axial attention, and finally obtain the segmented result map, specifically:

[0008] The first part includes two steps:

[0009] Step 1, download the public medical datasets GlaS and MoNuSeg. The original images in the dataset GlaS have different sizes, and most of them are RGB images with a size of 775

[0010] ×522, and the original images in the dataset MoNuSeg are all 1000×1000; then uniformly set the sample resolution of the datasets GlaS and MoNuSeg to 224×224;

[0011] Step 2, perform two operations of random flipping and random rotation on the images processed in Step 1 to enhance the training set samples and avoid overfitting, and then form the final required training set samples;

[0012]

[0013] The second part includes two steps:

[0014] Step 3: Process the training set samples obtained in Step 2 through 4 convolutional layers, and insert a max pooling layer and a ReLU activation function into each convolutional layer. The dimension of the convolutional kernel of each convolutional layer is 3×3, and the shallow semantic features E1, E2, E3, and E4 of the input image I are initially obtained; and reshape E1, E2, E3, and E4 into token sequences T1, T2, T3, and T4 with sizes of P, P / 2, P / 4, and P / 8 respectively, where the channel size is all changed to 128;

[0015] Step 4: Use the shallow semantic feature E4 obtained in Step 3 as the input of the dilated convolutional module DCM to obtain the output feature E5 containing more information, where DCM contains 5 convolutional kernels of 3×3, and the expansion parameters are 1, 2, 4, 8, and 16 in sequence, and the receptive field sizes are

[0016] N = 3×3, 7×7, 15×15, 31×31, and 63×63;

[0017] The third part includes three steps:

[0018] Step 5: Use the token sequences T1, T2, T3, and T4 of the shallow semantic features obtained in Step 3 as the input, and adopt the residual axial attention-based channel attention module RA-CAM to learn the correlation between different channels, and obtain the deep features T c1 、T c2 、T c3 、T c4 ;

[0019] Step 6: Use the deep features T c1 、T c2 、T c3 、T c4 obtained in Step 5 as the input, and adopt the residual axial attention-based spatial attention module RA-SAM to learn the importance of different spatial positions, and obtain the deep features T s1 、T s2 、T s3 、T s4 ;

[0020] Step 7: Perform LayerNorm layer normalization and MLP multi-layer perceptron on T s1 、T s2 、T s3 、T s4 in sequence to obtain the enhanced deep semantic features T

[0021] features T M1 、T M2 、T M3 、T M4 ,and then Ts1 , T s2 , T s3 , T s4 As a residual term, it improves the expression ability of the network;

[0022] The fourth part includes one step:

[0023] Step 8, taking the deep semantic features T M1 , T M2 , T M3 , T M4 obtained in step 7 as the input, and performing semantic

[0024] feature alignment operations using the semantic alignment module SRA to obtain the aligned features O1, O2, O3, O4;

[0025] The fifth part includes three steps:

[0026] Step 9, in the loss function, first simulate the dynamic model to construct the loss function Loss ACM , Loss ACM includes two parts, the length term Length and the region term Region. The length term Length constrains the segmentation boundary contour to make the contour line relatively smooth, and the region term Region maintains the shape based on the variance of the regions inside and outside the boundary; then for the problem that the lesion part in medical images is usually much smaller than the background part, the total loss function of the segmentation network Loss = Loss ACM + α·Loss Dice + β·Loss BCE ;

[0027] Step 10, inputting the training set samples in step 2 into the network from step 3 to step 9, setting the network hyperparameters: the learning rate is 0.001, the number of iterations adopts an early stopping strategy so as not to be fixed, the batch size is 4, the optimizer is Adam, using five-fold cross-validation, the loss function is Loss, and training the network to obtain the final pre-trained model for medical image segmentation;

[0028] Step 11, inputting the public test set into the pre-trained model obtained in step 10, and the network can segment the lesion part of the medical image.

[0029] The present invention provides a medical image segmentation method based on residual axial attention. First, to address the problems of blurred boundaries and complex variations in the lesion parts of medical images, a channel - spatial transformer based on residual axial attention is designed. This module fuses multi - scale features from multiple stages in the encoder, which can effectively enhance the captured deep features. Then, aiming at the semantic gap between the deep features obtained by the transformer and the up - sampled features, a semantic alignment module based on residual axial attention is designed, which uses the shallow features extracted by the decoder to guide the semantic information alignment of the deep features obtained by the transformer. Finally, a loss function combined with edge prediction is designed, taking the regional information as the geometric constraint for the image segmentation task to encourage the model to generate clear lesion boundaries. The present invention can achieve fine segmentation of medical images by using the transformer based on residual axial attention, the semantic alignment module, and combining the edge prediction loss function. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 is the overall network framework diagram of the present invention;

[0031] Figure 2 is the dilated convolution module DCM of the present invention;

[0032] Figure 3 is the residual axial attention RA - Attention of the present invention;

[0033] Figure 4 is the channel - spatial transformer RA - CST based on residual axial attention of the present invention;

[0034] Figure 5 is the semantic alignment module SRA based on residual axial attention;

[0035] Figure 6 are some medical images on the public test set;

[0036] Figure 7 are the results after processing using the present invention Figure 6 DETAILED DESCRIPTION OF THE EMBODIMENTS

[0037] To better understand the present invention, the medical image segmentation method based on residual axial attention of the present invention will be described in more detail below in conjunction with specific embodiments. In the following descriptions, the detailed descriptions of the currently existing technologies may dilute the subject matter of the present invention, and these descriptions will be ignored here.

[0038] ​Step 1: Download the publicly available medical datasets GlaS and MoNuSeg. The original images in the GlaS dataset have different sizes, and most of them are RGB images with a size of 775×522. The original images in the MoNuSeg dataset are all 1000×1000 in size. Then, uniformly set the sample resolution of the GlaS and MoNuSeg datasets to 224×224;

[0039] Step 2: Perform two operations on the images processed in Step 1, namely random flipping and random rotation, to enhance the training set samples and avoid overfitting, and then form the final required training set samples;

[0040] Figure 1 This is the overall framework diagram of the medical image segmentation method based on residual axial attention of the present invention. In this embodiment, it is carried out according to the following steps:

[0041] Step 3: Train the dataset obtained in Step 2 through a fully supervised training method. The purpose of this step is to extract the shallow semantic features of medical images. The specific implementation is as follows:

[0042] Step 3.1: The neural network uses UNet as the main framework, which contains 4 convolutional layers, and a max-pooling layer and a ReLU activation function are inserted into each convolutional layer. The dimension size of the convolutional kernel of each convolutional layer is 3×3, and the shallow semantic features E1, E2, E3, and E4 of the input image I are initially obtained;

[0043] Step 3.2: Reshape the shallow semantic features E1, E2, E3, and E4 of different scales into token sequences T1, T2, T3, and T4 with sizes of P, P / 2, P / 4, and P / 8 respectively, and all the channel sizes are changed to 128;

[0044] Step 4: The dilated convolutional module DCM is as Figure 2 shown. Use the shallow semantic feature E4 obtained in Step 3 as the input of DCM to obtain the output feature E5 containing more information. Among them, DCM contains 5 3×3 convolutional kernels, and the expansion parameters are 1, 2, 4, 8, and 16, and the receptive field sizes are N = 3×3, 7×7, 15×15, 31×31, and 63×63 respectively;

[0045] Step 5: Use the token sequences T1, T2, T3, and T4 of the shallow semantic features obtained in Step 3 as the input, and adopt a channel attention module based on residual axial attention to learn the correlation between different channels, and obtain the deep features T c1 、T c2 、T c3 、T c4 of enhanced image channels. The specific implementation is as follows:

[0046] Step 5.1, the residual axial attention-based channel attention module RA-CAM is as follows Figure 4 As shown, first, the token sequences T1, T2, T3, and T4 of the shallow semantic features obtained in Step 3 are concatenated along the channel axis to obtain T Σ ; Then, in order to enable the network to adaptively adjust the weights between feature channels, T Σ is used as the query, value, and key of the residual axial attention module, and the channel attention module is adopted to learn the relationships between channels and capture global dependencies T C ; Finally, T C is segmented into T c1 , T c2 , T c3 , T c4 ;

[0047] Step 5.2, aiming at the problem that Transformer lacks explicit perception of position information, the residual idea is introduced into the axial attention to form the residual axial attention RA-Attention;

[0048] Step 5.3, RA-Attention is as follows Figure 3 As shown, in the RA-Attention operation, first, the position information r V , r Q , r K is added to the axial attention weights W V , W Q , W K to retain the spatial structure and position relationship of the lesion; then T Σ will be passed as the residual term Res to the next attention layer; finally, RA-Attention is performed on the height axis and width axis of the feature map in turn;

[0049] Step 6, the residual axial attention-based spatial attention module RA-SAM is as follows Figure 4 As shown, taking the deep features T c1 , T c2 , T c3 , T c4 with enhanced channels obtained in Step 5 as the input, the residual axial attention-based spatial attention module is adopted to learn the importance of different spatial positions, and the deep features T s1 , T s2 , T s3 , T s4 with enhanced image space are obtained. The specific implementation is as follows:

[0050] Step 6.1, first, the deep features T c1 , T c2 , T c3, T c4 Stitch along the patch-axis to obtain T Σc ;

[0051] Step 6.2, to enable the network to dynamically adjust weights at different spatial positions, use T Σc as the value and key of the residual axial attention module, and T1, T2, T3, T4 as queries. Use the spatial attention module to learn the spatial correlation between multiple scale token sequences and perform cross-attention operations to obtain T s1 , T s2 , T s3 , T s4 ;

[0052] Step 7, perform LayerNorm layer normalization and MLP multi-layer perceptron on T s1 , T s2 , T s3 , T s4 in sequence to obtain enhanced deep features T M1 , T M2 , T M3 , T M4 , and then use T s1 , T s2 , T s3 , T s4 as residual terms to improve the expression ability of the network;

[0053] Step 8, use the deep features T M1 , T M2 , T M3 , T M4 obtained in Step 7 as inputs, and perform semantic feature alignment operations using the semantic alignment module SRA to obtain aligned features O1, O2, O3, O4 with consistent semantic information. The specific implementation is as follows:

[0054] Step 8.1, the semantic alignment module SRA is as Figure 5 shown. SRA reshapes the upsampled feature D i into a token sequence D Ti (i = 1, 2, 3, 4), where D4 is E5 obtained in Step 4; perform residual axial attention operations on D Ti and the deep features T M1 , T M2 , T M3 , T M4 , and use the upsampled feature D i to guide the deep features T M1 , T M2 , T M3 , T M4 to eliminate the semantic inconsistency between the two;

[0055] Step 8.2, concatenate the output after residual axial attention with the token sequence D of the upsampled features i in the channel dimension, then perform convolution with a 3×3 kernel, BatchNorm normalization, and ReLU non-linear activation function in sequence, and finally obtain aligned features O1, O2, O3, and O4 with consistent semantic information;

[0056] Step 9, in the loss function, first simulate the dynamic model to construct the loss function Loss ACM , Loss ACM consists of two parts: the length term Length and the region term Region. The length term Length constrains the segmentation boundary contour to make the contour line relatively smooth, and the region term Region maintains the shape based on the variance of the regions inside and outside the boundary:

[0057]

[0058] where v is the true value, u is the predicted value, C represents the contour curve, Ω represents the set of all points included in the entire image, c1 is the value of all pixels inside the GT image boundary, c2 is the value of all pixels outside the GT image boundary, and λ is a weight coefficient greater than zero; then for the problem that the lesion part in medical images is usually much smaller than the background part, the total loss function Loss of the segmentation network is:

[0059] Loss = Loss ACM + α·Loss Dice + β·Loss BCE

[0060] where, Loss Dice is the Dice loss function, Loss BCE is the cross-entropy loss, and α and β are weight coefficients greater than zero;

[0061] Step 10, input the training set samples in Step 2 into the network from Step 3 to Step 9, and set the network hyperparameters: the learning rate is 0.001, the number of iterations adopts an early stopping strategy and is not fixed, the batch size is 4, the optimizer is Adam, five-fold cross-validation is adopted, the loss function is Loss, and train the network to obtain the final pre-trained model for medical image segmentation;

[0062] Step 11, input the public test set into the pre-trained model obtained in Step 10, and the network can segment the lesion part of the medical image. Figure 6 are some medical images on the public test set, Figure 7 are the results of segmenting these images.

[0063] Starting from the problems of blurred boundaries and complex changes in the lesion parts in medical images, the present invention provides a medical image segmentation method based on residual axial attention. First, in terms of feature enhancement, the method adopts a channel-spatial transformer based on residual axial attention. This module fuses multi-scale features at multiple stages in the encoder, which can effectively enhance the captured deep features. In addition, to solve the problem of the semantic gap between the deep features obtained by the transformer and the upsampled features, the method designs a semantic alignment module based on residual axial attention. It uses the shallow features extracted by the decoder to guide the semantic information alignment of the deep features obtained by the transformer. Dilated convolution is used in the network to expand the receptive field, so that the shallow features contain more image information. Finally, combined with the loss function of edge prediction, the regional information is used as the geometric constraint for the image segmentation task to encourage the model to generate clear lesion boundaries. The algorithm of the present invention has excellent segmentation performance and is applicable to medical images with blurred lesion boundaries.

[0064] Although the above description has been made of the illustrative specific embodiments of the present invention, it should be clear that the present invention is not limited to the scope of the specific embodiments. For those of ordinary skill in the art of this technology, as long as various changes are within the spirit and scope of the present invention defined and determined by the appended claims, these changes are obvious, and all inventions and creations using the concept of the present invention are within the scope of protection.

Claims

1. A medical image segmentation method based on residual axial attention, characterized in that, The design enhances the deep features of images based on a residual axial attention-based channel-spatial converter, eliminates feature differences based on a residual axial attention-based semantic alignment module, and performs geometric constraints based on an edge prediction loss function, including five parts: dataset preprocessing, shallow feature extraction, deep feature extraction, semantic feature alignment, and network model training and testing. The first part includes two steps: Step 1, download the public medical datasets GlaS and MoNuSeg. The original images in dataset GlaS have different sizes, mostly RGB images of size 775×522, and the original images in dataset MoNuSeg are all of size 1000×1000. Then, uniformly set the sample resolution of datasets GlaS and MoNuSeg to 224×224. Step 2, perform two operations on the images processed in Step 1, namely random flipping and random rotation, to enhance the training set samples, and then form the final required training set samples. The second part includes two steps: Step 3, train the dataset obtained in Step 2 through a fully supervised training method. The purpose of this step is to extract the shallow semantic features of medical images. The specific implementation is as follows: (1) The neural network uses UNet as the main framework, which contains 4 convolutional layers, and a max-pooling layer and a ReLU activation function are inserted into each convolutional layer. The dimension size of the convolutional kernel in each convolutional layer is 3×3, and the shallow semantic features E1, E2, E3, and E4 of the input image I are initially obtained. (2) Reshape the features E1, E2, E3, and E4 of different scales into token sequences T1, T2, T3, and T4 with sizes of P, P / 2, P / 4, and P / 8 respectively, where the channel dimensions are all changed to 128. Step 4, use the shallow semantic feature E4 obtained in Step 3 as the input of the dilated convolutional module DCM to obtain an output feature E5 containing more information. DCM contains 5 convolutional kernels of 3×3, and the expansion parameters are 1, 2, 4, 8, and 16, and the receptive field sizes are N = 3×3, 7×7, 15×15, 31×31, and 63×63 respectively. The third part includes three steps: Step 5: Using the token sequences T1, T2, T3, and T4 of the shallow semantic features obtained in Step 3 as inputs, and adopting the residual axial attention-based channel attention module RA-CAM to learn the correlations between different channels, obtaining the deep features T of enhanced image channels c1 , T c2 , T c3 , T c4 , and the specific implementation is as follows: (1) First, connect the token sequences T1, T2, T3, and T4 of the shallow semantic features obtained in step 3 along the channel axis to obtain T Σ ; then, to enable the network to adaptively adjust the weights between feature channels, use T Σ as the query, value, and key of the residual axial attention module. Among them, the residual axial attention module refers to the problem that the Transformer lacks explicit perception of position information. Introduce the residual idea into the axial attention to form the residual axial attention RA-Attention; finally, use the channel attention module based on residual axial attention to learn the relationship between channels and capture the global dependence T C , and split T C into T c1 , T c2 , T c3 , T c4 ; (2) In the RA-Attention operation, first, the position information r V , r Q , r K are added to the axial attention weights W V , W Q , W K to retain the spatial structure and positional relationship of the lesion; then, T Σ is passed as the residual term Res to the next attention layer; finally, RA-Attention is performed on the height axis and width axis of the feature map in sequence; Step 6: Take the T of the channel-enhanced deep features obtained in Step 5 c1 、T c2 、T c3 、T c4 as inputs, and use the residual axial attention-based spatial attention module RA-SAM to learn the importance of different spatial positions, obtaining the deep features T s1 、T s2 、T s3 、T s4 of image spatial enhancement. The specific implementation is as follows: (1) First, splice the channel-enhanced deep features T c1 , T c2 , T c3 , T c4 along the patch-axis to obtain T Σc ; (2) To enable the network to dynamically adjust weights at different spatial positions, T Σc is used as the value and key of the residual axial attention module, T1, T2, T3, and T4 are queries, and the spatial attention module based on residual axial attention is adopted to learn the spatial correlation between multiple-scale token sequences, and cross-attention operation is performed to obtain T s1 , T s2 , T s3 , T s4 ; Step 7, perform LayerNorm layer normalization and MLP multi-layer perceptron on T s1 , T s2 , T s3 , T s4 in sequence to obtain enhanced deep features T M1 , T M2 , T M3 , T M4 , and then use T s1 , T s2 , T s3 , T s4 as residual terms to improve the expression ability of the network; The fourth part includes one step: Step 8: Take the deep features T M1 , T M2 , T M3 , T M4 as inputs, and perform semantic feature alignment operations using the semantic alignment module SRA to obtain aligned features O1, O2, O3, and O4 with consistent semantic information. The specific implementation is as follows: (1) The SRA reshapes the E5 obtained in step 4 into a token sequence D Ti , where i = 1, 2, 3, 4; and D Ti is subjected to a residual axial attention operation with the deep features T M1 , T M2 , T M3 , T M4 . The upsampling feature D Ti is used to guide the deep features T M1 , T M2 , T M3 , T M4 to eliminate the semantic inconsistency between the two; (2) Concatenate the output after residual axial attention with the token sequence D i along the channel dimension, then perform convolution with a 3×3 convolutional kernel, BatchNorm batch normalization, and ReLU non-linear activation function in sequence, and finally obtain the aligned features O1, O2, O3, O4 with consistent semantic information; The fifth part includes three steps: Step 9. In the loss function, first simulate the kinetic model to construct the loss function Loss ACM , Loss ACM includes two parts: the length term Length and the region term Region. The length term Length constrains the segmentation boundary contour to make the contour line relatively smooth, and the region term Region maintains the shape based on the variance of the regions inside and outside the boundary: where v is the true value, u is the predicted value, C represents the contour curve, Ω represents the set of all points contained in the entire image, c1 is the value of all pixels within the GT image boundary, c2 is the value of all pixels outside the GT image boundary, and λ is a weight coefficient greater than zero. Then, for the problem that the lesion part in medical images is usually much smaller than the background part, the total loss function Loss of the segmentation network is: Loss=Loss ACM +α·Loss Dice +β·Loss BCE Among them, Loss Dice is the Dice loss function, and Loss BCE is the cross-entropy loss, where α and β are weight coefficients greater than zero; Step 10, input the training set samples in Step 2 into the network from Step 3 to Step 9, set the network hyperparameters: the learning rate is 0.001, the number of iterations adopts an early stopping strategy and is not fixed, the batch size is 4, the optimizer is Adam, and five-fold cross-validation is adopted. The loss function is Loss, and train the network to obtain the final pre-trained model for medical image segmentation. Step 11: Input the public test set into the pre-trained model obtained in Step 10 to segment the lesion part of the medical image.

2. The medical image segmentation method based on residual axial attention according to claim 1, wherein Step 5 uses a RA-CAM module based on the residual structure to obtain deep features with enhanced image channels.

3. A medical image segmentation method based on residual axial attention according to claim 1, characterized in that, Step 6 uses a RA-SAM module based on the residual structure to obtain deep features with enhanced image space.

Images