Medical image segmentation method based on residual axial attention

Abstract

The invention provides a medical image segmentation method based on residual axial attention. Firstly, in order to solve the problems of fuzzy boundary and complex change of a focus part in a medical image, a channel-space converter based on residual axial attention is designed, the module fuses multi-scale features of multiple stages in an encoder, and captured deep features can be effectively enhanced; then, aiming at the problem that semantic difference exists between the deep features obtained by the converter and the up-sampling features, a semantic alignment module based on residual axial attention is designed, and shallow features extracted by a decoder are utilized to guide the deep features obtained by the converter to perform semantic information alignment; and finally, designing a loss function combined with edge prediction, and taking the regional information as a geometric constraint of an image segmentation task to encourage the model to generate a clear lesion boundary. According to the method, the converter based on residual axial attention and the semantic alignment module are utilized, and the edge prediction loss function is combined, so that fine segmentation of the medical image can be realized.

Description

Technical Field

[0001] The invention relates to image processing technology, and in particular 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 structures or lesion areas of interest from medical images. The rapid development of deep learning technology has brought revolutionary changes 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 feature representations in images, but also gradually extract abstract features in images through multi-layer convolution and pooling operations. At the same time, based on the label information in the training data, feature representations of different structures and tissues in the image are learned, thereby achieving accurate segmentation of medical images.

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

[0004] In this context, the attention mechanism can help the model better focus on important areas in the image, thereby improving the accuracy and efficiency of segmentation and providing new possibilities for medical image analysis. The visual transformer ViT based on the self-attention mechanism can effectively capture the long-distance dependencies between pixels. However, ViT requires a large amount of training data to achieve good performance and has high computational complexity, especially for high-resolution input images. Swin Transformer reduces computational complexity by introducing a hierarchical mechanism, which results in a small receptive field due to limited information interaction between windows. Most Transformer-based methods require pre-training on large datasets to achieve satisfactory results, perform poorly on small samples or unbalanced datasets, and have 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, in response to the above problems in medical image segmentation, some methods try to combine CNN and Transformer to enhance network capabilities. However, the introduction of attention mechanism will increase the computational complexity of the model, and still face the problem of balancing performance and computational cost. Therefore, designing a medical image segmentation model that is both efficient and accurate is of great research significance. Summary of the invention

[0006] The purpose of the present invention is to solve the problem of semantic gap between deep features and shallow features in medical image segmentation, and to provide a medical image segmentation method based on residual axial attention, which effectively extracts and fuses medical image features through a deep learning network to achieve accurate segmentation of the lesion part of the medical image.

[0007] In order to achieve the above-mentioned 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 data set; the second part is to extract shallow features of the medical image; the third part is to extract deep features of the medical image after the shallow features are extracted; the fourth part is to semantically align the shallow features and deep features extracted from the medical image; the fifth part is to train and test the medical image segmentation network model based on residual axial attention, and finally obtain the segmentation result map, which is specifically:

[0008] The first part consists of two steps:

[0009] Step 1: Download the public medical datasets GlaS and MoNuSeg. The original images of the GlaS dataset vary in size, most of which are 775

[0010] ×522 RGB images, the original image size of the MoNuSeg dataset is 1000×1000; then the sample resolution of the GlaS and MoNuSeg datasets is uniformly set to 224×224;

[0011] Step 2: Perform random flipping and random rotation on the image processed in step 1 to enhance the training set samples and avoid excessive

[0012] Fitting, and then forming the final required training set samples;

[0013] The second part consists of two steps:

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

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

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

[0017] The third part consists of three steps:

[0018] Step 5: Take the shallow semantic feature tag sequence T1, T2, T3, and T4 obtained in step 3 as input, and use the channel attention module RA-CAM based on residual axial attention to learn the correlation between different channels to obtain the deep feature T enhanced by the image channel. c1 、T c2 、T c3 、T c4 ;

[0019] Step 6: T of the channel-enhanced deep features obtained in step 5 c1 、T c2 、T c3 、T c4 As input, the spatial attention module RA-SAM based on residual axial attention is used to learn the importance of different spatial positions and obtain the deep features T of image spatial enhancement. s1 、T s2 、T s3 、T s4 ;

[0020] Step 7, T s1 、T s2 、T s3 、T s4 LayerNorm layer normalization and MLP multi-layer perceptron are performed in sequence to obtain enhanced deep semantic features.

[0021] T M1 、T M2 、T M3 、T M4 , then Ts1 、T s2 、T s3 、T s4 As a residual term, it improves the network's expressiveness;

[0022] The fourth part consists of a step:

[0023] Step 8: The deep semantic features T obtained in step 7 M1 、T M2 、T M3 、T M4 As input, the semantic alignment module SRA is used for semantic

[0024] Feature alignment operation to obtain aligned features O1, O2, O3, and O4;

[0025] Part Five consists of three steps:

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

[0027] 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 the early stopping strategy so that it is not fixed, the batch size is 4, the optimizer is Adam, five-fold cross validation is adopted, the loss function is Loss, and the network is trained to obtain the final medical image segmentation pre-training model;

[0028] 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.

[0029] The present invention provides a medical image segmentation method based on residual axial attention. First, in order to solve the problem of blurred boundaries and complex changes in lesion parts in medical images, a channel-space converter based on residual axial attention is designed. This module fuses multi-scale features of multiple stages in the encoder, which can effectively enhance the captured deep features; then, in view of the semantic gap between the deep features obtained by the converter and the upsampled features, a semantic alignment module based on residual axial attention is designed, and the shallow features extracted by the decoder are used to guide the deep features obtained by the converter to align semantic information; finally, a loss function combined with edge prediction is designed, and regional information is used as a geometric constraint for the image segmentation task to encourage the model to generate clear lesion boundaries. The present invention uses a converter based on residual axial attention and a semantic alignment module, combined with an edge prediction loss function, to achieve fine segmentation of medical images. BRIEF DESCRIPTION OF THE DRAWINGS

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

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

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

[0033] Figure 4 The residual axial attention based channel-space transformer RA-CST of the present invention;

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

[0035] Figure 6 Some medical images in the public test set;

[0036] Figure 7 To utilize the present invention to process Figure 6 The result after. DETAILED DESCRIPTION

[0037] In order to better understand the present invention, the medical image segmentation method based on residual axial attention of the present invention is described in more detail below in conjunction with specific embodiments. In the following description, the detailed description of the current prior art may dilute the subject content of the present invention, and these descriptions will be ignored here.

[0038] Step 1, download the public medical datasets GlaS and MoNuSeg. The original images of the dataset GlaS have different sizes, most of which are RGB images of size 775×522, and the original images of the dataset MoNuSeg have a size of 1000×1000. Then, the sample resolution of the datasets GlaS and MoNuSeg is uniformly set to 224×224.

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

[0040] Figure 1 : is the overall network framework diagram of the medical image segmentation method based on residual axial attention of the present invention. In this embodiment, the following steps are performed:

[0041] Step 3: Train the data set obtained in step 2 through full supervision training. The purpose of this step is to extract 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 each convolutional layer is inserted with a maximum pooling layer and a ReLU activation function. The convolution kernel dimension 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 tag sequences T1, T2, T3, and T4 of sizes P, P / 2, P / 4, and P / 8, respectively, where the channel sizes are all changed to 128;

[0044] Step 4, dilated convolution module DCM as Figure 2 As shown in the figure, the shallow semantic feature E4 obtained in step 3 is used as the input of DCM to obtain the output feature E5 containing more information, where DCM contains 5 3×3 convolution kernels, 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: Take the shallow semantic feature tag sequence T1, T2, T3, and T4 obtained in step 3 as input, and use the channel attention module based on residual axial attention to learn the correlation between different channels to obtain the deep feature T of image channel enhancement. c1 、T c2 、T c3 、T c4 , the specific implementation is as follows:

[0046] Step 5.1, the channel attention module RA-CAM based on residual axial attention is as follows Figure 4 As shown in the figure, firstly, the tag sequences T1, T2, T3, and T4 of the shallow semantic features obtained in step 3 are connected through the channel axis to obtain T Σ ; Then, in order to make the network adaptively adjust the weights between feature channels, T Σ As the query, value and key based on the residual axial attention module, a channel attention module is adopted to learn the relationship between channels and capture the global dependency T C ; Finally, T C Split into T c1 、T c2 、T c3 、T c4 ;

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

[0048] Step 5.3, RA-Attention Figure 3 As shown, in the RA-Attention operation, the position information r is first V 、r Q 、r K Add axial attention weight W V , W Q , W K , retaining the spatial structure and positional relationship of the lesion; then T Σ The residual term Res is passed to the next attention layer; finally, RA-Attention is performed on the height axis and width axis of the feature map respectively;

[0049] Step 6: The spatial attention module RA-SAM based on residual axial attention is as follows Figure 4 As shown, the channel-enhanced deep feature T obtained in step 5 is c1 、T c2 、T c3 、T c4 As input, a spatial attention module based on residual axial attention is used to learn the importance of different spatial positions and obtain the deep features T of image spatial enhancement. s1 、T s2 、T s3 、T s4 , the specific implementation is as follows:

[0050] Step 6.1: First, the channel-enhanced deep feature T obtained in step 5 is c1 、T c2 、T c3、T c4 Splice along the patch-axis to get T Σc ;

[0051] Step 6.2: In order to make the network dynamically adjust weights at different spatial locations, T Σc As the value and key based on the residual axial attention module, T1, T2, T3, and T4 are queries. The spatial attention module is used to learn the spatial correlation between multiple scale tag sequences, and the cross attention operation is performed to obtain T s1 、T s2 、T s3 、T s4 ;

[0052] Step 7, T s1 、T s2 、T s3 、T s4 LayerNorm layer normalization and MLP multi-layer perceptron are performed in sequence to obtain the enhanced deep feature T M1 、T M2 、T M3 、T M4 , then T s1 、T s2 、T s3 、T s4 As a residual term, it improves the network's expressiveness;

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

[0054] Step 8.1, the semantic alignment module SRA is as follows Figure 5 As shown, SRA will upsample the feature D i Reshape into a labeled sequence D Ti (i=1,2,3,4), where D4 is E5 obtained in step 4; Ti And the deep feature T obtained in step 7 M1 、T M2 、T M3 、T M4 Perform residual axial attention operation and use the upsampled feature D i To guide the deep feature T M1 、T M2 、T M3 、T M4 , eliminating the semantic inconsistency between the two;

[0055] Step 8.2: Compare the output after residual axis attention with the labeled sequence D of upsampled features i Splicing is performed in the channel dimension, and then 3×3 convolution kernel, BatchNorm normalization and ReLU nonlinear activation function are performed in sequence, and finally the alignment features O1, O2, O3, and O4 with consistent semantic information are obtained;

[0056] Step 9: In the loss function, first simulate the dynamics model and construct the loss function Loss ACM , Loss ACM It 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 relatively smooth, and the region term Region maintains the shape based on the variance of the area 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 contained in the entire image, c1 is the value of all pixels within the boundary of the GT image, c2 is the value of all pixels outside the boundary of the GT image, and λ is a weight coefficient greater than zero; then for the problem that the lesion part is usually much smaller than the background part in medical images, the total loss function Loss of the segmentation network is:

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

[0060] Among them, Loss Dice is the Dice loss function, Loss BCE is the cross entropy loss, α 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, set the network hyperparameters: the learning rate is 0.001, the number of iterations adopts the early stopping strategy so that it is not fixed, the batch size is 4, the optimizer is Adam, five-fold cross validation is adopted, the loss function is Loss, and the network is trained to obtain the final medical image segmentation pre-training model;

[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 These are some medical images from the public test set. Figure 7 These images are segmented.

[0063] The present invention starts from the problem that the boundary of the lesion part in medical images is blurred and the changes are complex, and a medical image segmentation method based on residual axial attention is given. First, the method adopts a channel-space transformer based on residual axial attention for feature enhancement. This module fuses the multi-scale features of multiple stages in the encoder, which can effectively enhance the captured deep features; in addition, in order to solve the problem of semantic gap between the deep features obtained by the converter 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 deep features obtained by the converter to align the semantic information. The network uses dilated convolution 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 of 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 suitable for medical images with blurred lesion boundaries.

[0064] Although the above describes an illustrative specific embodiment of the present invention, it should be clear that the present invention is not limited to the scope of the specific embodiment. For ordinary technicians in this technical field, as long as various changes are within the spirit and scope of the present invention as defined and determined by the attached claims, these changes are obvious, and all inventions and creations using the concept of the present invention are protected.

Claims

1. A medical image segmentation method based on residual axial attention, characterized in that: Design a channel-space transformer based on residual axial attention to enhance the deep features of the image, a semantic alignment module based on residual axial attention to eliminate feature differences and an edge prediction loss function to perform geometric constraints, including five parts: dataset preprocessing, shallow feature extraction, deep feature extraction, semantic feature alignment, network model training and testing: The first part consists of two steps: Step 1, download the public medical datasets GlaS and MoNuSeg. The original images of the dataset GlaS have different sizes, most of which are RGB images of size 775×522, and the original images of the dataset MoNuSeg have a size of 1000×1000. Then, the sample resolution of the datasets GlaS and MoNuSeg is uniformly set to 224×224. Step 2, randomly flipping and randomly rotating the image processed in step 1 to enhance the training set samples, and then form the final required training set samples; The second part consists of two steps: Step 3: Train the data set obtained in step 2 through full supervision training. The purpose of this step is to extract 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 each convolutional layer is inserted with a maximum pooling layer and a ReLU activation function. The convolution kernel dimension 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; (2) Features E1, E2, E3, and E4 of different scales are reshaped into tag sequences T1, T2, T3, and T4 of sizes P, P / 2, P / 4, and P / 8, respectively, where the channel sizes are all changed to 128; Step 4: Use the shallow semantic feature E4 obtained in step 3 as the input of the dilated convolution module DCM to obtain the output feature E5 containing more information, where DCM contains 5 3×3 convolution kernels, 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 consists of three steps: Step 5: Take the shallow semantic feature tag sequence T1, T2, T3, and T4 obtained in step 3 as input, and use the channel attention module RA-CAM based on residual axial attention to learn the correlation between different channels to obtain the deep feature T enhanced by the image channel. c1 , T c2 , T c3 , T c4 , the specific implementation is as follows: (1) First, the shallow semantic feature tag sequences T1, T2, T3, and T4 obtained in step 3 are connected through the channel axis to obtain T Σ ; Then, in order to make the network adaptively adjust the weights between feature channels, T Σ As the query, value and key based on the residual axial attention module, a channel attention module is adopted to learn the relationship between channels and capture the global dependency T C ; Finally, T C Split into T c1 , T c2 , T c3 , T c4 ; (2) To address the problem that Transformer lacks explicit perception of position information, the residual idea is introduced into axial attention to form residual axial attention RA-Attention; (3) In the RA-Attention operation, the position information r is first V 、r Q 、r K Add axial attention weight W V , W Q , W K , retaining the spatial structure and positional relationship of the lesion; then T Σ The residual term Res is passed to the next attention layer; finally, RA-Attention is performed on the height axis and width axis of the feature map respectively; Step 6: T of the channel-enhanced deep features obtained in step 5 c1 , T c2 , T c3 , T c4 As input, the spatial attention module RA-SAM based on residual axial attention is used to learn the importance of different spatial positions and obtain the deep features T of image spatial enhancement. s1 , T s2 , T s3 , T s4 , the specific implementation is as follows: (1) First, the channel-enhanced deep feature T obtained in step 5 c1 , T c2 , T c3 , T c4 Splice along the patch-axis to get T Σc ; (2) In order to enable the network to dynamically adjust weights at different spatial locations, T Σc As the value and key based on the residual axial attention module, T1, T2, T3, and T4 are queries. The spatial attention module is used to learn the spatial correlation between multiple scale tag sequences, and the cross attention operation is performed to obtain T s1 , T s2 , T s3 , T s4 ; Step 7, T s1 , T s2 , T s3 , T s4 LayerNorm layer normalization and MLP multi-layer perceptron are performed in sequence to obtain the enhanced deep feature T M1 , T M2 , T M3 , T M4 , then T s1 , T s2 , T s3 , T s4 As a residual term, it improves the network's expressiveness; The fourth part consists of a step: Step 8: The deep feature T obtained in step 7 M1 , T M2 , T M3 , T M4 As input, the semantic alignment module SRA is used to perform semantic feature alignment operations to obtain alignment features O1, O2, O3, and O4 with consistent semantic information. The specific implementation is as follows: (1) SRA upsamples the feature D i Reshape into a labeled sequence D Ti (i=1,2,3,4), where D4 is E5 obtained in step 4; Ti And the deep feature T obtained in step 7 M1 , T M2 , T M3 , T M4 Perform residual axial attention operation and use the upsampled feature D Ti To guide the deep feature T M1 , T M2 , T M3 , T M4 , eliminating the semantic inconsistency between the two; (2) Compare the output after residual axis attention with the label sequence D i Splicing is performed in the channel dimension, and then 3×3 convolution kernel, BatchNorm batch normalization and ReLU nonlinear activation function are performed in sequence, and finally the alignment features O1, O2, O3, O4 with consistent semantic information are obtained; Part 5 consists of three steps: Step 9: In the loss function, first simulate the dynamics model and construct the loss function Loss ACM , Loss ACM It 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 relatively smooth, and the region term Region maintains the shape based on the variance of the area 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 boundary of the GT image, c2 is the value of all pixels outside the boundary of the GT image, and λ is a weight coefficient greater than zero; then for the problem that the lesion part is usually much smaller than the background part in medical images, 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, Loss BCE is the cross entropy loss, α 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 the early stopping strategy so that it 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 the network is trained to obtain the final medical image segmentation pre-training model; 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.

2. The medical image segmentation method based on residual axial attention according to claim 1, characterized in that: Step 5 uses a RA-CAM module based on a residual structure to obtain deep features for image channel enhancement.

3. The 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 a residual structure to obtain deep features for image spatial enhancement.

4. The medical image segmentation method based on residual axial attention according to claim 1, characterized in that: Step 4 uses the DCM module to extract more effective features.

5. The method for medical image segmentation based on residual axial attention according to claim 1, characterized in that: Step 10 uses a new loss function Loss ACM , which optimizes the performance of the model in the segmentation task while taking image features into account.

Images