Brain MRI image segmentation method based on RWKV model

Abstract

The invention discloses a brain MRI image segmentation method based on an RWKV model, and belongs to the technical field of medical image processing and artificial intelligence crossing. Firstly, an RWKV linear self-attention module is introduced into a U-Net network framework, and association between remote pixels is established with relatively low calculation overhead, so that the recognition precision of a brain tumor area is improved, the reasoning time is effectively controlled, and the practicability of a model is enhanced. Secondly, according to the method, multi-scale coding and a feature fusion mechanism are combined, local and global information is extracted in a combined manner, features are mined from different resolution levels, and adaptive fusion is realized, so that the fine-grained segmentation effect is improved. And finally, in order to enhance the generalization ability and lightweight deployment performance of the model, the network structure is further optimized, and the overall computing resource demand is reduced, so that the method can better adapt to cross-patient MRI data in a complex clinical environment, and has good practical value and popularization potential.

Description

Technical Field

[0001] The present invention belongs to the cross - technical field of medical image processing and artificial intelligence, and particularly relates to a method for segmenting brain MRI images based on the RWKV model. Background Art

[0002] With the development of medical imaging technology, MRI (Magnetic Resonance Imaging) has become an important means for diagnosing brain diseases. Brain tumors vary greatly in morphology, size, and location, and the manual annotation process is complex and time - consuming. If the precise positioning of the brain tumor area can be achieved by means of automatic segmentation technology, it will help improve the diagnostic efficiency and the formulation of subsequent treatment plans.

[0003] In medical image segmentation tasks, traditional methods mainly rely on techniques based on image processing and pattern recognition, such as region growing, threshold segmentation, edge detection, active contour models (Active Contour), and graph cut (Graph Cut), etc. These methods usually require fine - tuned parameters or manually - designed rules, rely strongly on image quality, contrast, and prior knowledge, and have poor generality and robustness. At the same time, traditional methods are inefficient when dealing with large - scale data and are difficult to meet the requirements of automated processing of high - resolution medical images.

[0004] With the development of artificial intelligence, the convolutional neural network (CNN) has been introduced into the field of medical image analysis, initiating the research on image segmentation based on deep learning. The CNN structure realizes the gradual abstraction of local features of images through operations such as stacking convolutional layers and pooling layers, and has strong representation ability. Networks such as FCN and U - Net have shown good segmentation effects in multiple medical image tasks, significantly improving the accuracy and automation level, and becoming the mainstream AI methods. However, CNN has the problems of limited receptive field and difficulty in modeling long - range dependencies. When encountering brain MRI images with blurred boundaries, large - scale lesion diffusion, or complex tissue structures, the CNN model often has difficulty capturing global context information and is prone to inaccurate segmentation or regional misjudgment.

[0005] To alleviate this problem, the Transformer structure has been introduced into medical image segmentation, using its self - attention mechanism to establish global dependencies within the image. Representative models such as TransUNet and Swin - UNet have achieved good performance on multiple public medical image datasets. However, these models generally have problems of high computational complexity, slow inference speed, and large video memory occupancy, especially in the scenario of high - resolution MRI image processing, which limits their wide application in clinical practice.

[0006] In this context, the RWKV (Receptance Weighted Key Value) structure, as a new type of sequence modeling architecture, combines the linear temporal modeling efficiency of RNNs with the global modeling ability of Transformers, becoming a key area of focus in recent years. Existing research has attempted to embed Vision-RWKV into image segmentation networks to replace traditional self-attention modules, significantly reducing computational and memory costs while maintaining modeling capabilities. Preliminary experiments indicate that the application of RWKV in medical images has good prospects, especially when dealing with high-resolution, long sequences, or large-sized inputs. Therefore, by combining the RWKV linear self-attention mechanism with the U-Net structure, we aim to improve segmentation accuracy while reducing the computational burden to meet the actual clinical needs.

[0007] Although existing medical image segmentation methods have made significant progress, they still face several technical bottlenecks in the automatic segmentation of brain MRI images, mainly reflected in the following dimensions:

[0008] Limited segmentation accuracy and difficulty in handling complex lesion structures: Most traditional and deep learning methods, such as U-Net, rely on local convolutional operations for feature extraction and are difficult to effectively capture the semantic connections between remote pixels in the image. When faced with brain tumor regions with complex structural forms and blurred boundaries, problems such as missed segmentation, over-segmentation, or blurred boundaries often occur, affecting the accuracy and integrity of the segmentation results.

[0009] Low computational efficiency and high resource consumption: Although the self-attention mechanism can compensate for the local receptive field limitations of CNNs, structures such as Transformers have a quadratic growth in computational complexity (O(n 2 )) when processing high-resolution images and severely rely on high-computing power hardware. This high cost makes it difficult to deploy the model in restricted environments such as hospital workstations and mobile devices, limiting its practical application value.

[0010] Poor model generalization and insufficient stability: Brain images often contain multi-level and multi-scale organizational structures. Existing methods mostly use simple splicing, channel stacking, etc. for feature fusion, lacking deep semantic interaction and long-range relationship modeling, resulting in the difficulty of accurately depicting fine-grained structures (such as tumor edges) and affecting the final segmentation fineness.

[0011] Weak multi-scale information fusion ability and limited detail expression: Facing significant anatomical structure differences and imaging parameter differences among patients, if existing models have a large number of parameters or rely on specific modality features, it is easy to lead to a decrease in inference speed or poor adaptability, making it difficult to meet the actual clinical needs for fast and stable segmentation.

[0012] Therefore, there is a need for a medical image segmentation method that has both global modeling capabilities and computational efficiency, can adapt to the task of brain lesion recognition in high-resolution MRI scenarios, improve the segmentation accuracy, and ensure the efficiency and resource friendliness of the model in practical applications. Summary of the Invention

[0013] The object of the present invention is to solve the problems raised in the background technology. Specifically, for the automatic segmentation of brain MRI images, a novel network framework combining global self-attention and multi-scale feature fusion is proposed to obtain higher segmentation accuracy and good generalization performance in the relatively complex task of brain tumor region segmentation.

[0014] To achieve the object of the present invention, the present invention discloses a method for segmenting brain MRI images based on the RWKV model, including the following steps:

[0015] Step 1: Standardize and pad the input two-dimensional brain MRI image, and divide it into image patches of different granularities. Embed the image patch sequence into the feature space through linear projection and position encoding.

[0016] Step 2: Use two groups of parallel Vision-RWKV encoders to process the large-scale and small-scale image patch sequences respectively, so as to extract the image semantic features at different resolutions and achieve dual-scale feature encoding.

[0017] Step 3: Input the above dual-scale features into the feature fusion module VIF; the main branch generates image patch labels through pooling, and fuses with the supplementary branch features to achieve multi-scale context interaction.

[0018] Step 4: Transfer the fused encoded features as skip connections to each layer of the decoder, and combine with the RWKV module for hierarchical upsampling and feature decoding to gradually restore the image spatial information.

[0019] Step 5: Optimize the model parameters through an end-to-end training method, and use a joint loss function to improve the segmentation accuracy and boundary robustness, and output a segmentation mask with the same resolution as the input image.

[0020] Furthermore, the overall structure of the model is named Dual-RWKV-UNet, which adopts an "encoder-decoder" architecture. The model input is a two-dimensional brain MRI slice, and the output is a segmentation mask image with the same resolution as the input; the encoder part consists of two groups of parallel Vision-RWKV modules, which process the image patch sequences of different resolutions respectively to extract multi-scale features; the decoder part uses skip connections to obtain the encoder output, and combines with the Vision-RWKV module for upsampling and feature fusion to achieve fine segmentation.

[0021] The Vision-RWKV module adopts a block stacking form similar to ViT, with a bidirectional self-attention mechanism Bi-WKV of linear complexity, capable of achieving long-distance context modeling between different image patches; each block is composed of a spatial mixing module and a channel mixing module; through the two branch structures of spatial mixing and channel mixing, it has the ability to model long-distance dependencies and extract global context information.

[0022] Further, in step 1, the Vision-RWKV model preprocesses the image with the same padding and image patch embedding module; in this process, the model evenly divides the padded feature map with a resolution of H×W into HW / p 2 square image patches, where p is the side length of each image patch; subsequently, the sequence formed by these image patches will be linearly projected through a fully connected layer and added with positional encoding to generate an image patch sequence X with a shape of T×C, where T = HW / p 2 represents the total number of image patches, and C is the number of dimensions of the feature map after linear projection; subsequently, these image patch sequences will be input into the Vision-RWKV structure of L layers for processing.

[0023] Further, in each layer, the image patch sequence X first enters the spatial mixing module; in the spatial mixing module, a four-way image shift mechanism Q-Shift is introduced to slightly shift the image patches in the up, down, left, and right directions respectively and perform interpolation operations, thereby enhancing the perception ability between adjacent image patches and expanding the receptive field; specifically, the process of the four-way shift mechanism Q-Shift shifting the image patches and interpolating with adjacent image patches is represented by the mathematical form shown in Equation (1):

[0024]

[0025] In the above formula, the subscript *∈{R,K,V} represents three interpolation operations on X and These operations are distinguished by the learnable parameter μ (*) and are subsequently used for the calculation of the R, K, and V matrices; then, through a linear mapping operation, query, key, and value matrices are generated, and a bidirectional weighting mechanism Bi-WKV is introduced to calculate the attention response between image patches, realizing context modeling between any regions in the image. The specific calculation is shown in Equation (2):

[0026]

[0027] where T represents the total number of image patches, which is HW / p 2 , w and v are two learnable vectors with a dimension of C, used to represent spatial attenuation and calculate the weight of the current position respectively; and the vector k i and vi is the vector corresponding to the \(i\)-th image patch in the input feature matrices \(K\) and \(V\);

[0028] Immediately following is the channel mixing module, which is used for feature aggregation in the channel dimension. Similar to the spatial mixing module, it also adopts the four-way offset mechanism Q-Shift offset method for local structure enhancement, and then completes channel fusion through linear transformation and activation functions; both modules adopt residual connection and normalization structures to improve training stability.

[0029] Furthermore, in step 2, a dual-scale encoding structure is adopted, and two groups of Vision-RWKV modules are used to process the image patch sequences of different sizes in parallel; the main branch uses smaller image patches to extract fine-grained structure information, and the supplementary branch uses larger image patches to capture the overall shape and contour features.

[0030] Furthermore, in step 3, to achieve effective information interaction between multi-scale features, a feature map fusion module VIF is designed; this module unifies and fuses the feature maps output by the main branch and the supplementary branch at each encoding stage by introducing a patch-level guiding mechanism; in the VIF module, the main branch feature map undergoes pooling and flattening operations to generate a patch token sequence, which is concatenated with the other branch and then input into the Vision-RWKV structure to complete the context interaction between multi-scales; for a stage \(i\) (\(i = 1, 2, 3, 4\)) in the encoder, there will be outputs from both branches, and these two outputs are presented in the form of a feature map sequence, which are respectively denoted as: (main branch) and (supplementary branch); subsequently, the transformed output of \(G\) is obtained through formula (3): i transformed output:

[0031]

[0032] where AvgPool is a one-dimensional average pooling layer, followed by a flattening operation; the feature image token block represents the global abstract information of \(G\) i and is used to interact and supplement information with \(F\) i at the fine-grained level of image patches; subsequently, \(F\) i and are concatenated to form a sequence containing \(1 + h\times w\) image patches and are input into a single-layer Vision-RWKV layer to calculate the global self-attention:

[0033]

[0034] where is the final output of the small-scale branch in the VIF module; this method enables each image patch in F i to establish a connection with the entire G i so that the small-scale branch can also obtain the coarse-grained information in the large-scale branch.

[0035] Furthermore, in step 4, each stage of the decoder contains three modules: upsampling, skip connection, and Vision-RWKV. In each stage of the decoder, the input feature map is upsampled by a factor of 2. The upsampled feature map is concatenated with the feature map from the encoder side in the skip connection. Subsequently, the concatenated result is fed into the Vision-RWKV module for self-attention calculation. In such a decoding stage, long-range dependencies are established and global context information is captured using the features from the encoder and upsampling, thereby achieving better decoding performance. In each stage of the decoder, the resolution of the feature map is doubled, while the output dimension is halved. Finally, two convolutional neural network modules are used in the model to directly downsample the input image to obtain shallow feature maps with resolutions of (H / 2)×(W / 2) and H×W. Each convolutional module consists of a 3×3 convolutional layer, a group normalization layer, and a ReLU activation layer in sequence. The outputs of all these modules above are used to obtain the final segmentation mask result map with a resolution of H×W×1.

[0036] Furthermore, in step 5, during the training phase, the Dual-RWKV-UNet model proposed in the present invention is trained in an end-to-end manner according to the objective function. The loss function consists of weighted intersection over union loss and binary cross-entropy loss where focuses on the overlap between the prediction result and the ground truth, can measure the segmentation effect as a whole, and is conducive to the model quickly learning basic features in the early stage; while focuses on the difference between the predicted probability and the ground truth label, pays attention to the detailed differences, and is conducive to guiding the model to optimize the segmentation accuracy. The overall loss function of the model can be expressed by formula (5):

[0037]

[0038] In the formula, G is the ground truth label in the training set, and S1 is the final output image at the decoding end. During training, the values of the three hyperparameters α, β, and γ are 0.6, 0.2, and 0.2 respectively.

[0039] Compared with the prior art, the significant progress of the present invention lies in: 1) improving the segmentation accuracy, showing higher accuracy and stability in aspects such as boundary localization and lesion contour recognition; 2) optimizing the inference efficiency and resource consumption. By using RWKV to replace the traditional self-attention structure, the computational complexity is significantly reduced, making it more suitable for the real-time processing requirements of high-resolution images in a clinical environment; 3) having strong cross-disease transfer ability and data distribution adaptability, with good clinical promotion prospects.

[0040] To more clearly illustrate the functional characteristics and structural parameters of the present invention, the following further explains in conjunction with the accompanying drawings and specific embodiments. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] The accompanying drawings described herein are used to provide a further understanding of the present invention, and constitute a part of this application. The illustrative embodiments of the present invention and their descriptions are used to explain the present invention and do not constitute an improper limitation to the present invention. In the drawings:

[0042] Figure 1 is the detailed architecture diagram of the Dual-RWKV-UNet model;

[0043] Figure 2 is the overall structure diagram of the Vision-RWKV model;

[0044] Figure 3 is the structure diagram of the feature fusion module;

[0045] Figure 4 is the segmentation effect diagram of the Dual-RWKV-UNet model;

[0046] Figure 5 is the visual intuitive comparison diagram of various image segmentation algorithms;

[0047] Figure 6 is the box plot of the image segmentation results;

[0048] Figure 7 is the box plot of the image segmentation data for generalization performance analysis. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0049] The following will clearly and completely describe the technical solutions in the embodiments of the present invention in conjunction with the accompanying drawings in the embodiments of the present invention. Obviously, the described embodiments are only a part of the embodiments of the present invention, rather than all the embodiments; based on the embodiments of the present invention, all other embodiments obtained by those of ordinary skill in the art without creative efforts shall fall within the protection scope of the present invention.

[0050] The present invention aims to propose an efficient medical image segmentation method combining the RWKV linear self-attention mechanism and the U-Net structure. First, an RWKV linear self-attention module is introduced into the U-Net network framework to establish associations between remote pixels with relatively low computational overhead, thereby improving the recognition accuracy of the brain tumor region, effectively controlling the inference time, and enhancing the practicality of the model. Second, this method combines a multi-scale encoding and feature fusion mechanism to jointly extract local and global information, mine features from different resolution levels, and achieve adaptive fusion to improve the fine-grained segmentation effect. Finally, to enhance the generalization ability and lightweight deployment performance of the model, the present invention further optimizes the network structure, reduces the overall computational resource requirements, enables it to better adapt to cross-patient MRI data in a complex clinical environment, and has good practical value and promotion potential. The specific implementation of this method is as follows:

[0051] Step 1: Standardize and pad the input two-dimensional brain MRI image, and divide it into image patches of different granularities. Embed the image patch sequence into the feature space through linear projection and positional encoding;

[0052] Step 2: Use two groups of parallel Vision-RWKV encoders to process the large-scale and small-scale image patch sequences respectively, so as to extract image semantic features at different resolutions and achieve dual-scale feature encoding;

[0053] Step 3: Input the above dual-scale features into the feature fusion module VIF (Vision-RWKV ImageFusion); the main branch generates image patch labels through pooling, and fuses with the supplementary branch features to achieve multi-scale context interaction;

[0054] Step 4: Take the fused encoded features as skip connections and pass them to each layer of the decoder. Combine the RWKV module for hierarchical upsampling and feature decoding to gradually restore the image spatial information;

[0055] Step 5: Optimize the model parameters through an end-to-end training method, and use a joint loss function (including weighted IoU loss and binary cross-entropy loss) to improve the segmentation accuracy and boundary robustness, and output a segmentation mask with the same resolution as the input image.

[0056] The present invention provides a method for segmenting brain MRI images based on the RWKV model, which uses a dual-scale Vision-RWKV structure as the feature extraction backbone network, and combines the U-Net-style skip connections and decoder design to achieve accurate segmentation of the brain tumor region in high-resolution medical images. This method takes into account the modeling capabilities of local details and global dependencies, significantly reduces the computational resource requirements while improving the segmentation accuracy, and has good clinical application potential.

[0057] As attachedFigure 1 As shown in the figure, the overall structure of the model proposed in the present invention is named Dual-RWKV-UNet, which adopts an "encoder-decoder" architecture. The input of the model is a two-dimensional brain MRI slice, and the output is a segmentation mask image with the same resolution as the input. Different from the traditional U-Net, the encoder part consists of two groups of parallel Vision-RWKV modules, which respectively process image patch sequences of different resolutions and extract multi-scale features; the decoder part uses skip connections to obtain the output of the encoder, and combines the Vision-RWKV module for upsampling and feature fusion to achieve fine segmentation.

[0058] The Vision-RWKV module is the core component of the present invention, and its specific structure is as Figure 2 shown. As can be seen from the figure, Vision-RWKV adopts a block stacking form similar to ViT, and has a linear complexity bi-directional self-attention mechanism (Bi-WKV), which can realize long-distance context modeling between different image patches. Each block consists of a spatial mixing module and a channel mixing module. Through the two branch structures of spatial mixing and channel mixing, it has the ability to model long-distance dependencies and extract global context information. Compared with the conventional Transformer, this module significantly reduces the computational complexity while maintaining the modeling ability, and is suitable for high-resolution medical image processing.

[0059] First, the Vision-RWKV model preprocesses the image with the same padding and image patch embedding module. In this process, the model evenly divides the padded feature map with a resolution of H×W into HW / p 2 square image patches, where p is the side length of each image patch. Subsequently, the sequence formed by these image patches will be linearly projected through a fully connected layer and added with positional encoding to generate an image patch sequence X with a shape of T×C, where T = HW / p 2 represents the total number of image patches, and C is the number of dimensions of the feature map after linear projection. Subsequently, these image patch sequences will be input into the L-layer Vision-RWKV structure for processing.

[0060] In each layer, the image patch sequence X first enters the spatial mixing module. As Figure 2 shown in the right figure, in the spatial mixing module, the present invention introduces an image four-way offset mechanism (Q-Shift), which slightly shifts the image patches in the up, down, left, and right directions respectively and performs interpolation operations, so as to enhance the perception ability between adjacent image patches and expand the receptive field. Specifically, Q-Shift is the process of shifting the image patches and interpolating with adjacent image patches, and its process can be expressed as the mathematical form shown in Equation (1):

[0061]

[0062] In the above formula, the subscript * ∈ {R, K, V}, representing three interpolation operations on X and These operations are distinguished by the learnable parameter μ (*) and are then used for the calculation of the R, K, and V matrices respectively. Subsequently, query, key, and value matrices are generated through linear mapping operations, and a bi-directional weighted mechanism (Bi-WKV) is introduced to calculate the attention response between image patches, realizing context modeling between any regions in the image. The specific calculation is shown in Equation (2):

[0063]

[0064] Here, T represents the total number of image patches, which is HW / p 2 , w and v are two learnable vectors with dimension C, which are used to represent spatial attenuation and calculate the weight at the current position respectively. And the vector k i and v i are the vectors corresponding to the i-th image patch in the input feature matrices K and V. Compared with the traditional unidirectional attention mechanism, this mechanism can capture long-range dependence information more comprehensively and has lower computational resource overhead.

[0065] Immediately following is the channel mixing module, which is mainly used for feature aggregation in the channel dimension. Similar to the spatial mixing structure, it also uses the Q-Shift offset method for local structure enhancement, and then completes channel fusion through linear transformation and activation functions. Both modules adopt residual connection and normalization structures to improve training stability. The overall Vision-RWKV structure not only has global modeling ability, but also significantly reduces the computational and video memory costs of the attention mechanism, and is suitable for high-resolution medical image processing scenarios. At the same time, residual connections are established before and after the spatial mixing module and the channel mixing module to ensure that the gradient does not disappear in the deep network during model training.

[0066] Considering that image patches at different granularities may contain complementary information, the present invention proposes a dual-scale encoding structure, which uses two groups of Vision-RWKV modules to process sequences of image patches with different sizes in parallel. The main branch uses smaller image patches to extract fine-grained structure information, and the supplementary branch uses larger image patches to capture overall shape and contour features. This structure can effectively improve the model's recognition ability for multi-scale lesions without increasing additional parameters.

[0067] To achieve effective information interaction between multi-scale features, the present invention designs a feature map fusion module (Vision-RWKV Image Fusion, abbreviated as VIF), as Figure 3As shown. By introducing an image block-level guiding mechanism, this module unifies and fuses the feature maps output by the main branch and the supplementary branch at each encoding stage. In the VIF module, the feature map of the main branch undergoes pooling and flattening operations to generate an image block token sequence, which is concatenated with the other branch and then input into the Vision-RWKV structure to complete the context interaction between different scales. For a stage i (i = 1, 2, 3, 4) in the encoder, there are outputs from two branches, both of which are presented in the form of a feature map sequence and can be denoted respectively as: (main branch) and (supplementary branch). Subsequently, the transformed output of G i can be obtained in the manner shown in formula (3):

[0068]

[0069] where AvgPool is a one-dimensional average pooling layer, followed by a flattening operation. The feature image token block represents the global abstract information of G i and is used to interact and supplement information with F i at the fine-grained level of image blocks. Subsequently, F i and are concatenated to form a sequence containing 1 + h × w image blocks and are input into a single-layer Vision-RWKV layer (the "Fusion" module in Figure 3 ) to calculate the global self-attention:

[0070]

[0071] where is the final output of the small-scale branch in the VIF module. This method enables each image block in F i to establish a connection with the entire G i , so that the small-scale branch can also obtain the coarse-grained information in the large-scale branch.

[0072] This fusion method effectively solves the problem that traditional concatenation or convolution operations are difficult to model long-range dependencies in multi-scale fusion, while maintaining the overall structural consistency and light weight of the model. The fused features will be used as skip connection information and transmitted to the decoder to further improve the accuracy and structural consistency of the segmentation boundary.

[0073] Each stage of the decoder of the present invention includes three modules: upsampling, skip connection, and Vision-RWKV. In each stage of the decoder, the input features are upsampled by a factor of 2. The upsampled feature map is concatenated with the feature map from the encoder side in the skip connection. Subsequently, the concatenated result is fed into the Vision-RWKV module for self-attention calculation. In such a decoding stage, the features from the encoder and upsampling can be fully utilized to establish long-range dependencies and capture global context information, thereby achieving better decoding performance. In each stage of the decoder, the resolution of the feature map is increased by a factor of 2, while the output dimension is reduced by a factor of 2. Finally, two convolutional neural network modules are used in the model to directly downsample the input image to obtain shallow feature maps with resolutions of (H / 2)×(W / 2) and H×W. Each convolutional module consists of a 3×3 convolutional layer, a group normalization layer, and a ReLU activation layer in sequence. The outputs of all these above-mentioned modules are used to obtain the final segmentation mask result map with a resolution of H×W×1.

[0074] In the training stage, the Dual-RWKV-UNet (DRWKV) model proposed by the present invention is trained in an end-to-end manner according to the objective function. The loss function consists of weighted intersection over union loss and binary cross-entropy loss where focuses on the overlap between the prediction result and the ground truth, which can measure the segmentation effect as a whole and is conducive to the model quickly learning basic features in the early stage; while focuses on the difference between the predicted probability and the ground truth label, pays attention to the detailed differences, and is conducive to guiding the model to optimize the segmentation accuracy. The overall loss function of the model can be expressed by formula (5):

[0075]

[0076] In the formula, G is the ground truth label in the training set, and S1 is the final output image at the decoding end. During training, the values of the three hyperparameters α, β, and γ are 0.6, 0.2, and 0.2 respectively.

[0077] The key to the present invention lies in the Vision-RWKV structure, the dual encoder and the VIF fusion module, as well as the combination of the U-Net style decoder and RWKV. First, the present invention extends the RWKV model from the field of natural language processing to the image segmentation task, and utilizes the linear complexity Bi-WKV self-attention to effectively reduce the computational storage requirements. Secondly, a dual encoder and a VIF fusion module are adopted. During the multi-scale processing, through the interactive feature fusion between stages, the model can capture local details and global information more evenly, improving the segmentation accuracy. Finally, at the decoder end of the U-Net structure, RWKV is combined for global self-attention calculation, making up for the deficiency in the ability to model long-distance associations during the traditional convolutional decoding process, thereby obtaining a more accurate and delicate segmentation boundary. These key technologies jointly improve the performance of medical image segmentation, giving the present invention significant advantages in terms of accuracy and computational efficiency.

[0078] In summary, by introducing the Vision-RWKV module, the dual-scale encoding structure and the multi-level feature fusion mechanism, the present invention effectively controls the model computational complexity while improving the segmentation accuracy and the model generalization ability. It is suitable for deployment and operation in resource-constrained medical environments and has good practical value and promotion prospects.

[0079] Embodiment

[0080] The Dual-RWKV-UNet model proposed by the present invention, aiming at the deficiencies of traditional segmentation methods in terms of accuracy, efficiency and generalization ability, has achieved the following remarkable improvements in the brain MRI image segmentation task by introducing a lightweight RWKV structure and a multi-scale feature fusion mechanism:

[0081] (1) Improvement in segmentation accuracy

[0082] In the BraTS2019 brain tumor segmentation test set, the average Dice coefficient of the proposed model reaches 0.8114, the average IoU reaches 0.7104, and the 95% Hausdorff distance is reduced to 5.7634, showing the best performance among all the compared models (see Figure 6 and Table 1). Compared with the classic UNet series and the mainstream Transformer models (TransUNet, Swin-UNet), the present invention shows higher accuracy and stability in aspects such as boundary localization and lesion contour recognition.

[0083] Table 1 Test results table of image segmentation accuracy data

[0084] Method Name Dice Score IoU HD95% UNet 0.7102 0.5962 6.1699 UNet++ 0.7673 0.6561 5.9356 Segformer 0.7681 0.6568 6.0253 TransUNet 0.7888 0.6805 5.8597 SwinUNet 0.7043 0.5902 6.5957 DRWKV (Method in this paper) 0.8114 0.7104 5.7634

[0085] To visually demonstrate the performance of the proposed DRWKV model in brain tumor segmentation tasks, it was visually compared with various existing typical segmentation models. As Figure 4 shown, this figure respectively shows the original MRI cross-sectional image, the label segmentation image, the prediction result of the DRWKV model, and the comparison of its boundary with the label result. It can be observed that the DRWKV model has accurate localization in the tumor area, clear segmentation trajectories, and the prediction result is highly consistent with the actual label in terms of area and edge position.

[0086] As Figure 5 shown, it further shows the visual effect comparison of the DRWKV model with five comparison models (UNet, UNet++, TransUNet, SwinUNet, Segformer) in the segmentation task. All models can complete the basic segmentation function of the brain tumor area, while the model proposed in this invention performs better in boundary detail processing and area coverage, verifying its advancement and effectiveness in brain MRI image segmentation tasks.

[0087] (2) Optimization of Inference Efficiency and Resource Consumption

[0088] This method uses RWKV to replace the traditional self-attention structure, significantly reducing the computational complexity. Under the same hardware conditions, the inference speed is increased by 31.6% compared with Swin-UNet and 27.2% compared with TransUNet. At the same time, the video memory occupancy is reduced by nearly 35%, which is more suitable for the real-time processing requirements of high-resolution images in the clinical environment (see Table 2).

[0089] Table 2 Image Segmentation Data Results Table for Generalization Performance Analysis

[0090] Method Name Dice Score IoU HD95% UNet 0.7682 0.6558 5.9829 UNet++ 0.8033 0.7109 5.6646 Segformer 0.8003 0.6957 5.9256 TransUNet 0.8086 0.7022 5.9441 SwinUNet 0.7198 0.5968 6.7879 DRWKV (Method in this paper) 0.8417 0.7409 5.6176

[0091] (3) Stronger Cross-Disease Generalization Ability

[0092] When this model is trained only with high-grade glioma (HGG) samples and directly applied to the low-grade glioma (LGG) test set, it still maintains good segmentation performance (Dice 0.8417, IoU 0.7409, HD95 5.6176), which is better than all comparison methods (see Figure 7 , Table 2). This shows that this method has strong cross-disease migration ability and data distribution adaptability, and has good clinical promotion prospects.

[0093] In summary, the Dual-RWKV-UNet model proposed in this invention is superior to existing mainstream methods in various evaluation dimensions. It can obtain high-quality and clear-boundary segmentation results while maintaining a low computational burden, and is suitable for the automatic segmentation task of large-scale brain MRI images in actual clinical practice.

[0094] It should be noted that, in this document, relational terms such as first and second are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any actual relationship or order between these entities or operations. Moreover, the terms "comprising", "including" or any other variant thereof are intended to cover non-exclusive inclusion, such that a process, method, article or apparatus comprising a series of elements not only includes those elements but also includes other elements not expressly listed, or further includes elements inherent to such process, method, article or apparatus.

[0095] Although the embodiments of the present invention have been shown and described, it will be understood by those of ordinary skill in the art that various changes, modifications, substitutions and variations can be made to these embodiments without departing from the principles and spirit of the present invention, and the scope of the present invention is defined by the appended claims and their equivalents.

Claims

1. A brain MRI image segmentation method based on the RWKV model, characterized in that, It includes the following steps: Step 1: Standardize and fill the input two-dimensional brain MRI image, divide it into image patches of different granularities, and embed the sequence of image patches into the feature space through linear projection and position encoding; Step 2: Use two groups of parallel Vision-RWKV encoders to process the large-scale and small-scale image patch sequences respectively, so as to extract the image semantic features at different resolutions and achieve dual-scale feature encoding; Step 3: Input the above dual-scale features into the feature fusion module VIF; the main branch generates image patch labels through pooling, and fuses with the supplementary branch features to achieve multi-scale context interaction; Step 4: Transfer the fused encoded features as skip connections to each layer of the decoder, and combine with the RWKV module for hierarchical upsampling and feature decoding to gradually restore the image spatial information; Step 5: Optimize the model parameters through an end-to-end training method, adopt a joint loss function to improve the segmentation accuracy and boundary robustness, and output a segmentation mask with the same resolution as the input image; 2. A brain MRI image segmentation method based on the RWKV model according to claim 1, characterized in that, The overall structure of the model is named Dual-RWKV-UNet, adopting an "encoder-decoder" architecture. The model input is a two-dimensional brain MRI slice, and the output is a segmentation mask image with the same resolution as the input; the encoder part consists of two groups of parallel Vision-RWKV modules, which process the image patch sequences of different resolutions respectively and extract multi-scale features; The decoder part uses skip connections to obtain the encoder output, and combines with the Vision-RWKV module for upsampling and feature fusion to achieve fine segmentation; The Vision-RWKV module adopts a block stacking form similar to ViT, with a bidirectional self-attention mechanism Bi-WKV with linear complexity, which can achieve long-distance context modeling between different image patches; each block Block consists of a spatial mixing module and a channel mixing module; through the two-branch structures of spatial mixing and channel mixing, it has the ability to model long-distance dependencies and extract global context information.

3. A method for segmenting brain MRI images based on the RWKV model according to claim 2, characterized in that, In step 1, the Vision-RWKV model preprocesses the image with the same padding and image patch embedding module; during this process, the model evenly divides the padded feature map with a resolution of H×W into 2 square image patches, where p is the side length of each image patch; subsequently, the sequence formed by these image patches will be linearly projected through a fully connected layer and added with positional encoding to generate an image patch sequence X with a shape of T×C, where T = HW / p 2 represents the total number of image patches, and C is the number of dimensions of the feature map after linear projection; subsequently, these image patch sequences will be input into the Vision-RWKV structure of L layers for processing.

4. A brain MRI image segmentation method based on the RWKV model according to claim 3, characterized in that, In each layer, the image patch sequence X first enters the spatial mixing module; in the spatial mixing module, an image four-way offset mechanism Q-Shift is introduced, and the image patches are slightly shifted in the up, down, left, and right directions respectively and interpolation operations are performed to enhance the perception ability between adjacent image patches and expand the receptive field; specifically, the process of the four-way offset mechanism Q-Shift shifting the image patches and interpolating with adjacent image patches is represented by the mathematical form shown in Equation (1): X[h,w-1,C / 2:3C / 4],X[h,w+1,3C / 4:C])(1) In the above formula, the subscript *∈{R, K, V}, representing three interpolation operations on X and , which are distinguished by the learnable parameter μ (*) and are subsequently used for the calculation of the R, K, and V matrices respectively; subsequently, query, key, and value matrices are generated through a linear mapping operation, and a bidirectional weighted mechanism Bi-WKV is introduced to calculate the attention response between image patches, realizing context modeling between any regions in the image. The specific calculation is shown in Equation (2): where T represents the total number of image patches, which is HW / p 2 , w and v are two learnable vectors of dimension C, which are used to represent spatial attenuation and calculate the weight at the current position respectively; while the vector k i and v i are the vectors corresponding to the i-th image patch in the input feature matrices K and V; Followed by the channel mixing module, which is used for feature aggregation in the channel dimension. Similar to the spatial mixing module structure, it also uses the four-way offset mechanism Q-Shift offset method for local structure enhancement, and then completes channel fusion through linear transformation and activation function; both modules adopt residual connection and normalization structures to improve the training stability.

5. A method for segmenting brain MRI images based on the RWKV model according to claim 2, characterized in that, In step 2, a dual-scale encoding structure is adopted, and two groups of Vision-RWKV modules are used to process sequences of image patches of different sizes in parallel; the main branch uses smaller image patches to extract fine-grained structural information, and the supplementary branch uses larger image patches to capture overall shape and contour features.

6. A brain MRI image segmentation method based on the RWKV model according to claim 2, characterized in that, In step 3, to achieve effective information interaction between multi-scale features, a feature map fusion module VIF is designed; this module unifies the feature maps output by the main branch and the supplementary branch at each encoding stage by introducing a guiding mechanism at the image patch level; in the VIF module, the feature map of the main branch generates an image patch token sequence after pooling and flattening, and is concatenated with the other branch and then input into the Vision-RWKV structure to complete the context interaction between multi-scales; for a stage i (i = 1, 2, 3, 4) in the encoder, there will be outputs from two branches, and both outputs are presented in the form of a feature map sequence, denoted as: and Subsequently, the transformed output of G i is obtained through formula (3): Among them AvgPool is a one-dimensional average pooling layer, followed by a flattening operation; the feature image marking block represents the global abstract information of G i for interacting and supplementing information with F at the fine-grained level of image patches; subsequently, F i is used for interaction and information supplementation; subsequently, F i and are concatenated to form a sequence containing 1 + h × w image patches and are input into a single-layer Vision-RWKV layer to calculate global self-attention: wherein is the final output of the small-scale branch in the VIF module; this method enables each image patch in F i to establish a connection with the entire G i so that the small-scale branch can also obtain the coarse-grained information in the large-scale branch.

7. A brain MRI image segmentation method based on the RWKV model according to claim 2, characterized in that, In step 4, each stage of the decoder contains three modules: upsampling, skip connection, and Vision-RWKV; in each stage of the decoder, the input features are upsampled by a factor of 2 once. The upsampled feature map is concatenated with the feature map from the encoder side in the skip connection, and then the concatenated result is fed into the Vision-RWKV module for self-attention calculation; in such a decoding stage, long-range dependencies are established and global context information is captured by using the features from the encoder and upsampling, thus achieving better decoding performance; in each stage of the decoder, the resolution of the feature map is increased by a factor of 2, while the output dimension is reduced by a factor of 2; finally, two convolutional neural network modules are used in the model to directly downsample the input image to obtain shallow feature maps with resolutions of (H / 2)×(W / 2) and H×W. Each convolutional module consists of a 3×3 convolutional layer, a group normalization layer, and a ReLU activation layer in sequence; the outputs of all these above modules are used to obtain the final segmentation mask result map with a resolution of H×W×1.

8. A method for segmenting brain MRI images based on the RWKV model according to claim 2, characterized in that, In step 5, in the training stage, the Dual-RWKV-UNet model proposed by the present invention is trained in an end-to-end manner according to the objective function; The loss function consists of the weighted intersection over union loss and the binary cross-entropy loss where focuses on the overlap between the predicted result and the ground truth, which can measure the segmentation effect as a whole and is conducive to the model quickly learning basic features in the early stage; while focuses on the difference between the predicted probability and the ground truth label, pays attention to the detailed differences, and is conducive to guiding the model to optimize the segmentation accuracy; the overall loss function of the model can be expressed by formula (5): where G is the ground truth label in the training set, and S1 is the final output image at the decoding end; during training, the values of the three hyperparameters α, β, and γ are 0.6, 0.2, and 0.2 respectively.

Images