A method for MRI image segmentation based on high-resolution network and boundary enhancement
By using a high-resolution network based on Swin Transformer and a boundary enhancement method, the problems of reduced resolution and difficulty in distinguishing boundaries in brain tumor MRI image segmentation are solved, achieving high-precision brain tumor segmentation results.
Patent Information
- Application Number
- CN202310370910.8
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2023-04-10
- Publication Date
- 2026-02-13
- Estimated Expiration
- 2043-04-10
AI Technical Summary
Existing methods for segmenting brain tumor MRI images face problems such as reduced resolution and difficulty in distinguishing boundary regions, especially deep learning-based methods which struggle to accurately segment tumors when the boundaries are blurred.
We employ a high-resolution network based on the Swin Transformer and a boundary enhancement method. Through adaptive weights and knowledge distillation, we combine a multi-resolution feature fusion module with patch merging and patch expanding modules to maintain the high resolution of the image and enhance the segmentation ability of the boundary region.
It achieves high resolution in brain tumor image segmentation tasks, improves boundary discrimination ability and model accuracy, and simplifies the model training process.
Smart Images

Figure CN116452619B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The application belongs to the technical field of image segmentation, and relates to an MRI image segmentation method based on a high-resolution network and boundary enhancement. BACKGROUND
[0002] Deep learning is a series of algorithms in the field of machine learning that attempt to use multiple nonlinear transformations to abstract data in multiple layers, not only learning the nonlinear mapping between input and output, but also learning the hidden structure of the input data vector to intelligently identify or predict new samples. With the development of deep learning, deep learning methods have shown their powerful ability in image processing tasks. The potential of deep learning methods makes them the main choice for image segmentation, especially medical image segmentation. In recent years, image segmentation based on deep learning technology has received extensive attention, and the breakthrough of deep learning in medical image segmentation is crucial for the development of the medical field.
[0003] Most mainstream segmentation methods are based on the FCN architecture, in which the encoder gradually reduces the spatial resolution to obtain a larger receptive field and learn more high-level semantic information. However, for dense tasks such as segmentation, the reduction of resolution will inevitably cause information loss, which will affect the segmentation effect. This is reflected in the brain tumor image segmentation task, where the intra-class of glioma has variable shape, position, and regularity. For some small volume enhanced tumors, they will be missed as the model resolution decreases, thus reducing the segmentation performance. On the other hand, the new network architecture Transformer has shown strong competitiveness in the segmentation field due to its natural long-distance characteristics. This model can capture the global view without reducing the resolution, so it is expected to solve the information loss problem caused by resolution reduction. Some scholars introduce Transformer as the backbone network, but this network has problems such as large number of parameters and difficulty in training, so the mainstream approach has two methods, 1) Patch, the Patch size of the Vit model is 16x16, which is undoubtedly also a reduction in resolution. 2) Local Self-Attention, Swin Transformer model uses a hierarchical scheme of local self-attention plus patch merging, which is also a method of reducing resolution layer by layer. How to propose a new scheme to preserve the original resolution as much as possible to complete the segmentation task is the key to solving the information loss problem.
[0004] The existing brain tumor MRI image segmentation method is mainly based on a deep learning model, and the brain tumor image segmentation task based on deep learning faces two problems: 1) the U-Net network is usually used as the backbone network for medical image segmentation, and the model uses a convolutional neural network (CNN) as an encoder. Due to the limited receptive field of the Conv operator of the CNN, in order to expand the attention area of the network, multiple layers (convolution-pooling) need to be stacked, which will cause the resolution to be reduced, thereby causing information loss; 2) because the tumor is usually surrounded by healthy brain tissue such as cerebrospinal fluid, gray matter, white matter, etc., the tumor boundary is blurred, making it difficult for methods based on deep learning technology to distinguish the boundary region. SUMMARY
[0005] To achieve the above purpose, the present application provides an MRI image segmentation method based on a high-resolution network and boundary enhancement, which solves the problems of resolution reduction and difficulty in distinguishing the boundary region during the segmentation of the existing brain tumor MRI image.
[0006] The technical scheme adopted by the present application is an MRI image segmentation method based on a high-resolution network and boundary enhancement, comprising the following steps:
[0007] Step S1, for the MRI image to be segmented, a pretreatment operation is performed;
[0008] Step S2, the pretreated image data is input into a first high-resolution segmentation network model Model1 and a second high-resolution segmentation network model Model2, both of which are high-resolution segmentation network models HRSwinNet based on Swin Transformer, and Model2 adopts a boundary enhancement method;
[0009] The boundary enhancement method is an adaptive weight-based boundary enhancement learning method and / or a knowledge distillation-based boundary enhancement learning method;
[0010] Step S3, the final output result Output1 of Model1 and the final output result Output2 of Model2 are averaged to obtain the final segmentation result.
[0011] Further, the pretreatment in step S1 includes three methods of image cropping, resampling and standardization.
[0012] Further, the HRSwinNet model comprises coding stages Stage, each Stage is sequentially connected, and a multi-resolution feature fusion module MRFF is inserted between each Stage; the nth Stage comprises n parallel Swin Transformer modules and patch merging modules except for the last Stage, and the last Stage comprises n parallel Swin Transformer modules and n-1 patch merging modules; the output of the last multi-resolution feature fusion module MRFF is sequentially connected with a residual module, a patch expanding module and a Conv module; wherein the patch merging module and the patch expanding module are respectively used for down-sampling and up-sampling of the feature map.
[0013] Further, the Swin Transformer module comprises two cascaded layers, the first layer is a window-based multi-head self-attention module W-MSA, and the second layer is a shift window-based multi-head self-attention module SW-MSA, one W-MSA module and one SW-MSA module are sequentially connected;
[0014] The calculation process of the continuous Swin Transformer module is as follows:
[0015]
[0016]
[0017]
[0018]
[0019] In the formula, and z l respectively represent the output features of the W-MSA module and the MLP module at the l layer, l-1 represents the output of the previous layer, and l+1 represents the output of the next layer; LN is layer normalization, and MLP is a multi-layer perceptron.
[0020] Further, the loss function of the HRSwinNet model is:
[0021]
[0022] wherein I represents the number of image voxels; J is the category number; Y i,j and G i,j respectively represent the output probability of class j at voxel i, and a one-hot encoded true value.
[0023] Further, the boundary enhancement learning method based on adaptive weights in the step S2 is specifically:
[0024] First, the high-resolution segmentation network model HRSwinNet is regularly trained, and the trained model is denoted as Model1; the training data is input into Model1 and the to-be-trained model Model2, the output result of Model1 is logits, and the probability value prob is obtained through Softmax(logits); the cross entropy C=CrossEntropy(output,target) of prob and the true value target is calculated, and then the value of C is mapped between (0, δ) through the Sigmoid function, and the value obtained at this time is the weight w value to be weighted to Model2; Model2 is weighted by using the weight w, and an enhanced model Model2 is obtained.
[0025] When calculating the loss in the Model2 training process, the inner product is calculated with the weight w, and the final result is the loss after weighting.
[0026] Further, the boundary enhancement learning method based on knowledge distillation in the step S2 is specifically:
[0027] First, the high-resolution segmentation network model HRSwinNet is regularly trained, and the trained model is denoted as Model1; the training data is input into Model1 and the to-be-trained model Model2, the output result of Model1 is logits1, and then the Softmax operation is performed on logits1 to obtain a four-dimensional vector of a pixel point, and the length of each direction of the vector represents the probability that the pixel point is the corresponding class; the variance variance of the values of each direction of the four-dimensional vector is calculated, and the greater the variance, the greater the probability of being located in the boundary area; the hard label value of the point with the variance value variance less than the hyperparameter v is replaced with a soft label, and the loss of Model2 and the new soft label target2 is calculated, so that a stronger segmentation model Model2 is obtained.
[0028] The beneficial effects of the present application are:
[0029] 1) The high-resolution composite network based on the Swin Transformer proposed in the present application can achieve true high resolution, thereby improving the performance of brain tumor image segmentation tasks.
[0030] 2) The segmentation method has strong boundary distinguishing ability, simple model training, and high model precision. BRIEF DESCRIPTION OF DRAWINGS
[0031] In order to more clearly illustrate the technical solutions in the embodiments of the present application or the prior art, the drawings needed to be used in the embodiments or prior art description will be briefly introduced. Obviously, the drawings in the following description are only some embodiments of the present application, and other drawings can be obtained by those skilled in the art without creative labor.
[0032] Figure 1 is the HRSwinNet model architecture diagram of the embodiment of the present application.
[0033] Figure 2 is the encoder architecture diagram of the embodiment of the present application.
[0034] Figure 3 is the Swin Transformer module calculation process diagram of the embodiment of the present application.
[0035] Figure 4 is the decoder architecture diagram of the embodiment of the present application.
[0036] Figure 5 is the adaptive weight-based boundary enhancement learning method flowchart of the embodiment of the present application.
[0037] Figure 6 is the relationship diagram between model capturing "knowledge" and total amount of model parameters of the embodiment of the present application.
[0038] Figure 7 is the knowledge diagram contained in the soft label of the embodiment of the present application.
[0039] Figure 8 is the knowledge distillation-based boundary enhancement learning method flowchart of the embodiment of the present application. DETAILED DESCRIPTION
[0040] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only some embodiments of the present application, not all embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor are within the scope of protection of the present application.
[0041] The embodiment of the present application provides a high-resolution brain tumor MRI (Magnetic Resonance Imaging) image segmentation method based on high-resolution network and boundary enhancement, comprising the following steps:
[0042] Step 1,
[0043] For the brain tumor MRI image that needs to be segmented, first, a preprocessing operation is performed on the brain tumor MRI image, the preprocessing operation can remove useless information, ensure data consistency and unify data distribution, and the like, and the three preprocessing schemes of image cropping, resampling and standardization are adopted in the embodiment of the application.
[0044] Image cropping is to crop a three-dimensional medical image to its non-zero region, and the specific method is to find a smallest three-dimensional bounding box in the image, the values outside the bounding box region are 0, and the image is cropped using the bounding box. Compared with before cropping, the cropped image has no effect on the final segmentation result, but can reduce the image size, avoid useless calculation and improve the calculation efficiency. Image cropping is more effective for some data sets such as brain data sets, because there are relatively more black backgrounds in the peripheral brain MRI data set.
[0045] The purpose of resampling is to solve the problem that the actual space size represented by a single voxel in different images in some three-dimensional medical image data sets is inconsistent. Because the convolutional neural network only operates in the voxel space, the size information in the actual physical space is ignored. In order to avoid this difference, the size of different image data in the voxel space needs to be adjusted to ensure that the actual physical space represented by each voxel in different image data is consistent.
[0046] The purpose of standardization is to make the gray value of each image in the training set have the same distribution. For brain MRI data, the embodiment adopts z-scoring (i.e. subtracting the mean value and dividing by the standard deviation) for standardization, and the mean value and variance are calculated using the gray information of a single image during standardization.
[0047] Step 2,
[0048] The MRI data is input into the first high-resolution segmentation network model Model1 and the second high-resolution segmentation network model Model2, both of which are the high-resolution segmentation network model HRSwinNet based on Swin Transformer proposed in the application. Model1 is trained in a conventional training manner, and Model2 adopts the two boundary enhancement methods proposed in the application, and Model2 enhances the segmentation capability of the boundary region on the basis of Model1.
[0049] Step 3,
[0050] The final output results Output1 and Output2 of Model1 and Model2 are simply averaged to obtain Output, which is the final segmentation result.
[0051] The application provides a high-resolution segmentation network model HRSwinNet based on Swin Transformer. The model is composed of an encoder and a decoder. The encoder is responsible for encoding data and maintaining high-resolution representation during the encoding process. The decoder is responsible for decoding the encoded features into segmentation results. The specific details of the model are shown below.
[0052] As shown in Figure 1 HRSwinNet contains multiple encoding stages, called Stages. Each Stage is sequentially connected, and a multi-resolution feature fusion module is inserted between each Stage. The high-resolution features of the image remain unchanged in all Stages. The nth Stage contains n parallel Swin Transformer modules and patch merging modules. These Swin Transformer modules will produce n parallel low-resolution features with richer semantics. The input 3D medical image is fed into a 3D convolution module (Conv module), which divides the 3D medical image into small 3D blocks and represents each 3D block as a vector in the feature space. The vector of the 3D block goes through multiple Stages, and the output of the last Stage is fed into the last multi-resolution feature fusion module. After upsampling the low-resolution features, the outputs of the above feature fusion modules are connected. The final segmentation mask is generated by sequentially passing the connected features through a residual module, a patch expanding module, and a Conv module. The patch merging module and the patch expanding module are used for downsampling and upsampling the feature map, respectively. The definition of the patch merging module is the same as that of the Swin Transformer module, while the patch expanding module is the same as the VT-UNet. The detailed information of the Swin Transformer module, the HRSwinNet Stage, and the multi-resolution feature fusion module (MRFF) will be introduced in the following subsections.
[0053] Patch merging module
[0054] Swin Transformer is a new Transformer network model for image classification and object detection, which introduces a Patch merging module to improve the model's representation ability for multi-scale features. In Swin Transformer, the Patch merging module is responsible for fusing small image blocks (patches) after segmentation into larger feature maps. Unlike traditional patch merging, the Patch merging module in Swin Transformer adopts a multi-level fusion strategy in stages, called "Shifted Window Attention". Specifically, the Patch merging module in Swin Transformer is divided into several stages, each of which fuses small blocks from the previous stage into larger blocks, and uses the Shifted Window Attention module for feature fusion in each stage. Shifted Window Attention divides the input small blocks into multiple sub-blocks and performs self-attention calculation within each sub-block, then concatenates the resulting sub-block features into a larger feature block. This feature fusion method can effectively handle multi-scale information in images and has good computational and spatial efficiency.
[0055] In practice, the Patch merging module of Swin Transformer can help the model better handle multi-scale information and improve the performance of image classification and object detection tasks. At the same time, it also has high flexibility and can be trained and inferred at different resolutions, different image sizes, and different input block sizes.
[0056] Patch expanding module
[0057] In the Swin Transformer, the Patch Expanding module expands large chunks of the feature map into smaller patches. This process, the opposite of Patch Merging, aims to extract more detailed information from high-resolution regions of the feature map, improving the model's accuracy and robustness. Specifically, the Swin Transformer's Patch Expanding module consists of multiple staged blocks, each employing deconvolution and non-local attention mechanisms. Within each block, deconvolution expands the large chunks into smaller ones, and the non-local attention mechanism performs attention calculations on the expanded patches to obtain more refined feature representations. In practice, the Swin Transformer's Patch Expanding module helps the model better handle detailed information in high-resolution regions, improving performance in image classification and object detection tasks. It also offers high flexibility, allowing training and inference at different resolutions, image sizes, and input chunk sizes.
[0058] Swin Transformer Block
[0059] like Figure 3 The illustrated Swing Transformer module comprises two cascaded layers. The first layer is a windowing multi-head self-attention (W-MSA) module, which computes self-attention within the small window (W-MSA) of the input 3D block. The second layer is a shifted windowing multi-head self-attention (SW-MSA) module, which computes the output of the previous layer's self-attention within the shifted window (SW-MSA) of the input 3D block. A W-MSA block and an SW-MSA block are sequentially connected to obtain global spatial correlation in a more efficient manner.
[0060] Assuming the input 3D medical image is Where K is the number of channels in the input 3D data. The first step of SwinTransformer is to divide X into small 3D blocks, with the block size set to P×P×P, so we have Block. Order Represents a collection of 3D blocks. For example... Figure 2 As shown, 3D block set The data is fed into a Conv block, whose kernel size and stride are set to be equal to the size of the 3D block. The output of the Conv block is represented as... Since each pixel in Z represents a 3D patch embedding vector, the feature map Z is adjusted to a two-dimensional matrix Z ∈ R S×C where each row of the matrix Z represents a 3D patch embedding vector, and there are s rows in total. After encoding, the set of patch encoding vectors is fed into a Swin Transformer block, and the successive Swin Transformer Block computation process is as follows:
[0061]
[0062]
[0063]
[0064]
[0065] In the formula, and z l represent the output features of the (S)W-MSA module and the MLP module at layer l, respectively, l-1 represents the output of the previous layer, and l+1 represents the output of the next layer. LN (Layer Normalization) is a layer normalization, which is used to normalize the output of each layer to reduce the influence of internal covariate shift. MLP (Multi-Layer Perceptron) is a common feedforward neural network composed of multiple fully connected layers, which is used for nonlinear transformation and feature extraction, and its output can be used as the input of the subsequent layer. In the Transformer, MLP is usually used to perform nonlinear transformation on the attention output to improve the expression ability and generalization performance of the model.
[0066] HRSwinNet Stage
[0067] Each Stage generates multi-resolution feature maps, except for the last Stage, the nth Stage contains n Swin Transformer modules and n patch merging modules, while the last Stage contains n-1 patch merging blocks. In the nth Stage, n parallel Swin Transformer modules process n feature maps of different resolutions, respectively. Each Swin Transformer module is followed by a 3D patch merging module to downsample the output feature map of the Swin Transformer module.
[0068] Taking the third Stage (Stage-3) as an example, the input of this Stage is three feature maps with resolutions of The three Swin Transformer modules process the three feature maps respectively, and the output of the Swin Transformer module has the same resolution as the input feature map. The output of the Swin Transformer is fed into the patch merging layer, thereby generating three feature maps with resolutions of and The feature map with resolution is directly sent to the fourth stage (Stage-4), while the outputs of the other 3D block merging layers and the outputs of the Swin Transformer modules are sent to the multi-resolution feature fusion block.
[0069] Multi-resolution feature fusion module (MRFF)
[0070] As shown in Figure 4 , the multi-resolution feature fusion block is used to exchange information from feature maps of different resolutions and generate new feature maps with more accurate spatial and richer semantic information. Figure 4 The MRFF block after the fourth stage is shown, which fuses four features with different resolutions.
[0071] The fourth stage of HRSwinNet contains four Swin Transformer blocks, which process four feature maps with different resolutions respectively. There are four feature maps with resolutions of and generated by the four Swin Transformer modules. The feature maps with resolutions of and are called f4, f8, f16, and f32 respectively. Each of the four feature maps will be down-sampled or up-sampled to the other three resolutions. Taking the feature map f4 as an example, f4 is sequentially passed through three patch merging layers to generate three new feature maps with resolutions of The three newly generated feature maps are denoted as f 4_8 , f 4_16 and f 4_32 . It is easy to infer the generated feature maps of f8, f 16 , f 32 , and the generated feature maps are denoted as {f 8_4 ,f 8_16 ,f 8_32}, {f 16_4 ,f 16_8 ,f 16_32}, {f 32_4 ,f 32_8 ,f 32_16}. Then the feature maps with the same resolution are concatenated, for example, [f4, f 8_4 16_4 32_4 ], where the symbol [] represents feature concatenation. After concatenating the feature maps, the concatenated features will be sent to the residual module through the feedforward neural network to generate the output of the MRFF block, thus generating four feature maps with resolutions . The output of the MRFF block after the 2nd Stage and the 3rd Stage can be easily derived. As shown in Figure 2 , the output of the last MRFF block with resolution lower than is upsampled. After upsampling, all feature maps are concatenated and used to generate the output segmentation mask.
[0072] The upsampling operation is implemented by transpose convolution, and applying a 3*3*3 convolution to the concatenated features will reduce the dimension by half. Finally, the Softmax and Argmax operations will be performed on the fused features of all resolutions to obtain the final segmentation map.
[0073] For the medical image segmentation task, the loss function is an index for evaluating the error between the segmentation result and the true value in the training process, which will be used to continuously guide the neural network training. The soft Dice loss function is used in this embodiment, which is calculated in a voxel manner as follows:
[0074]
[0075] where I represents the number of image voxels; J is the class number; Y i,j and G i,j represent the output probability of class j at voxel i and a one-hot encoded true value, respectively.
[0076] Existing medical image segmentation algorithms either downsample high-resolution images or crop them into small blocks for separate processing. In either way, the loss of local details or global context information will lead to limited segmentation accuracy. The high-resolution composite network based on SwinTransformer proposed in the present application will achieve true high resolution, thus improving the performance in brain tumor image segmentation tasks.
[0077] 2 Boundary enhancement learning algorithm
[0078] 2.1 Boundary enhancement learning method based on adaptive weight
[0079] The segmentation task is actually a pixel-by-pixel classification task. Inspired by the AdaBoost algorithm, model 2 is introduced, which is the same as model 1 that has been trained. It is assumed that models 1 and 2 are both weak classifiers. Model 2 will increase the weight of the boundary that model 1 failed to correctly classify based on model 1. The increased weight value is determined by the cross-entropy of the predicted value of the pixel point and the true value target. The algorithm flowchart is shown in Figure 5 .
[0080] First, the high-resolution network HRSwinNet is fully trained. The model after training can achieve a high level. The model is denoted as Model 1. Then, the trained model Model 2 is introduced, and the training of Model 2 is started. The training data enters Model 1 and Model 2 at the same time. The output result of Model 1 is logits, and the probability value prob is obtained through Softmax(logits). The cross-entropy C = CrossEntropy(prob, target) of prob and target is calculated. Then, the value of C is mapped to (0, δ) through the Sigmoid function. The value obtained at this time is the weight value to be added.
[0081] The calculation method of the increased weight w is as follows:
[0082] w = 1 + δ · Sigmoid(CrossEntropy(Softmax(logits), target))
[0083] The inner product with the weight w is calculated when the loss of the model Model 2 is calculated. The final loss is the loss after weighting, so Model 2 will pay more attention to the learning of the error area (usually the boundary area). Finally, Model 1 and Model 2 are combined, that is, the output results of the two models are simply added and averaged to obtain a stronger final segmentation result. The final output result of the test set data x is:
[0084] output = Avg(Model1(x) + Model2(x))
[0085] The calculation method of the label weight of each image pixel is as follows:
[0086]
[0087]
[0088] where c is the number of classes, w is the increased weight, Y i,j is the original label value, is the label value after weighting.
[0089] 2.2 Boundary enhanced learning method based on knowledge distillation
[0090] For the boundary region, the present application introduces a soft label to replace the hard label, forcing the model to learn more hidden information. The flow chart of the boundary enhanced learning method based on knowledge distillation is as shown in Figure 8 .
[0091] First, the high-resolution network is trained completely, and the model after training can achieve a high level. The model is recorded as Model1, and then the to-be-trained model Model2 is introduced, and the training of Model2 is started. The training data enters Model1 and Model2 at the same time, and the output result of Model1 is logits1, and then the Softmax operation is performed on logits1. The result predicted by Model1 is a probability value. Specifically, the value of the pixel point (x, y) is a four-dimensional vector, and the length of each direction of the vector represents the probability that the point is the corresponding class. The variance variance of the value of each direction of the vector is calculated. The greater the variance of the point, the greater the probability that the point is located in the boundary region. In order to introduce dark knowledge, the hard label value of the point with a variance value smaller than the hyperparameter v (0.3 in this embodiment) is replaced by a soft label (0.04, 0.6, 0.35, 0.01), such as (0, 1, 0, 0). Model2 will calculate the loss with the new soft label target2, so as to obtain a stronger segmentation model Model2.
[0092] It should be noted that logits and logits1 in the present application are intermediate results, which refer to the original prediction results output by the last layer of the HRSwinNet model without being processed by the softmax activation function during the training process. In the medical image segmentation task, logits is a tensor with dimensions (batch_size, height, width, num_classes), where batch_size is the number of input images, height and width are the height and width of the image, and num_classes is the number of target classes to be segmented. For each pixel position and each target class, logits gives the original probability prediction value of the pixel belonging to the target class. During training, the model compares logits with the true label using the cross-entropy loss function, and then updates the model parameters through backpropagation to optimize the performance of the model. During testing, logits is input into the softmax function to obtain a normalized probability distribution, which is used to predict the segmentation result. Logits is only an intermediate result of a model and is not the same as the final result Output.
[0093] The best segmentation model in the prior art does not consider the segmentation boundary problem, and for the brain tumor image segmentation task, the boundary is blurred and difficult to distinguish. In view of this problem, the present application has carried out the following two works: (1) the present application adopts an adaptive weighting method to adaptively adjust the attention degree of the model to the region, and the weighted value is determined according to the variance of the model. (2) The present application also introduces "dark knowledge" into the boundary region, that is, soft label. This not only makes the model training more simple, but also increases the accuracy of the model.
[0094] For high-resolution networks:
[0095] For CNN-based models, there is a problem with the Conv operator that the receptive field is relatively limited. In order to expand the attention area of the network, multiple layers (convolution-pooling) need to be stacked, which actually brings the problem that the "effective / real" receptive field is also Gaussian decayed outward from a certain center. Therefore, the effective attention point of CNN is usually one or two important parts in the image. In the segmentation model, it is reflected that the resolution of the feature map is gradually reduced, so that the encoder can learn more abstract / high-level semantic information through the gradually increasing receptive field. However, this also exposes a problem: due to the limited receptive field, CNN-based models are difficult to learn long-distance feature dependency, which is usually crucial for high-level computer vision tasks.
[0096] The long-distance feature of the Transformer naturally makes it possible to use global effective information from shallow to deep, and the multi-head attention mechanism ensures that the network can pay attention to multiple important features. In fact, each head is an independent attention mechanism, which is one of the differences between Transformer and CNN. However, considering the size of the model and the problem of overfitting, the existing Transformer model does not completely use the full-resolution model to complete the computer vision task. The existing Transformer model usually uses the patch cutting method to reduce the model size.
[0097] The present application constructs a real high-resolution network, parallel decoding streams of different resolutions, and interacts and transmits different resolution feature information, thereby realizing high-resolution and high-semantic feature modeling.
[0098] For boundary regions:
[0099] Boosting is a common statistical learning method. In the classification problem, it learns multiple classifiers (usually weak classifiers) by changing the weights of training samples, and linearly combines these classifiers to improve the performance of the classifier. For this boosting method, two questions need to be answered: one is how to change the weight or probability distribution of the training sample in each round; the other is how to combine weak classifiers into a strong classifier. AdaBoost (Adaptive Boosting) algorithm is a kind of Boosting, which can well solve the above two problems. For the first question, AdaBoost increases the weight of the samples that are misclassified by the weak classifier in the previous round, and reduces the weight of the samples that are correctly classified. Thus, the samples that are misclassified will be paid more attention by the weak classifier in the next round due to their increased weight.
[0100] People usually intuitively believe that the number of parameters of a model basically determines the amount of "knowledge" that it can capture from the data.
[0101] This idea is basically correct, but it needs to be noted that:
[0102] (1) The relationship between the number of parameters of a model and the amount of "knowledge" it can capture is not a stable linear relationship (1 in the figure), but a growth curve close to diminishing marginal returns (2 and 3 in the figure). Figure 6 Figure 6
[0103] (2) The same model architecture and the same number of model parameters, using the same training data, do not necessarily capture the same amount of "knowledge". Another key factor is the training method. A suitable training method can make it possible to obtain more "knowledge" with a smaller total number of model parameters (3 vs. 2 in the figure). Figure 6
[0104] In the MNIST handwritten digit recognition task, there are 10 output categories. Suppose a certain input "2" looks more like "3" (as shown in Figure 7 ), the probability of "3" in the softmax output value is 0.1, and the values of other negative labels are very small. Another "2" looks more like "7", and the probability of "7" is 0.1. The hard labels of the two "2"s are the same, but their soft labels are different. Therefore, the soft label contains more information than the hard label. When the entropy of the soft label distribution is relatively high, the knowledge it contains is more abundant.
[0105] For brain tumor, there are four classes: background, edema, non-enhancing tumor, and enhancing tumor. The boundaries of different classes are difficult to distinguish because tumors are usually surrounded by healthy brain tissue, such as cerebrospinal fluid, gray matter, and white matter, resulting in ambiguous tumor boundaries, which makes it difficult for deep learning networks to distinguish these boundary regions.
[0106] Conventional model training usually uses hard labels for learning. For a certain point in an image, its hard label is 2. During learning, the model knows only that the point is class 2, without any other information. Inspired by knowledge distillation, the present application considers using soft labels for learning in the boundary region. Suppose the soft label of the point is [3%, 20%, 70%, 7%]. The meaning of this soft label is that the probability of the point being label 0 is 3%, the probability of the point being label 1 is 20%, the probability of the point being label 2 is 70%, and the probability of the point being label 3 is 7%. This soft label contains more knowledge than the hard label and is meaningful for distinguishing boundaries.
[0107] Therefore, the method of the present application strengthens the model's learning of the boundary in two ways: (1) strengthening the weight of the boundary region; and (2) introducing "dark knowledge" for the boundary region, thereby improving the model's ability to distinguish boundaries.
[0108] Test design:
[0109] This example uses two data sets, brain tumors in the MSD data set and the BraTS2021 data set.
[0110] BraTS2021 contains 1251 MRI scan images with a shape of 240x240x155, and follows VT-UNet, dividing the 1251 MRI scan images into 834, 208, and 209 for training, validation, and testing, respectively. The task is to segment three semantically meaningful tumor classes, namely enhancing tumor (ET), tumor core (TC) region, and whole tumor (WT) region. The brain data set of MSD consists of 484 MRI images, each including four channels, namely FLAIR, T1w, T1gd, and T2w. The corresponding target ROI is consistent with BraTS2021, with the training / validation / testing set ratio being 80%, 15%, and 5%, respectively. The data set division is shown in the following table.
[0111] Table 1: Data set division
[0112] Dataset Training set (80%) Validation set (15%) Test set (5%) BraTS2021 1000 188 63
[0113] During model training, the depth of each Swin Transformer block is 2, and the resolution is and The attention heads of the Swin Transformer block of HRSwinNet are 3, 6, 12, and 24, respectively. AdamW optimizer is used to optimize the parameters of HRSwinNet with a learning rate of 1e-4. Linear warm-up cosine decay learning rate scheduler is used to adjust the value of the learning rate, and the warm-up epoch is 50. The sum of the dice loss and the cross-entropy loss is used as the loss function. For fair comparison, the data set preprocessing and model evaluation methods are the same as VT-UNet on the BraTS2021 data set. The batch size is 1, and the MRI image scans of BraTS2021 are cropped to a fixed size of 128x128x128, and the cropped MRI scans are fed into HRSwinNet. To illustrate its superior performance, HRSwinNet is compared with a variety of CNN-based segmentation methods and Transformer-based methods on the BraTS2021 data set. The above methods are trained for 300 epochs on the training data set of BraTS2021 in this embodiment, and the performance of these methods on the test data set of BraTS2021 is compared.
[0114] The brain data set training details of MSD are the same as the training on the BraTS2021 data set.
[0115] The model with the best Dice-Similarity coefficient (DSC) on the validation data set is saved for final evaluation, and model evaluation techniques such as model ensemble are not used in this paper. All experiments in this embodiment are performed using two RTX 3090 GPUs. For fair comparison, the same data set preprocessing method, training device, and evaluation method are used to retrain the three comparison methods, and the results of the retraining are reported.
[0116] Dice score and Hausdorff distance (HD) are used to quantitatively evaluate the segmentation results of the model on the two data sets, and for fair comparison, the method provided by nnUNet is used to calculate the Hausdorff distance in this embodiment.
[0117] (1) BraTS2021 experimental results
[0118] Table 2: BraTS2021 experimental results
[0119]
[0120]
[0121] Dice score and Hausdorff distance (HD) are used to quantitatively evaluate the segmentation results of the model on the two data sets, and for fair comparison, the method provided by nnUNet
[77] The provided method calculates the Hausdorff distance.
[0122] As shown in Table 2, the experimental results on the BraTS2021 dataset demonstrate that the proposed HRSwinNet achieves the best average Dice score and average Hausdorff distance, with its average Dice score comparable to VT-UNet. For enhanced tumor segmentation, VT-UNet-B exhibits the best Dice score, but its Hausdorff distance is relatively large. The Dice score of Swin UNETR is comparable to that of UNETR. In summary, the experimental results shown in Table 2 indicate that HRNet-like networks can achieve better performance than UNet-like architectures, suggesting that HRNet-like network design is a valuable research direction.
[0123] (2) MSD Experiment Results
[0124] To demonstrate the model's generalization performance, this embodiment also conducted similar experiments on the MSD brain dataset.
[0125] Table 3: MSD Experimental Results
[0126]
[0127] As shown in Table 3, the experimental results on the MSD dataset indicate that the proposed HRSwinNet achieves the best average Hausdorff distance, and its average Dice score is also highly competitive, although slightly inferior to VT-UNet. Similar to the results on the BraTS2021 dataset, VT-UNet-B still performs best in Dice score for enhanced tumor segmentation, but its Hausdorff distance is larger. The reason for the poor performance of the proposed method in terms of average Dice score is that the MSD dataset is much smaller than BraTS2021, while the proposed model is relatively larger, leading to this result.
[0128] Both the adaptive weight-based boundary enhancement method and the self-distillation-based boundary enhancement method are boundary enhancement algorithms, so they are analyzed together in this embodiment. To better illustrate the effectiveness of the two methods in this study, the following ablation experiment results are presented in this embodiment, which were conducted on the BraTS2021 dataset:
[0129] Table 4: Ablation Experiment Results for Each Module
[0130] Models AvgHD(95)↓ AvgDice↑ BasedModel 12.84 82.87 HRSwinNet 10.93 85.56 HRSwinNet+AW 9.43 86.64 HRSwinNet+AW+SA 9.25 87.49
[0131] where Base Model is the model of HRSwinNet without the high-resolution module, AW is the boundary enhancement module of adaptive weight, and SA is the boundary enhancement module of self-distillation.
[0132] According to the data in the table, it is obvious that the high-resolution network and the two boundary enhancement methods improve the segmentation effect to some extent. The high-resolution model improves the segmentation effect comprehensively, and both evaluation indicators are improved to a certain extent. The improvement of the two boundary enhancement modules to the segmentation result is more on the boundary, and the boundary area accounts for a very small proportion of the whole tumor, so the improvement of the Dice score is very small. However, both of them have a certain improvement on the HD95 index which can reflect the boundary segmentation situation.
[0133] Effectiveness of the boundary enhancement method based on self-distillation:
[0134] Table 5: Comparison of model and computational complexity
[0135] Models #Params(M) FLOPs(G) InferenceTime(s) nnUNet 19.07 412.65 10.28 CoTr 46.51 399.21 19.21 SwinUNETR 62.29 192.10 18.63 VT-UNet-B 65.93 158.49 15.72 TransUNet 96.07 48.34 26.97 UNETR 92.58 41.19 12.08 Ours 99.57 215.60 22.39
[0136] In Table 5, the number of FLOPs, parameters, and average inference time of the model in the BraTS2021 test set are shown. The number of FLOPs and inference time are calculated based on the input size of 128x128x128 and using the sliding window method. According to the benchmark, HRSwinNet is a medium-sized model with 99.57M parameters and 215.60G FLOPs. For comparison, other transformer-based methods such as CoTr, VT-UNet-B, and TransUNet have 46.51M, 65.93M, and 96.07M parameters and 399.21G, 158.49G, and 48.34G FLOPs, respectively. HRSwinNet shows comparable model complexity while significantly outperforming these models in the BraTS2021 test. nnUNet is a CNN-based segmentation model with 19.07M parameters and 412.65G FLOPs. Similarly, HRSwinNet outperforms CNN-based models while having no excessive model complexity.
[0137] The above only describes the preferred embodiments of the present application and is not used to limit the protection scope of the present application. Any modification, equivalent replacement, improvement, etc. made within the spirit and principles of the present application shall be included in the protection scope of the present application.
Claims
1. A method for MRI image segmentation based on high-resolution networks and boundary enhancement, characterized in that, Includes the following steps: Step S1: Perform preprocessing operations on the MRI images that need to be segmented; Step S2: The preprocessed image data is then input into the first high-resolution segmentation network model. and the second high-resolution segmentation network model Both models are high-resolution segmentation network models based on the Swing Transformer, namely HRSwinNet. Employ boundary reinforcement methods; The boundary enhancement method is a boundary enhancement learning method based on adaptive weights and / or a boundary enhancement learning method based on knowledge distillation; Step S3, Final output and Final output The final segmentation result is obtained by averaging the results. The boundary reinforcement learning method based on adaptive weights in step S2 is specifically as follows: First, the high-resolution segmentation network model HRSwinNet is trained using standard methods. The trained model is denoted as... Input the training data and the model to be trained The output result is ,go through Obtain the probability value ,calculate and the true value cross-entropy Then the value of C is passed through the Sigmoid function and mapped to... In between, the value obtained at this time is to be... Weighted weights Values, using weights right By weighting, we obtain the enhanced model. ; The During training, the loss is calculated by comparing it with the weights. Calculate the inner product, and the final result is the weighted loss; The boundary reinforcement learning method based on knowledge distillation in step S2 is specifically as follows: First, the high-resolution segmentation network model HRSwinNet is trained using standard methods. The trained model is denoted as... Training data enters simultaneously and the model to be trained The output result is Subsequently Performing a Softmax operation yields a four-dimensional vector for each pixel. The length of each direction of the vector represents the probability that the pixel belongs to its corresponding category. Calculate the variance of the values in each direction of the four-dimensional vector. , The larger the variance, the greater the probability of it being located in the boundary region; Less than hyperparameters Replace the hard label values of the points with soft labels, With the new soft label Calculate the loss to obtain a stronger segmentation model. .
2. The MRI image segmentation method based on high-resolution networks and boundary enhancement according to claim 1, characterized in that, The preprocessing in step S1 includes three methods: image cropping, resampling, and normalization.
3. The MRI image segmentation method based on high-resolution networks and boundary enhancement according to claim 1, characterized in that, The HRSwinNet model consists of encoding stages, which are sequentially connected. A multi-resolution feature fusion (MRFF) module is inserted between each stage. Except for the last stage, the subsequent stages... Each Stage contains The system consists of two parallel Swing Transformer modules and a patch merging module. The last stage contains... A parallel SwinTransformer module and The last multi-resolution feature fusion module (MRFF) has an output that is connected to the residual module, the patch expanding module, and the Conv module in sequence. The patch merging module and the patch expanding module are used to downsample and upsample the feature maps, respectively.
4. The MRI image segmentation method based on high-resolution networks and boundary enhancement according to claim 3, characterized in that, The Swin Transformer module contains two cascaded layers: the first layer is a window-based multi-head self-attention module W-MSA, and the second layer is a shift-window-based multi-head self-attention module SW-MSA. A W-MSA module and a SW-MSA module are connected in sequence. The continuous calculation process of the Swing Transformer module is as follows: In the formula, and Represent Module, Module in The output features of the layer This indicates the output of the previous layer. This represents the output of the next layer; LN stands for layer normalization, and MLP stands for multilayer perceptron.
5. The MRI image segmentation method based on high-resolution networks and boundary enhancement according to claim 3, characterized in that, The loss function of the HRSwinNet model is: Where I represents the number of image voxels; J is the category number; and Let i represent the output probability of class j at voxel i and the true value of a one-hot encoded voxel i, respectively.