Transformer-based enhanced medical image segmentation network
Through the dual branch encoder and enhanced feature fusion module, combined with convolutional neural network and transformer, the limitations of medical image segmentation model in global relationship modeling and feature extraction are solved, and the segmentation accuracy and robustness are improved, especially In complex image processing.
Patent Information
- Application Number
- CN202411722331.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-27
- Publication Date
- 2025-01-14
AI Technical Summary
Existing medical image segmentation models are not good in processing complex textures and blurred boundaries, and convolutional neural networks and transformer modules have limitations in global relationship modeling and feature extraction, and fail to fully consider the contextual correlation of input data. Fusion with feature information.
Using a dual branch encoder structure, combining convolutional neural network and transformer, the local and global feature extraction capabilities are improved through the area perception attention module and the enhanced feature fusion module, and the loss function module is used to optimize feature consistency and cross-dimensional modeling. .
It improves the accuracy and robustness of medical image segmentation, especially when processing complex images and blur boundaries, enhances the feature representation ability and the fusion of multi-scale features, and improves the performance of fine-grained segmentation tasks.
Smart Images
Figure CN119313909A_ABST
Abstract
Description
Technical Field
[0001] The invention belongs to the technical field of medical image processing and relates to a transformer-based enhanced medical image segmentation network. Background Art
[0002] Medical image segmentation is to divide medical images into several non-overlapping connected regions based on certain similar features of medical images. The results of medical image segmentation can not only realize the recognition and location of abnormal areas, but also provide a strong reference for doctors to diagnose diseases and formulate surgical plans.
[0003] At present, convolutional neural networks (CNNs) are mainly used as the main tool for analyzing medical images, and gradually developed into U-Net models that introduce attention mechanisms, U-Net++ models with dense skip connections and nested U-Nets, R2U-Net architectures that introduce cyclic residual connections, DoubleU-Net models with two cascaded U-Net structures, and EGE-Unet models that perform attention mechanisms in different axes. However, the local receptive fields of these models limit the ability to model long-distance dependencies, especially when dealing with targets with complex textures and fuzzy boundaries, the segmentation effect is still insufficient.
[0004] In order to solve the defect of convolutional neural networks in extracting global features from images. Vaswani et al. proposed the first transformer network to solve the problem of long-range dependency and position information representation in sequence modeling. Due to the excellent performance of transformer networks in the field of natural language processing, Dosovitskiy proposed a visual transformer architecture that can better capture the global information of images. Subsequently, Chen et al. proposed TransUNet, which uses transformer encoders as feature extractors of U-Net. The model uses a self-attention mechanism to improve the segmentation effect. However, the convolution and self-attention mechanisms in this method do not work in the entire network architecture, which poses a challenge to continuously establish local and global representation models. In order to solve the problem of limited local area interaction information and affecting edge segmentation accuracy in the global modeling process, Zhang et al. proposed the TransFuse architecture, which effectively utilizes the respective advantages of convolutional neural networks and transformers. When performing multi-organ segmentation tasks with complex backgrounds, TransFuse cannot effectively extract and process global features in complex image data.
[0005] That is, there are the following problems in the prior art: first, the existing transformer module fails to fully consider the contextual relevance of the input data when establishing the global relationship model, resulting in unsatisfactory global semantic feature extraction. Second, the feature information fusion between the convolutional neural network branch and the transformer branch ignores the interaction of the channel and spatial dimensions, resulting in the loss of some feature information. Summary of the invention
[0006] In view of the above problems, the present invention proposes a transformer-based enhanced medical image segmentation network, which well solves the problems in the prior art.
[0007] In order to achieve the above object, the technical solution adopted by the present invention is as follows:
[0008] A transformer-based enhanced medical image segmentation network, including an encoder and a decoder,
[0009] The encoder includes a dual-branch encoder consisting of a convolutional neural network encoder branch and a transformer encoder branch, each of the encoder branches includes three processing stages, each of the transformer encoder branch processing stages includes a region-aware attention module block and a block merging layer, and an enhanced feature fusion FFR module is provided between the convolutional neural network encoder branch and the transformer encoder branch:
[0010] The encoder further comprises a ResNet-50 encoder branch, wherein the ResNet-50 encoder branch and the transformer encoder branch are used together for local context feature extraction;
[0011] The decoder includes a convolutional neural network module and an upsampling module, the convolutional neural network module includes a convolution layer, a batch normalization layer and an activation function ReLU layer, and the upsampling module includes a transposed convolution layer;
[0012] Among them, the formulas of the convolutional neural network encoder branch and the transformer encoder branch are:
[0013]
[0014]
[0015]
[0016]
[0017]
[0018]
[0019] in, represents the input image, Patch(·) is the patch embedding layer, ST j (·) represents the output feature of the jth layer of the transformer encoder branch, SF j (·) represents the output characteristics of the jth layer of the convolutional neural network encoder branch, j∈{1,2,3}.
[0020] Optionally, the transformer encoder branch includes a region-aware attention calculation RAA unit, which is used to generate a context score pv. The region-aware attention calculation RAA unit includes three vector branches Q, K, and V. The vector branch Q obtains a probability distribution matrix ps through a Softmax layer. The context score pv is calculated as follows:
[0021]
[0022] y=pv*ReLU(I V T V )T O
[0023] Among them, I K Indicates branch K, I V represents branch V, Relu is the activation function, T K , T V , T O Represents the weight value corresponding to the linear layer, and ps is the probability distribution matrix.
[0024] Optionally, an enhanced fully connected EFC module is further included, and the overall process of the enhanced fully connected EFC module is:
[0025] y 1 =FC(RELU(LN(DWConv(FC(X in ))+FC(x in ))))
[0026] Among them, x in is the output of the region-aware attention calculation RAA unit, DWConv is a DW convolutional layer with a convolution kernel size of 3×3, FC is a fully connected layer, and LN is layer normalization.
[0027] Optionally, the enhanced feature fusion FFR module formula is:
[0028] M c (F1)=Sigmoid(MLP(AvgPool(F1))+MLP(MaxPool(F1)))
[0029] M s (F2)=Sigmoid(Conv(AvgPool(F2)+MaxPool(F2)))
[0030] F′=F1*M c (F1)+F2*M s (F2)
[0031] Among them, F1 is the feature output of the convolutional neural network encoder branch, F2 is the feature output of the transformer encoder branch, is the global average pooling, MaxPool is the global maximum pooling, Conv is the convolution layer, and Sigmoid is the activation function.
[0032] Optionally, a loss function module is also included, and the loss function module includes a cross entropy loss function CrossEntrop and a sample similarity function Dice, and the formula is:
[0033] L=λ 1 L Dice +λ 2 L CE
[0034] Among them, L is the total loss value, and λ1 and λ2 represent different weighting coefficients.
[0035] Compared with the prior art, the present invention has the following beneficial effects:
[0036] 1. By fully combining the advantages of convolutional neural network encoder and transformer encoder, through the synergy of the two, the model can learn rich local and global features at different resolutions, thereby improving segmentation accuracy and feature learning ability;
[0037] 2. By setting up the regional awareness attention calculation RAA unit, the multi-scale global information of the image is efficiently obtained, and the correlation of different tags in the context is strengthened through feature adjustment mapping to enhance the feature representation ability. In addition, combined with the enhanced fully connected EFC module, the consistency of features and cross-dimensional modeling capabilities are optimized, making the model more accurate and robust in processing complex images and fuzzy boundaries;
[0038] 3. By setting up an enhanced feature fusion FFR module, the interaction between spatial and channel features is strengthened, which effectively improves the fusion and utilization of multi-scale features. In particular, when dealing with small organs with large-scale shape changes, it has higher perception capabilities, thereby improving the performance of the model in fine-grained medical image segmentation tasks. BRIEF DESCRIPTION OF THE DRAWINGS
[0039] Figure 1 It is a schematic diagram of the overall process of an embodiment of the present invention;
[0040] Figure 2 2 is a schematic diagram of the structure of a regional awareness attention calculation RAA unit according to an embodiment of the present invention;
[0041] Figure 3 is a schematic diagram of the structure of an enhanced fully connected EFC module according to an embodiment of the present invention;
[0042] Figure 4It is a schematic diagram of the structure of the enhanced feature fusion FFR module of an embodiment of the present invention. DETAILED DESCRIPTION
[0043] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.
[0044] See also Figure 1 , is a transformer-based enhanced medical image segmentation network disclosed in an embodiment of the present invention, including an encoder and a decoder, the encoder including a dual-branch encoder consisting of a convolutional neural network (CNN) encoder branch and a transformer (Transformer) encoder branch, the encoder branches each including three processing stages, each transformer encoder branch processing stage including a regional awareness attention module (RAFormer) block and a patch merge layer (Patch merge), an enhanced feature fusion FFR module is arranged between the convolutional neural network encoder branch and the transformer encoder branch; the encoder also includes a ResNet-50 encoder branch, the ResNet-50 encoder branch and the transformer encoder branch are used together for local context feature extraction; the decoder includes a convolutional neural network module and an upsampling module, the convolutional neural network module includes a convolutional layer, a batch normalization layer and an activation function ReLU layer, and the upsampling module includes a transposed convolutional layer;
[0045] Among them, the formulas of the convolutional neural network encoder branch and the transformer encoder branch are:
[0046]
[0047]
[0048]
[0049]
[0050]
[0051]
[0052] in, represents the input image, Patch(·) is the patch embedding layer, STj(·) represents the output features of the jth layer of the transformer encoder branch, SF j (·) represents the output characteristics of the jth layer of the convolutional neural network encoder branch, j∈{1,2,3}.
[0053] Specifically, the convolutional neural network captures the local features of medical images through convolution operations and hierarchical feature representation, and the transformer network realizes the global feature extraction of medical images through cascaded self-attention mechanism and matrix operations with contextual interaction. Through the setting of dual-branch encoder, the local features are fused with the global features, making full use of the local details and global semantic features of medical images, which is more conducive to medical image segmentation.
[0054] In some feasible embodiments, the dual-branch encoders are connected through an enhanced feature fusion FFR module, which combines the local and global feature information obtained from the dual-branch encoders and passes valuable information from the parallelized encoder to the decoder, thereby obtaining pixel-level segmentation results.
[0055] In some feasible methods, the three processing stages of the encoder branch include the first stage capturing low-level texture and edge information; the second and third stages gradually extract higher-order semantic information. In practical applications, the three-stage structure can extract multi-scale features well while maintaining the high efficiency of the model. Adding more stages will seriously increase the complexity of the calculation, and the improvement of the segmentation effect is limited.
[0056] The transformer encoder branch first divides the image into 4*4 patches to maintain local continuity with the convolutional neural network encoder branch. In the processing stage, the region-aware attention module block utilizes the region-aware attention mechanism to learn long-range dependencies and generates downsampled features through the block merging layer.
[0057] The ResNet-50 encoder branch is used together with the transformer encoder branch for local context feature extraction, enabling the model to learn low-resolution information rich in semantic details. Subsequently, the output features of size H / 16×W / 16, H / 8×W / 8, and H / 4×W / 4 are forwarded to the enhanced feature fusion FFR module to fuse tightly integrated local and global features at different resolutions. The fused features are upsampled three times on the extended path, i.e., the skip connection, to provide rich spatial information, and the feature information lost during the downsampling process is compensated by the skip connection, which is conducive to display feature matching.
[0058] After the features enter the decoder, the feature map can be upsampled to twice the original resolution through the transposed convolution layer. After upsampling to twice the resolution, the spatial structure information of the image can be gradually restored, so that the feature map generated by the decoder at each step is closer to the resolution of the original image, providing a basis for the final high-precision segmentation. Then the output is segmented pixel by pixel and linear projection prediction is performed.
[0059] As a specific implementation of the transformer-based enhanced medical image segmentation network provided by the application, please refer to Figure 2 , the transformer encoder branch includes a region-aware attention calculation RAA unit, which is used to generate a context score pv. The region-aware attention calculation RAA unit includes three vector branches Q, K, and V. The vector branch Q obtains the probability distribution matrix ps through the Softmax layer. The context score pv is calculated as:
[0060]
[0061] y=pv*ReLU(I V T V )T o
[0062] Among them, I K Indicates branch K, I V represents branch V, Relu is the activation function, T K , T V , T o Represents the weight value corresponding to the linear layer, and ps is the probability distribution matrix.
[0063] Among them, the context score represents the weight distribution, reflecting the relevance of the current input to the entire image or other parts of the sequence. The context score tells the model which areas of information should be paid more attention to. Through the context score, the model can dynamically adjust attention, giving greater weight to areas that are more relevant to the current task and ignoring unimportant information. The final output is the weighted result of the value vector V, which combines the context score and the feature information of different regions in the image. The output not only contains local detail information, but also integrates global correlation features, which can more comprehensively characterize the characteristics of the input image and provide high-quality feature expression for subsequent tasks such as segmentation or classification.
[0064] Overall, the attention mechanism in the existing transformer encoder model is not enough to extract the semantic features of medical images, and it is easy to ignore the features of the captured focus area. Through the mechanism provided by the region-aware attention calculation RAA unit, the multi-dimensional information in the channel is efficiently processed, and the contextual information of the medical image is reasonably used to realize long-range dependency modeling, which has a good feature representation effect for medical images with high-density noise. The region-aware attention calculation RAA unit mechanism can effectively process the multi-dimensional information in the channel, and reasonably use the contextual information to realize long-range dependency modeling, which has a good feature representation effect for medical images with high-density noise, low contrast, and complex backgrounds.
[0065] In some feasible methods, the input medical image enters the region-aware attention calculation RAA unit after linear calculation of the block merging layer, and obtains three vector branches Q, K, and V. The similarity between the token and the input is calculated by performing an inner product operation on the linear projection. Then, the Q vector obtains a probability distribution matrix through the Softmax layer, which represents the similarity measure between the token and the input. Then, the K vector is multiplied by the probability distribution matrix to obtain the context score, and the V vector is multiplied by the context score to obtain the final output. The region perception capability of the region-aware attention calculation RAA unit not only establishes a long-distance spatial mapping, but also captures more fine-grained spatial information compared to traditional processing, further learns features of different scales and directions, and improves feature representation capabilities.
[0066] As a specific implementation of the transformer-based enhanced medical image segmentation network provided in the application, it also includes an enhanced fully connected EFC module. The overall process of the enhanced fully connected EFC module is:
[0067] y 1 =FC(RELU(LN(DWConv(FC(X in ))+FC(xi n ))))
[0068] Among them, xi n It is the vector after the output of the RAA unit of the region-aware attention calculation is processed by Layer Norm. DWConv is the DW convolution layer with a convolution kernel size of 3×3, FC is the fully connected layer, and LN is the layer normalization.
[0069] For specific usage scenarios, please refer to Figure 3 An enhanced fully connected EFC module is set for the existing multi-layer perceptron MLP structure, which combines the advantages of deep convolution and fully connected layers. The training process is stabilized by layer normalization after deep convolution, which effectively enhances the feature expression ability. The nonlinear conversion ability of the model and the richness of feature expression are improved through a two-layer fully connected design, which improves feature alignment and generates richer feature maps.
[0070] In some feasible ways, the enhanced fully connected EFC module is first processed by layer normalization LN after the depthwise separable convolution DWConv to better capture local context information. The output of the deep convolution layer is added with the result of the fully connected layer FC processing to enhance connectivity and avoid information loss. By embedding the enhanced fully connected EFC module into the transformer encoder, different feature distributions and consistencies are constructed, giving the model stronger feature representation capabilities.
[0071] As a specific implementation of the transformer-based enhanced medical image segmentation network provided in the application, the formula of the enhanced feature fusion FFR module is:
[0072] M c (F1)=Sigmoid(MLP(AvgPool(F1))+MLP(MaxPool(F1)))
[0073] M s (F2)=Sigmoid(Conv(AvgPool(F2)+MaxPool(F2)))
[0074] F′=F1*M c (F1)+F2*M s (F2)
[0075] Among them, F1 is the feature output of the convolutional neural network encoder branch, F2 is the feature output of the transformer encoder branch, is the global average pooling, MaxPool is the global maximum pooling, Conv is the convolution layer, and Sigmoid is the activation function.
[0076] Please understand that Figure 4 , the enhanced feature fusion FFR module weights it from the two dimensions of channel and space, activates the features for the segmentation results, suppresses irrelevant features through the activation function, and uses the local and global information of the dual-branch encoder to achieve multi-dimensional effective feature fusion. For the channel attention branch, the original output of the convolutional neural network encoder branch is first subjected to global maximum pooling and global average pooling respectively. The output result is sent to the fully connected layer, and finally the channel representation is nonlinearly transformed using the Sigmoid activation function to obtain the weight matrix on the channel; for the spatial attention branch, the original output of the transformer encoder branch is subjected to global maximum pooling and global average pooling, and the channel information representation is obtained using average pooling. The two results are concatenated, Conv is the convolution kernel size of the convolution layer set to 7*7, 7*7 can cover a larger image area, help capture more global features, and more effectively extract global information in the initial stage of the model. The number of channels is compressed to 1 through a 7*7 convolution block, and finally the channel features are nonlinearly transformed using the Sigmoid activation function to obtain the weight matrix on the space. The original feature map is multiplied with the original feature map to obtain the final output.
[0077] As another specific implementation of the transformer-based enhanced medical image segmentation network provided in the application, it also includes a loss function module, and the loss function module includes a cross entropy loss function Cross Entrop and a sample similarity function Dice, and the formula is:
[0078] L=λ 1 L Dice +λ 2 L CE
[0079] Among them, L is the total loss value, and λ1 and λ2 represent different weighting coefficients.
[0080] In some feasible embodiments, both λ1 and λ2 may be set to 0.5.
[0081] By setting the loss function module, the cross entropy loss function and the sample similarity loss function Dice are used together to achieve comprehensive optimization of global and local features. Cross entropy focuses on global pixel-by-pixel classification accuracy, while Dice focuses on regional consistency. The combination of the two can achieve a balance between global and local feature representation. Dice loss is particularly prominent in small target segmentation tasks and can effectively make up for the shortcomings of cross entropy in the case of imbalanced categories. The loss function improves the stability of the network training process, allowing the model to converge faster in the early stages and achieve better final segmentation performance.
[0082] Although the present invention has been described in detail with reference to the aforementioned embodiments, it is still possible for those skilled in the art to modify the technical solutions described in the aforementioned embodiments, or to make equivalent substitutions for some of the technical features therein. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the protection scope of the present invention.
Claims
1. A transformer-based enhanced medical image segmentation network, comprising an encoder and a decoder, characterized in that: The encoder includes a dual-branch encoder consisting of a convolutional neural network encoder branch and a transformer encoder branch, each of the encoder branches includes three processing stages, each of the transformer encoder branch processing stages includes a region-aware attention module block and a block merging layer, and an enhanced feature fusion FFR module is provided between the convolutional neural network encoder branch and the transformer encoder branch; The encoder further comprises a ResNet+50 encoder branch, wherein the ResNet+50 encoder branch and the transformer encoder branch are used together for local context feature extraction; The decoder includes a convolutional neural network module and an upsampling module, the convolutional neural network module includes a convolution layer, a batch normalization layer and an activation function ReLU layer, and the upsampling module includes a transposed convolution layer; Among them, the formulas of the convolutional neural network encoder branch and the transformer encoder branch are: in, represents the input image, Patch(·) is the patch embedding layer, ST j (·) represents the output feature of the jth layer of the transformer encoder branch, SF j (·) represents the output characteristics of the jth layer of the convolutional neural network encoder branch, j∈{1,2,3}.
2. The transformer-based enhanced medical image segmentation network according to claim 1, characterized in that: The transformer encoder branch includes a region-aware attention calculation RAA unit, which is used to generate a context score pv. The region-aware attention calculation RAA unit includes three vector branches Q, K, and V. The vector branch Q obtains a probability distribution matrix ps through a Softmax layer. The context score pv is calculated as follows: y=pv*ReLU(I V T V )T o Among them, I K Indicates branch K, I V represents branch V, Relu is the activation function, T K , T V , T O Represents the weight value corresponding to the linear layer, and ps is the probability distribution matrix.
3. The transformer-based enhanced medical image segmentation network according to claim 2, characterized in that: It also includes an enhanced fully connected EFC module, the overall process of the enhanced fully connected EFC module is: y1=FC(RELU(LN(DWConv(FC(x in ))+FC(x in )))) Among them, x in is the output of the region-aware attention calculation RAA unit, DWConv is a DW convolutional layer with a convolution kernel size of 3×3, FC is a fully connected layer, and LN is layer normalization.
4. The transformer-based enhanced medical image segmentation network according to claim 3, characterized in that: The formula of the enhanced feature fusion FFR module is: M c (F1)=Sigmoid(MLP(AvgPool(F1))+MLP(MaxPool(F1))) M s (F2)=Sigmoid(Conv(AvgPool(F2)+MaxPool(F2))) F′=F1*M c (F1)+F2*M s (F2) Among them, F1 is the feature output of the convolutional neural network encoder branch, F2 is the feature output of the transformer encoder branch, is the global average pooling, MaxPool is the global maximum pooling, Conv is the convolution layer, and Sigmoid is the activation function.
5. The transformer-based enhanced medical image segmentation network according to claim 4, characterized in that: It also includes a loss function module, which includes a cross entropy loss function Cross Entrop and a sample similarity function Dice, and the formula is: L=λ1L Dice +λ2L CE Among them, L is the total loss value, and λ1 and λ2 represent different weighting coefficients.