Brain tumor MRI image segmentation model and method based on improved Swin UNETR network
By improving the Swin UNETR network, a variety of convolutional fusion modules and attention mechanism modules were introduced, and the problems of insufficient segmentation accuracy and low computational efficiency of brain tumor MRI images in the prior art were solved, and higher segmentation accuracy and computational efficiency were achieved.
Patent Information
- Application Number
- CN202411660643.9
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-20
- Publication Date
- 2025-03-11
AI Technical Summary
The existing medical image segmentation technology based on deep learning has problems such as insufficient accuracy, large calculation amount and low image segmentation efficiency in brain tumor MRI image segmentation, especially when dealing with different types of tumors and variability, it is difficult to achieve accurate segmentation.
Improve the Swin UNETR network, and by introducing a variety of convolutional fusion modules and attention mechanism modules, the model can capture tumor characteristics and reduce computing overhead and video memory requirements. Specifically, it includes the use of depth separable convolution and void convolution for feature extraction and focusing on key information through the SimAM attention mechanism.
It significantly improves the segmentation accuracy of complex morphology and irregular boundaries of brain tumors, improves the computing efficiency and adaptability of the model, can run efficiently in resource-constrained environments, and maintains high segmentation accuracy under different tumor types and complex image conditions.
Smart Images

Figure CN119600043A_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the technical field of image segmentation, and more specifically, relates to a brain tumor MRI image segmentation model and method based on an improved Swin UNETR network. Background Art
[0002] Medical image segmentation is a key processing step in precision medicine. It is used to automatically separate different anatomical structures, lesion areas or organs in medical images. Accurate segmentation and positioning of diseased areas are crucial during surgery. In recent years, with the rapid development of deep learning, deep learning-based methods have replaced traditional methods with their excellent performance and are widely used in brain tumor image segmentation.
[0003] Currently, in the field of medical image segmentation, the mainstream deep learning-based image segmentation methods include: (1) FCN (Fully Convolutional Network), which is one of the early applications in segmentation tasks. It replaces the fully connected layers in traditional convolutional neural networks with convolutional layers, allowing the model to process input images of any size and generate segmentation maps; (2) U-Net: It is a network structure designed specifically for biomedical image segmentation. It uses an encoder-decoder architecture. The encoder extracts the feature information of the image, and the decoder fuses the encoder's high resolution through jump connections to restore the spatial detail information of the image. This U-shaped structure design enables U-Net to effectively retain local details in medical images. It is currently a classic method for medical image segmentation; (3) Swin Transformer: It is a network structure that combines the self-attention mechanism with hierarchical convolution. It can effectively process large-size medical images. Its hierarchical structure allows the capture of local features and all information.
[0004] Although the above technologies have made significant progress in medical image segmentation, they also have some shortcomings. Convolutional neural networks (CNNs) are limited to local receptive fields and cannot effectively capture long-range dependencies in images, which is particularly disadvantageous when processing brain tumor data with complex shapes and varying sizes. Although Swin Transformer overcomes this limitation and can better capture long-range dependencies and global context information, its complex network structure brings high computational overhead. Especially when processing high-resolution images, although the layered strategy reduces some of the computational workload, compared with the lightweight CNN model, Swin Transformer still has higher requirements for computing resources and video memory, which limits its application in resource-constrained environments. Summary of the invention
[0005] The purpose of the present invention is to provide a brain tumor MRI image segmentation model and method based on an improved Swin UNETR network, so as to effectively solve the problems of insufficient accuracy of brain tumor MRI image segmentation in existing deep learning-based medical image segmentation technology, especially when facing different types of tumors and variability, it is difficult to achieve accurate segmentation; or the amount of calculation is large and the image segmentation efficiency is low.
[0006] In order to achieve the above object, the technical solution provided by the present invention is:
[0007] The first aspect of the present invention provides a brain tumor MRI image segmentation model based on an improved Swin UNETR network. The model is improved by using Swin UNETR as a baseline model, and specifically includes:
[0008] An encoder, used for encoding and extracting features of the MRI image to obtain a feature region of the MRI image;
[0009] Skip connections, used to pass image features extracted by the encoder to the corresponding levels of the decoder; and
[0010] The decoder is used to fuse and restore the feature information transmitted by the jump connection to obtain the image segmentation result;
[0011] The decoder includes multiple convolution fusion modules and attention mechanism modules, which are used to fuse feature information of different levels in the decoder, and the feature information obtained by multiple convolutions is focused by the attention mechanism module. Therefore, the combination of multiple convolution fusion modules and attention mechanism modules can effectively enhance the model's ability to capture tumor features, so that it can still accurately identify tumor boundaries in complex backgrounds.
[0012] According to any of the technical solutions described in the first aspect of the present invention, the convolution fusion module includes two branches, the first branch adopts depthwise separable convolution, and the second branch adopts dilated convolution, and the output features of the two branches are fused and input into the attention mechanism module.
[0013] According to any technical solution described in the first aspect of the present invention, the attention mechanism module adopts SimAm attention mechanism. By combining the deep separable convolution and SimAM attention mechanism, the computational overhead and video memory requirements of the model can also be reduced, so that it can run efficiently in a resource-constrained environment.
[0014] According to any technical solution described in the first aspect of the present invention, the encoder is composed of four levels, each level contains two Transformer blocks for extracting multi-scale image features, and the Transformer block replaces the original MLP module with the designed SMLP module, and the SMLP module includes four parallel branches, three of which are used to capture the feature information of the vertical, horizontal and front-back directions respectively through convolution operations, and the fourth branch directly retains the original information of the input features through residual connections; finally, the feature information of the four branches is fused together and processed through three-dimensional convolution to keep the number of channels unchanged. The SMLP module effectively reduces the occupation of computing resources and improves the computing efficiency of the model by optimizing the architectural design of the model.
[0015] According to any technical solution described in the first aspect of the present invention, the model uses the Adam optimizer during training, the initial learning rate is set to 0.0001, and the loss function uses soft dice loss and cross entropy loss. The calculation formula of the loss function is as follows:
[0016]
[0017] L(G,Y)=L dice (G,Y)+L ce (G,Y)
[0018] Where I represents the number of voxels, J represents the number of categories, and G i,j and Y i,j They represent the output probability and one-hot encoding true value of class j on voxel i respectively.
[0019] The second aspect of the present invention provides a brain tumor MRI image segmentation method based on an improved Swin UNETR network, comprising:
[0020] Preprocessing the brain tumor MRI image to be segmented to obtain a preprocessed brain tumor MRI image;
[0021] The preprocessed brain tumor MRI image is input into the brain tumor MRI image segmentation model based on the improved Swin UNETR network to obtain the segmentation result;
[0022] Among them, the brain tumor MRI image segmentation model based on the improved Swin UNETR network adopts the model in any technical solution of the first aspect of the present invention.
[0023] According to any technical solution described in the second aspect of the present invention, multiple convolution fusion modules are used to perform feature fusion, specifically including: extracting features from an image using a dual-branch hybrid decoder, wherein:
[0024] The first branch is used to extract and normalize the input feature map. First, a deep convolution with a convolution kernel size of 3 is used for feature extraction, followed by normalization, and the LeakyRelu activation function is used to ensure the expression of nonlinear features. Then, a point-by-point convolution with a convolution kernel size of 1 is used to ensure the consistency of the output channel with the second branch.
[0025] The second branch uses a dilated convolution with a convolution kernel of 5 to extract features from the input feature map, followed by normalization.
[0026] The output results of the first branch and the second branch are concat-joined, and then the maximum pooling and average pooling are performed respectively. Subsequently, two weight values are obtained through a 1*1*1 convolution and a sigmoid activation function. The two weight values are applied to the above two branches respectively to obtain the product of the two branches and the two weight values. The two products are added together to obtain the final output result.
[0027] According to any technical solution described in the second aspect of the present invention, the brain tumor MRI image segmentation model based on the improved Swin UNETR network segments the MRI image, and the specific process includes:
[0028] The SMLP module and the Swin Transformer block of the original Swin UNETR network are used to encode and extract features of MRI images and reduce the network parameters to obtain the characteristic regions of MRI images, which specifically include three regional features: necrotic tumor core region, enhanced tumor region, and peritumoral edema region.
[0029] Use multiple convolution fusion modules to fuse feature information at different levels to gradually restore the image resolution and reduce the number of network parameters;
[0030] The attention mechanism module is used to focus the feature information obtained by various convolutions to enhance the focus on key information, and the robustness of the model is improved through adaptive feature weight allocation;
[0031] The final segmentation result is generated through the sigmoid activation function, and the brain tumor segmentation prediction map is output.
[0032] According to any technical solution described in the second aspect of the present invention, the SMLP module includes four parallel branches, and the process of feature extraction includes:
[0033] The first branch processes feature information along the vertical direction. First, the feature tensor of the input data is reshaped into the size of WDC*H, and linearly mapped through a linear layer of a weight matrix, and then the mapped result is reshaped into (B, C, H, W, D);
[0034] The second branch processes feature information in the horizontal direction. First, the feature tensor of the input data is reshaped into the size of HDC*W, and linearly mapped through a linear layer of a weight matrix. Then, the mapped result is reshaped into (B, C, H, W, D).
[0035] The third branch processes feature information along the forward and backward directions. First, the feature tensor of the input data is reshaped into the size of HWC*D, and linearly mapped through a linear layer of a weight matrix, and then the mapped result is reshaped into (B, C, H, W, D);
[0036] The fourth branch uses residual connection to directly retain the original information;
[0037] Finally, the output results of the four branches are concat- ed, and the number of output channels becomes four times the original number. Then, a convolution with a convolution kernel size of 1*1*1 is used to make the number of channels consistent with the original feature channels.
[0038] According to any technical solution described in the second aspect of the present invention, the size of the preprocessed brain tumor MRI image is 128×128×128, and the number of channels of the output brain tumor segmentation prediction image is 3.
[0039] In summary, the technical solution provided by the present invention can achieve the following beneficial effects compared with the prior art:
[0040] (1) The present invention improves on the existing Swin UNETR network model by introducing the DDSA module, that is, by using a variety of convolutional fusion modules and attention mechanism modules in combination, thereby significantly improving the segmentation accuracy of brain tumors with complex morphology and irregular boundaries, especially being able to be effectively used for processing different types of tumors and variability.
[0041] (2) The present invention optimizes the encoder structure of the model, that is, replaces the original MLP module of the Transformer block with the designed SMLP module, thereby effectively reducing the occupation of computing resources and improving the computing efficiency of the model while ensuring the performance and stability of the model.
[0042] (3) The present invention can extract and enhance global and local features through the coordinated use of DDSA and SMLP modules, so that the network can better capture global context information and local details, effectively improving the feature extraction ability of the model; and the model has stronger adaptability, enhanced adaptability to different tumor types, especially under complex image conditions, and can maintain high segmentation accuracy. At the same time, the improved network greatly reduces the computational complexity, and through lightweight designs such as deep separable convolution and dilated convolution, the efficiency of the model in processing high-resolution images is improved, which is suitable for large-scale medical image processing. BRIEF DESCRIPTION OF THE DRAWINGS
[0043] Figure 1 This is an example diagram of the model structure provided in the embodiments of the present application.
[0044] Figure 2 This is an example diagram of the BraTS2020 brain tumor data provided in the embodiments of the present application.
[0045] Figure 3 A schematic diagram of the loss change trend of the model provided in the embodiment of the present application during the training phase.
[0046] Figure 4 Schematic diagram of the DDSA technology of the model provided in the embodiment of the present application.
[0047] Figure 5 Schematic diagram of the SMLP technology of the model provided in the embodiment of the present application.
[0048] Figure 6 A visualization diagram of the segmentation technology provided in an embodiment of the present application. DETAILED DESCRIPTION
[0049] The main purpose of the present invention is to solve at least one of the following technical problems existing in the prior art: (1) Insufficient accuracy of brain tumor MRI image segmentation: The existing brain tumor MRI image segmentation method has limited accuracy when dealing with tumors with complex morphology and irregular boundaries, especially when facing different types of tumors and variability, it is difficult to achieve accurate segmentation; (2) The feature extraction capability of the deep learning network is not strong enough: When processing high-dimensional data, conventional convolutional neural networks (CNNs) have difficulty in fully extracting global and local features, resulting in deficiencies in the recognition and segmentation of brain tumor morphology; (3) Low computational efficiency: Existing segmentation models may have high computational complexity, especially when processing large-scale MRI data, there are problems of slow processing speed and high resource consumption, which limits its feasibility in clinical applications; (4) Insufficient combination of local features and global context information: When processing brain tumor images, many existing models do not deeply combine local features and global information, and cannot effectively improve the accuracy and robustness of image segmentation.
[0050] Based on the above purposes, the present invention can improve the network's ability to recognize complex morphologies of brain tumors by combining a multi-scale attention mechanism and a multi-layer perception module, thereby solving the problem of insufficient accuracy in the segmentation of brain tumor images in the prior art; and by simultaneously introducing DDSA (Dilated Depthwise Separable Convolution with SimAM Attention) and SMLP (Special MLP) modules in the improved Swin UNETR network, it is beneficial to better extract global and local feature information and enhance the ability to segment the boundaries and internal structures of brain tumors. At the same time, by introducing improved modules, the computational complexity of the model can also be reduced, and the operational efficiency of the model can be improved while ensuring accuracy, which is suitable for large-scale medical image data processing; through the improved network structure, the global context information and local detail features can be better integrated, further improving the adaptability of the model under different tumor types and complex image conditions.
[0051] Specifically, the brain tumor MRI image segmentation method based on the improved Swin UNETR network of the embodiment of the present invention includes:
[0052] Preprocessing the brain tumor MRI image to be segmented to obtain a preprocessed brain tumor MRI image;
[0053] The preprocessed brain tumor MRI image is input into the brain tumor MRI image segmentation model based on the improved Swin UNETR network to obtain the segmentation result.
[0054] Among them, the brain tumor MRI image segmentation model based on the improved Swin UNETR network includes an encoder, a decoder and skip connections. The encoder is used to extract multi-scale image features, and the skip connection is used to transfer the image features extracted by the encoder to the corresponding level of the decoder. The decoder is used to perform feature fusion and restoration to obtain image segmentation results.
[0055] Specifically, the brain tumor MRI image segmentation model of the embodiment of the present invention uses the Swin UNETR network as the baseline model, and optimizes the model to improve the segmentation accuracy of brain tumor MRI images. Figure 1 As shown in the figure, input represents the input brain tumor data. The entire network adopts a classic U-shaped structure design, which aims to generate a detailed segmentation result by encoding and decoding the image features layer by layer. The input of the model undergoes multi-level feature extraction and reconstruction by the encoder and decoder to complete the segmentation of the brain tumor area.
[0056] In order to facilitate a clearer understanding of the solution of the present invention, the Swin UNETR model is briefly introduced below. First, the MRI images of the four modalities are connected together to obtain an input H×W×D×4, and then the input image is divided into multiple non-overlapping patches of size H′×W′×D′, where the patch size is (2,2,2).
[0057] For each patch block, it is mapped to a token with embedding dimension C, resulting in a size of The three-dimensional token sequence is fed into the encoder of the model for feature extraction. The encoder contains 4 layers of Swin Transformer Block. Global feature extraction is performed through 4 layers of Swin Transformer Block. Each layer of Block contains two parts: linear embedding block and patch merging. The calculation formula of linear embedding block is as follows:
[0058]
[0059] W-MSA and SW-MSA are multi-head attention modules divided by window and moving window respectively, MLP and LN represent multi-layer perceptron and layer normalization respectively, l represents the current layer, l-1 represents the previous layer, and l+1 represents the next layer. represents the output of W-MSA and SW-MSA, Represents the output after MLP and residual operation. The calculation formula of self-attention is:
[0060]
[0061] Where Q, K, and V represent query, key, and value respectively; d represents the size of the query and key.
[0062] Based on the existing Swin UNETR network model, combined with Figure 1 As shown, in some embodiments, the encoder is composed of four levels, each level includes two Transformer blocks for extracting multi-scale image features. Unlike the traditional Transformer block, the present invention introduces an improvement in the Transformer block, replacing the original MLP layer (multi-layer perceptron) with the designed SMLP layer, thereby effectively improving the feature expression ability and computational efficiency of the network. Each level processes the input features through a patch merging layer, gradually reducing the spatial resolution and increasing the channel dimension of the features. After each encoding layer, the spatial size of the feature map is halved, and the number of channels is doubled, which enables the model to gradually capture more abstract semantic information while extracting low-level details.
[0063] Specifically, Figure 5 As shown in the figure, the SMLP module includes four parallel branches, three of which are used to capture the feature information in the vertical, horizontal and front-to-back directions respectively through convolution operations, and the fourth branch directly retains the original information of the input features through residual connections; finally, the feature information of these four branches is fused together and processed through three-dimensional convolution to keep the number of channels unchanged.
[0064] Furthermore, suppose X in ∈R B×H×W×D×C Represents the tensor input to the SMLP module, where B represents the batch size. During the model training process, one sample data is processed each time. The first branch of the SMLP module processes the feature information along the vertical direction. First, the feature tensor of the input data is reshaped to the size of WDC*H and passed through a weight matrix W W ∈R W ×W×W The linear layer performs linear mapping, and then reshapes the mapped result into (B, C, H, W, D);
[0065] The second branch processes feature information in the horizontal direction. First, the feature tensor of the input data is reshaped into the size of HDC*W and then passed through a weight matrix W W ∈R W×W×W The linear layer performs linear mapping, and then reshapes the mapped result into (B, C, H, W, D);
[0066] The third branch processes feature information along the forward and backward directions. First, the feature tensor of the input data is reshaped into the size of HWC*D and then passed through a weight matrix W W ∈R W×W×W The linear layer is linearly mapped, the weights of these three weight matrices are different, and then the mapped results are reshaped into (B, C, H, W, D);
[0067] The fourth branch uses residual connection to directly retain the original information;
[0068] Finally, the output results of the four branches are concat- ed, and the number of output channels becomes four times the original number, that is, a tensor with a shape of B×H×W×D×4C is formed; in order to maintain the consistency of the number of channels, a convolution with a convolution kernel size of 1*1*1 is then used to make the number of channels consistent with the original feature channels, that is, compress the number of channels from 4C to C, so as to obtain the final output. This not only effectively integrates feature information from different directions, but also reduces the computational complexity, thereby improving the training efficiency of the model.
[0069] The total number of parameters of an SMLP module is H 2 +W2 +D 2 +4C 2 , of which 4C 2 represents the parameters of the four-path fusion operation. In contrast, the number of parameters of the MLP module is 2α(HWD) 2 , where α represents the expansion layer in the MLP layer and is usually set to 4. Since the SMLP module introduces linear mapping and fusion operations in the direction in its design, its computational complexity is significantly lower than that of the traditional MLP layer. Specifically, the computational complexity of the SMLP module is:
[0070] Ω(SMLP)=HWDC(H+W+D+4C)
[0071] The computational complexity of the original MLP layer is:
[0072] Ω(MLP)=2α(HWD) 2 C
[0073] Assuming that N represents the product of HWD, the complexity of the traditional MLP module is O(N 2 ), and the complexity of the modified SMLP is From the comparison of complexity, we can see that the SMLP module has significant performance advantages when processing large-scale data. This shows that compared with the traditional MLP module, it significantly reduces the computational complexity. The SMLP module can not only efficiently process larger data sets, but also show better optimization effects in terms of computing resources and time consumption. The loss function uses soft dice loss and cross entropy loss. The calculation formula of the loss function is as follows:
[0074]
[0075] L(G,Y)=L dice (G,Y)+L ce (G,Y)
[0076] Where I represents the number of voxels, J represents the number of categories, and G i,j and Y i,j They represent the output probability and one-hot encoding true value of class j on voxel i respectively.
[0077] The decoder part of the model in the present invention corresponds to the encoder and has the same deep structure, ensuring that features can be seamlessly transferred between different levels of resolution. The features extracted in the encoder are transferred to the corresponding level of the decoder through SkipConnections. This design can effectively fuse high-resolution detail information and low-resolution semantic information.
[0078] In some embodiments, during the decoding phase, the features of the jump connection are first processed by an improved DDSA module (Dual-Dimensional Self-Attention). The DDSA module integrates an attention mechanism that can learn and integrate features in both spatial and channel dimensions. Through this module, the model can more efficiently capture global information and improve the quality of feature expression. The introduction of the DDSA module is to enhance the model's understanding of multi-scale information during the decoding process and ensure that the semantic features extracted by the encoder can be fully utilized.
[0079] During the decoding process of each layer, the model concatenates the features decoded from the current layer with the features of the previous layer after upsampling. Through this concatenation operation, the decoder can not only utilize the information of the current layer, but also fuse the lower-resolution features of the previous layer to ensure the transmission and integration of multi-scale information. Each upsampling operation is implemented through deconvolution, which doubles the spatial size of the feature map and halves the number of channels. In this process, the model gradually restores the spatial resolution and detail information of the original image, ensuring that the final segmentation result can retain more image details.
[0080] Specifically, Figure 4 As shown in the figure, the DDSA module includes multiple convolution fusion modules and attention mechanism modules. Multiple convolutions are used to fuse feature information at different levels in the decoder of the network, gradually restore the resolution of the image, and reduce the number of network parameters; the attention mechanism is used to focus the feature information obtained by multiple convolutions, enhance the focus on key information, and improve the robustness of the model through adaptive feature weight allocation; the last layer of the decoder maps the decoded features to the output space through a 1×1×1 convolution operation. The function of this convolution layer is to generate the required number of channels and further refine the features. Finally, the model generates the final segmentation result through the sigmoid activation function, and outputs a brain tumor segmentation prediction map with a channel number of 3. This output result corresponds to different types of tumor areas in the image.
[0081] The multiple convolution fusion module includes two branches, the first branch uses depthwise separable convolution, and the second branch uses dilated convolution. The output features of the two branches are fused and input into the attention mechanism module. In order to facilitate understanding of the solution of the present invention, the depthwise separable convolution is introduced first.
[0082] Depthwise convolution decomposes the standard convolution into two simpler steps: depthwise convolution and pointwise convolution. This operation greatly reduces the number of parameters and the amount of calculation, so it is widely used in lightweight neural network architectures. The operation of standard convolution is: assuming that the input feature map has multiple channels, each channel will be convolved with a convolution kernel, and then these results will be aggregated to generate the output feature map. Assuming that the input feature map has C in channels, and the output feature map has C out channels, the size of the convolution kernel is K×K, then the computational complexity of the standard convolution is:
[0083] K×K×C in ×C out ×H×W
[0084] Where H and W are the spatial sizes of the input feature map.
[0085] Depthwise separable convolution decomposes the above standard convolution into two steps:
[0086] (1) Depthwise Convolution: This step uses a convolution kernel to perform convolution on each input channel separately, that is, an independent K×K convolution kernel is applied to each channel. In this way, different input channels will not be mixed together, and the number of output feature maps is the same as the number of input channels. The computational complexity of this operation is:
[0087] K×K×C in ×H×W
[0088] (2) Pointwise convolution: This step uses a 1×1 convolution kernel (size 1×1) to linearly combine all channels to change the number of output channels. The computational complexity of pointwise convolution is:
[0089] C in ×C out ×H×W
[0090] By decomposing the standard convolution into depthwise convolution and pointwise convolution, the total computation becomes:
[0091] K×K×C in ×H×W+C in ×C out ×H×W
[0092] Compared with the computational complexity of standard convolution, the computational complexity of depthwise separable convolution is significantly reduced, especially when K is large (such as 3×3 convolution). Moreover, more than 90% of the main convolutions in the model are convolutions with a kernel of 3.
[0093] The dilated convolution is a special convolution operation that expands the receptive field by introducing holes (i.e., gaps) inside the convolution kernel, thereby capturing a wider range of contextual information while keeping the amount of computation low. The size of the convolution kernel is "expanded" by inserting gaps between the elements of the convolution kernel. The dilation rate determines the spacing between the elements of the convolution kernel. For example, in a 3×3 convolution kernel, if the dilation rate is 1, then this is an ordinary convolution; if the dilation rate is 2, then the convolution kernel will insert a one-pixel gap between each element, and its receptive field is equivalent to a 5×5 convolution. Assuming the size of the convolution kernel is K×K and the dilation rate is d, the actual receptive field is:
[0094] (K+(K-1)×(d-1))×(K+(K-1)×(d-1))
[0095] By adjusting the dilation rate d, the receptive field can be increased without increasing the number of convolution kernel parameters.
[0096] In the embodiment of the present invention, multiple convolution fusion modules are used to perform feature fusion, specifically including: extracting features from an image using a dual-branch hybrid decoder, wherein:
[0097] The first branch is used to extract and normalize the input feature map. First, a deep convolution with a convolution kernel size of 3 is used for feature extraction, followed by normalization, and the LeakyRelu activation function is used to ensure the expression of nonlinear features. Then, a point-by-point convolution with a convolution kernel size of 1 is used to ensure the consistency of the output channel with the second branch.
[0098] The second branch uses a dilated convolution with a convolution kernel of 5 to extract features from the input feature map, followed by normalization.
[0099] The output of the first branch is concatenated with the output of the second branch, and then the maximum pooling and average pooling are performed respectively. Then, two weight values are obtained through a 1*1*1 convolution and a sigmoid activation function, which are applied to the above two branches respectively to obtain the product of the two branches and the two weight values, and the two products are added together to obtain the final output result. Through the combination of the above two branches, the decoder can not only process image features more efficiently, but also capture important information in a larger range, thereby improving the overall segmentation performance.
[0100] As a preferred implementation, in some embodiments, the attention mechanism adopts the SimAM attention mechanism, which is a lightweight and parameter-free attention mechanism designed to help the model focus on the key areas of the image. Generally, there is a strong similarity between adjacent pixels in an image, while the similarity between distant pixels is weaker. SimAM uses this feature to generate a three-dimensional attention weight by calculating the similarity between each pixel of the feature map and its adjacent pixels. The calculation process of SimAM is as follows:
[0101]
[0102] where w t represents the attention weight of the t-th pixel, k is a normalization constant, N t is the set of neighboring pixels of the t-th pixel, s(f t ,f u ) is the Euclidean distance, and the formula is as follows:
[0103]
[0104] By introducing the SimAM attention mechanism, the model can not only selectively enhance the key features in the image, but also does not increase the number of parameters of the original network, nor does it bring additional training burden to the model. The parameter-free design of this mechanism ensures the improvement of computational efficiency and model performance. In the processing of the decoder, the features extracted by the residual connection and the dual branch are fused through a 3×3×3 convolution and normalization operation. Specifically, the residual connection retains the original feature information, while the dual branch structure extracts features at different scales and perspectives. Through convolution and normalization, these features are effectively integrated to ensure the comprehensiveness and consistency of the information. Finally, this fusion process generates the output of the decoder layer, which provides high-quality feature representation for subsequent feature processing and detail recovery.
[0105] It should be noted that the main improvement of the present invention is that on the basis of the Swin UNETR model, the SMLP module and the DDSA module are introduced, wherein only the SMLP module can be introduced, only the DDSA module can be introduced, or both the SMLP module and the DDSA module can be introduced at the same time, preferably both the SMLP module and the DDSA module are introduced at the same time, and the preferred MRI image segmentation method specifically includes the following operations:
[0106] The SMLP module and the Swin Transformer block of the original Swin UNETR network are used to encode and extract features of MRI images and reduce the network parameters to obtain the characteristic regions of MRI images, which specifically include three regional features: necrotic tumor core region, enhanced tumor region, and peritumoral edema region.
[0107] Use multiple convolution fusion modules to fuse feature information at different levels to gradually restore the image resolution and reduce the number of network parameters;
[0108] The attention mechanism module is used to focus the feature information obtained by various convolutions to enhance the focus on key information, and the robustness of the model is improved through adaptive feature weight allocation;
[0109] The final segmentation result is generated through the sigmoid activation function, and the brain tumor segmentation prediction map is output.
[0110] Model testing experiment:
[0111] For the experiment, the model was trained from scratch without using a pre-trained model, verified every 5 rounds, and the model checkpoints were saved. At the beginning of model training, the weights were initialized and the random number seed was set to ensure the accuracy of the experiment.
[0112] The embodiment of the present invention uses the BraTS2020 dataset for training. BraTS2020 is a medical image dataset for brain tumor segmentation and survival prediction, which is jointly provided by The Cancer Imaging Archive (TCIA) and BraTS. The BraTS2020 dataset includes 369 training data and 125 validation data sets. Each sample data contains four modalities of magnetic resonance imaging, namely, T1-weighted image (T1), contrast-enhanced T1-enhanced T1-weighted image (T1ce), T2-weighted image (T2), and fluid attenuation recovery image (FLAIR). The dimension of each modality is 240×240×155, and they are aligned to the same space. The labels of the dataset are divided into four types: label 0 represents the background, label 1 represents the necrotic and non-enhancing tumor area, label 2 represents the edema area around the tumor, and label 4 represents the enhancing tumor area. Finally, three sub-areas were identified, of which labels 1, 2 and 4 were used to evaluate the whole tumor (whole tumor, WT), labels 1 and 4 were combined to determine the tumor core (tumor core, TC), and label 4 was used to determine whether the tumor was enhancing (enhancing tumor, ET).
[0113] During the training process, the training set was randomly divided into a training set and a validation set for training at a ratio of 4:1. During the training process, the validation was performed every five rounds. Finally, the trained model was used to perform segmentation prediction on 125 official validation sets. Figure 2 An example of this dataset is shown. Since the validation set does not provide true labels, the trained model is used to segment the validation set to generate predicted segmentation maps. These segmentation results are then uploaded to the CBICA online evaluation platform, where the accuracy of the predicted segmentation maps is evaluated.
[0114] Specifically, the original 3D brain tumor MRI images were randomly cropped into a size of 128×128×128 as the input size of the model, and the training iterations were 300 times. The Adam (Adaptive Moment Estimation) optimizer was used, and the initial learning rate was set to 0.0001. In order to accelerate convergence and improve model performance, a linear warm-up of 50 epochs and a cosine annealing algorithm learning rate scheduler were used to adjust the learning rate. The batch size of the GPU was set to 1; the sliding window method was used for inference, and the overlap of adjacent voxels was 0.6; random cropping, flipping, scaling, and offset were used for data enhancement. The specific hyperparameter settings are shown in Table 1:
[0115] Table 1. Training hyperparameters
[0116]
[0117] exist Figure 3 The figure shows the loss and accuracy trend of the model during the training phase based on the above training configuration. It can be clearly observed from the figure that as the training progresses, the loss curve gradually stabilizes, indicating that the model is gradually converging. In addition, during the training process, the model weights and parameters are saved every five rounds, which lays a solid foundation for the subsequent model reasoning steps.
[0118] Under the same experimental setting, some models based on existing convolutional neural networks and Transformer were trained and tested. All results were submitted to the official website for online verification. The evaluation indicators adopted the official standards of the BraTS challenge, including the Dice similarity coefficient and 95% Hausdorff distance in the brain tumor segmentation task, where G represents the true label, P represents the segmentation result of the model, and TP, FP, TN and FN represent true positive, false positive, true negative and false negative. The Dice coefficient is a measure of the similarity between two sets. The Dice similarity coefficient formula is as follows:
[0119]
[0120] The Hausdorff distance represents the maximum distance between the predicted segmentation region boundary and the actual tumor region boundary. The smaller the value, the smaller the brain tumor boundary segmentation error and the better the quality. The Hausdorff distance formula is as follows:
[0121]
[0122] According to the comparative experiments, it can be seen that in terms of segmentation performance, the model of this application performs better in the entire tumor, tumor core and enhanced tumor area. As shown in Table 2 below, the Dice scores of the model of this application (SMLP+DDSA) reached 76.11%, 89.53% and 81.58% in the three regions of ET, WT and TC, respectively, which are 1.22%, 1.11% and 2.47% higher than the basic model. In terms of the HD95 index, the distances of the model of this application in the ET, WT and TC regions are 21.78, 4.91 and 9.14, respectively, which are 8.61, 3.39 and 3.1 less than the basic model in the ET, WT and TC regions, respectively, showing better boundary segmentation effect. It can be seen from the results in the table that the model of this application has significantly improved the segmentation performance compared with traditional convolutional neural network-based methods (such as 3D U-Net and Attention U-Net). At the same time, compared with models that combine convolutional neural networks and Transformer architectures (such as TransBTS, UNETR, and SDV-TUNet), the model of this application also shows more obvious performance advantages. These results show that the model of this application has stronger adaptability and boundary processing capabilities in tumor segmentation tasks.
[0123] Table 2 Experimental results of different models on the BraTS2020 validation set
[0124]
[0125] In order to better verify that the proposed module can optimize the segmentation performance of the model, ablation experiments were performed on the BraTS2020 validation set. All ablation experiments used the same experimental environment and hyperparameters to ensure the reliability and comparability of the results.
[0126] As shown in Table 3 below, the introduction of different modules has a significant impact on the segmentation performance of the model. By adding the DDSA module, the Dice scores of the three regions of brain tumors increased by 0.32%, 1.16% and 2.23%. This improvement shows that the DDSA module can effectively enhance the model's ability to capture tumor features, so that it can still accurately identify tumor boundaries in complex backgrounds. This module uses different convolution mechanisms to extract key features of the image, while the spatial attention mechanism can focus on important local areas, thereby further improving the segmentation accuracy. In addition to the improvement of segmentation performance, the DDSA module also significantly reduces the number of model parameters. Compared with traditional convolution, it reduces redundant calculations and optimizes the model's computational efficiency without sacrificing accuracy through efficient feature extraction. In addition, the introduction of the SMLP module also plays a key role in the optimization of the model. The traditional multi-layer perceptron structure often requires a large number of parameters to capture nonlinear features, while the SMLP module has brought about a significant reduction in the number of model parameters through a more streamlined structural design, and the performance of the model remains stable. This shows that the SMLP module effectively reduces the occupation of computing resources and improves the computational efficiency of the model by optimizing the architectural design of the model. At the same time, by combining the DDSA and SMLP modules, the overall number of parameters of the model is reduced by one third compared to the original. Although the number of parameters is significantly reduced, the segmentation performance of the model is not negatively affected, but improved instead, which is most obvious in the core area of the tumor.
[0127] Table 3 The impact of different modules on the number of model parameters and Dice score results
[0128]
[0129] Since the BraTS2020 dataset does not provide the true labels of the validation set, the data of the training set was used for visualization analysis. In order to more intuitively compare the segmentation effects of different models, Figure 6 The segmentation visualization effect of the present invention compared with several other representative technical models is demonstrated.
[0130] Among them, the first row shows the axial slices of brain tumors, the second row is the sagittal slices, and the third row is the coronal slices. Different colors in the image represent different lesion areas: the red area represents the necrotic tumor core (NCR / NET), the yellow area represents the enhanced tumor (ET), and the green area represents the peritumoral edema (ED). It can be clearly seen from the figure that compared with other methods, the segmentation prediction results of the present invention are closer to the true segmentation results. Among them, in the horizontal plane direction, 3D U-Net and Attenion-Unet have inaccuracies in the segmentation of brain tumor boundaries, and in the sagittal and coronal planes, the model also has a large difference between the segmentation of the necrotic tumor core and the true label. Compared with other methods, the model is more accurate in the segmentation of boundary information. This advantage is due to the DDSA module introduced in the model, which not only effectively expands the receptive field of the network, but also reduces the problem of inaccurate segmentation caused by the limitation of the receptive field. Moreover, it is outstanding in the retention of edge information, the segmentation results are more in line with the true label, and the detail information is better retained, thereby improving the accuracy of the overall segmentation.
Claims
1. A brain tumor MRI image segmentation model based on an improved Swin UNETR network, characterized in that: This model is improved based on Swin UNETR as the baseline model, including: An encoder, used for encoding and extracting features of the MRI image to obtain a feature region of the MRI image; Skip connections, used to pass image features extracted by the encoder to the corresponding levels of the decoder; and The decoder is used to fuse and restore the feature information transmitted by the jump connection to obtain the image segmentation result; Among them, the decoder includes multiple convolution fusion modules and attention mechanism modules, which use multiple convolution fusion modules to fuse feature information of different levels in the decoder, and use the attention mechanism module to focus the feature information obtained by multiple convolutions.
2. According to claim 1, the brain tumor MRI image segmentation model based on the improved Swin UNETR network is characterized in that: The convolution fusion module includes two branches. The first branch adopts depthwise separable convolution, and the second branch adopts dilated convolution. The output features of the two branches are fused and input into the attention mechanism module.
3. The brain tumor MRI image segmentation model based on the improved Swin UNETR network according to claim 1 is characterized in that: The attention mechanism module adopts the SimAm attention mechanism.
4. The brain tumor MRI image segmentation model based on the improved Swin UNETR network according to any one of claims 1 to 3, characterized in that: The encoder consists of four layers, each of which contains two Transformer blocks for extracting multi-scale image features. The Transformer block replaces the original MLP module with the designed SMLP module, which includes four parallel branches, three of which are used to capture feature information in the vertical, horizontal, and front-to-back directions through convolution operations, and the fourth branch directly retains the original information of the input features through residual connections; finally, the feature information of these four branches is fused together and processed through three-dimensional convolution to keep the number of channels unchanged.
5. The brain tumor MRI image segmentation model based on the improved Swin UNETR network according to claim 4 is characterized in that: The model uses the Adam optimizer during training, with the initial learning rate set to 0.0001. The loss function uses soft dice loss and cross entropy loss. The calculation formula of the loss function is as follows: L(G,Y)=L dice (G,Y)+L ce (G,Y) Where I represents the number of voxels, J represents the number of categories, and G i,j and Y i,j They represent the output probability and one-hot encoding true value of class j on voxel i respectively.
6. A brain tumor MRI image segmentation method based on an improved Swin UNETR network, characterized in that: include: Preprocessing the brain tumor MRI image to be segmented to obtain a preprocessed brain tumor MRI image; The preprocessed brain tumor MRI image is input into the brain tumor MRI image segmentation model based on the improved Swin UNETR network to obtain the segmentation result; Wherein, the brain tumor MRI image segmentation model based on the improved Swin UNETR network adopts the model described in any one of claims 1-5.
7. The image segmentation method according to claim 6, characterized in that: Multiple convolutional fusion modules are used for feature fusion, specifically including: using a dual-branch hybrid decoder to extract features from the image, where: The first branch is used to extract and normalize the input feature map. First, a deep convolution with a convolution kernel size of 3 is used for feature extraction, followed by normalization, and the LeakyRelu activation function is used to ensure the expression of nonlinear features. Then, a point-by-point convolution with a convolution kernel size of 1 is used to ensure the consistency of the output channel with the second branch. The second branch uses a dilated convolution with a convolution kernel of 5 to extract features from the input feature map, followed by normalization. The output results of the first branch and the second branch are concat-joined, and then the maximum pooling and average pooling are performed respectively. Subsequently, two weight values are obtained through a 1*1*1 convolution and a sigmoid activation function. The two weight values are applied to the above two branches respectively to obtain the product of the two branches and the two weight values. The two products are added together to obtain the final output result.
8. The image segmentation method according to claim 6 or 7, characterized in that: The brain tumor MRI image segmentation model based on the improved Swin UNETR network segments the MRI image, and the specific process includes: The SMLP module and the Swin Transformer block of the original Swin UNETR network are used to encode and extract features of MRI images and reduce the network parameters to obtain the characteristic regions of MRI images, which specifically include three regional features: necrotic tumor core region, enhanced tumor region, and peritumoral edema region. Use multiple convolution fusion modules to fuse feature information at different levels to gradually restore the image resolution and reduce the number of network parameters; The attention mechanism module is used to focus the feature information obtained by various convolutions to enhance the focus on key information, and the robustness of the model is improved through adaptive feature weight allocation; The final segmentation result is generated through the sigmoid activation function, and the brain tumor segmentation prediction map is output.
9. The image segmentation method according to claim 8, characterized in that: The SMLP module includes four parallel branches, and the process of feature extraction includes: The first branch processes feature information along the vertical direction. First, the feature tensor of the input data is reshaped into the size of WDC*H, and linearly mapped through a linear layer of a weight matrix, and then the mapped result is reshaped into (B, C, H, W, D); The second branch processes feature information in the horizontal direction. First, the feature tensor of the input data is reshaped into the size of HDC*W, and linearly mapped through a linear layer of a weight matrix. Then, the mapped result is reshaped into (B, C, H, W, D). The third branch processes feature information along the forward and backward directions. First, the feature tensor of the input data is reshaped into the size of HWC*D, and linearly mapped through a linear layer of a weight matrix, and then the mapped result is reshaped into (B, C, H, W, D); The fourth branch uses residual connection to directly retain the original information; Finally, the output results of the four branches are concat- ed, and the number of output channels becomes four times the original number. Then, a convolution with a convolution kernel size of 1*1*1 is used to make the number of channels consistent with the original feature channels.
10. The image segmentation method according to claim 6 or 7, characterized in that: The size of the preprocessed brain tumor MRI image is 128×128×128, and the number of channels of the output brain tumor segmentation prediction image is 3.
Citation Information
Patent Citations
Medical image segmentation method based on joint encoder network GPA-TUNet
CN117197151A
Eye fundus blood vessel image segmentation method, terminal device and storage medium
CN117437249A
Medical image segmentation system
CN117593275A
Brain tumor image segmentation method based on multi-scale convolution and Mama structure
CN118447244A
System and method for 3D medical image segmentation
US20240362788A1
Cited By
Intracranial artery lumen wall segmentation method based on HRMRI image
CN120259343A
A method for intracranial artery lumen and wall segmentation based on HRMRI images
CN120259343B
Eye fundus image lesion segmentation method and device, computer equipment and medium
CN120318249A
MRI brain tumor segmentation method
CN120451568A
Construction method of medical image classification and segmentation model and related application thereof
CN120783101A