Lightweight 3D heart nuclear magnetic resonance image segmentation algorithm based on three-dimensional sparse MLP and 3D efficient deformable attention

Abstract

The invention discloses a lightweight 3D heart MRI segmentation algorithm based on 3D efficient partial attention and three-dimensional sparse MLP, and the algorithm obtains an MRI pixel block with a detailed target region feature through a 3D efficient deformable attention module EDA, and fits the edge distribution of a segmented region with lower calculation complexity through reference point translation. Efficient attention calculation ensures feature information extraction and greatly reduces resource consumption; the method comprises the following steps: constructing a cross-shaped receptive field at extremely low cost through information in transverse, longitudinal and depth directions of a three-dimensional sparse MLP module (triSMLP) mixed voxel block, and outputting a fine characteristic spectrum with a global receptive field through three-time superposition; the feature maps output by the EDA and the triSMLP recover down-sampling lost detail features at an encoder through a partial convolution feature fusion decoder. Compared with most of the leading edge 3D segmentation algorithms at present, the method has the advantage that a better segmentation effect is obtained with extremely small computing resources.

Description

Technical Field

[0001] The present invention relates to a deep neural network, in particular to an algorithm for improving the progress and speed of 3D medical image segmentation by using three-way sparse MLP and 3D deformable attention, and belongs to the field of semantic segmentation. Background Art

[0002] Among urban and rural residents who died of diseases, the number of deaths from cardiovascular diseases was higher than that from other diseases. In 2020, cardiovascular diseases accounted for 48.00% and 45.86% of deaths in rural and urban areas respectively. Two out of every five deaths were caused by cardiovascular diseases. It can be seen that cardiovascular diseases are the primary threat to people's lives and health. Prevention and treatment of cardiovascular diseases are urgent tasks. Rapid and effective diagnosis of whether a patient has cardiovascular diseases is the prerequisite for prevention and treatment, and it is also one of the current hot medical issues.

[0003] In medical institutions, the main method for detecting cardiovascular diseases is to perform magnetic resonance imaging of the heart to obtain a cardiac magnetic resonance image (MRI). MRI imaging has almost no effect on the human body, and has high imaging quality and high organ differentiation. In the field of cardiovascular diseases, cardiac magnetic resonance imaging technology is considered to be the gold standard image diagnosis method for diagnosing cardiovascular diseases.

[0004] Segmentation based on cardiac MRI is a prerequisite for the diagnosis of cardiovascular diseases. Segmentation of the target area can provide some parameters of cardiac ventricular function, such as ejection fraction (EF), ventricular end-diastolic volume (VEDV), ventricular end-systolic volume (VESV), and ventricular myocardial mass (VMW). These are very important for determining whether the heart is abnormal. Therefore, cardiac organ segmentation is a prerequisite and an important step in the diagnosis and treatment of cardiovascular diseases. In actual clinical medicine, segmentation of various cardiac tissues is performed by experienced radiologists through visual judgment and manual segmentation. This method is time-consuming and labor-intensive. Automatic segmentation methods are very important for quantitative index analysis of the heart.

[0005] There are many methods based on deep learning, but they also face many challenges: 1. It is difficult to achieve excellent results in the detailed segmentation of organs. Since the myocardium and ventricles are intertwined in the area of ​​heart segmentation, the edge division of the myocardium is one of the biggest difficulties today; 2. Most of the current segmentation algorithms are based on 2D images. MRI data itself is 3D data. Most studies are sliced ​​and segmented at the 2D level, which lacks high-dimensional semantic information and is also clinically limited; 3. It has high requirements on computing resources. The attention-based 3D segmentation network that has become popular in recent years has made great progress in segmentation effects, but the number of parameters, network size, and computing time of the network itself have also increased dramatically, and there is still a long way to go in practical applications. Summary of the invention

[0006] In order to solve the current problems in heart segmentation, the present invention integrates a three-way sparse MLP module and a 3D efficient deformable attention module. This method can effectively improve the 3D cardiac MRI segmentation accuracy, has good segmentation performance, and reduces the operating requirements compared to other 3D models. The number of parameters is 3 / 100 of the models with similar effects.

[0007] The objective of the present invention is achieved through the following technical scheme: a feature fusion 3D cardiac magnetic resonance segmentation algorithm based on attention and improved MLP, the algorithm comprising a continuous small kernel convolution module for position encoding, a 3D efficient deformable attention and sparse MLP encoder and a partial convolution feature fusion decoder.

[0008] The continuous small-size convolution kernel transforms the input large-size 3D cardiac MRI image data by position encoding, divides the input image into small blocks of fixed size, and then maps it to a high-dimensional space, embedding the position information between the blocks.

[0009] The 3D efficient deformable attention module includes: efficient attention calculation of reference points, offset embedding and relative position deviation integration. The input of this module is the cardiac MRI voxel block with embedded position information, and the output is the MRI voxel block with segmented region enhancement. EDAT previously processed the original input F (H×W×D×C) with continuous small-size convolution patch embedding to obtain a low-resolution block F L , the specific calculation is:

[0010] F L =Conv(Conv(Conv(F)))

[0011] Voxel Block F P (H P ×W P ×D P ×C), in F P Random sampling FS (H P / r×W P / r×D P / r×C), where r is a hyperparameter controlling the scaling, and the reference voxel block is in a 3D normal grid The reference points are linearly arranged in 3D space coordinates, and the 3D normal grid used to store the voxel block coordinates is:

[0012] {(p x ,p y ,p d )|(0,0,0),…,(W G -1,H G -1,D G -1)}

[0013] The 3D normal grid sampling converts the sampling point coordinates into the input ratio (H P ×W P ×D P ) is scaled to [-1,1], where [-1,-1,-1] represents the point in the eighth quadrant with the farthest Euclidean distance from the origin, and [1,1,1] represents the point in the first quadrant with the farthest Euclidean distance from the origin.

[0014] The offset embedding network: The network input is the original input feature map, which is mapped through a three-dimensional query matrix to obtain the global q vector used to calculate the attention, and the output is multiple sets of three-dimensional vectors Δp, representing the offset direction and distance of the offset point. The offset calculation is performed by a simple neural network θ offset The composition mainly includes feature learning module and feature dimension reduction module;

[0015] Δp=θ offset (q),

[0016] The area near the offset reference point is sampled to obtain the area x for actual attention calculation, which is mapped by the three-dimensional matrix. and The specific calculation is:

[0017]

[0018] Where W k and W v is a three-dimensional matrix used for mapping, and the resulting Compute efficient self-attention with q.

[0019] The Efficient attention module, in the input Before calculating q, change the order of calculation, that is,

[0020]

[0021] In order to make the result obtained after the efficient attention changes the calculation method consistent with the original attention, it is necessary to And q normalization, that is, the softmax function is applied along each row, column and depth of the voxel, and the specific calculation method is:

[0022]

[0023] The three-way sparse MLP module includes reconstruction fusion in three different dimensions. The sparse MLP module takes the attention feature map output by 3DEDAT as input, and through the concatenation of three dimensions and the superposition of three sparse MLP modules, it realizes the diffusion of the receptive field and obtains a high-dimensional feature image with edge semantic features.

[0024] The reconstruction fusion in the three-way sparse MLP is to convert the input feature F (H×W×D×C) into HC×W×D in the height dimension, i.e., the H dimension. A linear layer acts on the H and C dimension information of the mixed input feature F to mix the HC tensor of each row. In the same way, the three dimensions H, W, and D of the input feature are reconstructed and linearly projected in turn. Finally, the reconstruction results of the three dimensions are fused to obtain the cross-shaped receptive field within the voxel. The specific calculation method is:

[0025] X H =FW H , X W =FW W , X D =FW D

[0026] X out =FC(concat(X H ,X W ,X D ))

[0027] Furthermore, a single sparse MLP module provides a cross-shaped receptive field centered on a voxel in the voxel block, and the receptive field is diffused globally by stacking multiple sparse MLP modules.

[0028] The partial convolution feature fusion decoder includes an upsampling channel, a partial convolution module, a deep convolution module and a residual connection. For the encoder's multi-layer input F i (i=1,2,3,4), restored to the output size through upsampling channels Concatenate to get the decoder input F D .

[0029] Furthermore, for the input feature map F D (H D ×W D ×D D ×C), through partial convolution, and through a linear layer to expand the number of channels to Will Divided into two parts by channel and After separating the height, width and depth, a depth convolution is performed to merge the height, width and depth into and Matrix multiplication to obtain The channel is restored through a linear layer Finally, the residual module converts the original input F D and The final high-quality output is obtained by stitching. The specific calculation method is:

[0030]

[0031] Where W1 and W2 represent linear projection; R() and F() represent Reshape and Flatten operations; PConv() and DWConv() refer to partial convolution and depth convolution operations respectively; Represents matrix multiplication.

[0032] Furthermore, the experimental data is the ACDC (Automated Cardiac Diagnosis Challenge) dataset, which was obtained from real clinical examinations by the Dijon University Hospital. The dataset covers examination results from 150 different patients, evenly divided into 4 pathological groups and 1 healthy subject group. Each patient's examination result is also accompanied by additional information, including weight, height, and end-diastolic and end-systolic records. The database is provided to participants through two datasets on a dedicated online evaluation website after personal registration: (1) The training dataset is analyzed by clinical experts and includes data from 100 patients and other reference data; (2) The test dataset consists of 50 new patients without manual annotations, but contains the above patient information. The original images are saved in Nifti format. The purpose of this challenge dataset is to compare the performance of automatic methods for segmenting the left ventricular endocardium and epicardium into the right ventricular endocardium in both end-diastolic and end-systolic conditions.

[0033] Furthermore, a heart segmentation algorithm is constructed for model training, the data set is loaded for training, and then the optimal model is saved, and the trained model is used to predict the test set of the data.

[0034] The beneficial effects of the present invention are:

[0035] The lightweight 3D cardiac MRI segmentation algorithm based on three-way sparse MLP and 3D efficient partial attention disclosed in the present invention obtains MRI pixel blocks with detailed target area features through a 3D efficient deformable attention module (EDA), and uses reference point translation to fit the edge distribution of the segmented area with lower computational complexity. The efficient attention calculation ensures feature information extraction while greatly reducing resource consumption; through the three-way sparse MLP module (triSMLP), the information of the horizontal, vertical and depth directions of the voxel block is mixed to construct a cross-shaped receptive field at a very low cost, and a fine feature map with a global receptive field is output through three superpositions; the feature map output by EDA and triSMLP is restored at the encoder through a partial depth convolution recovery network to restore the detailed features lost by downsampling. Compared with most of the current cutting-edge 3D segmentation algorithms, the present invention obtains a more superior segmentation effect with very small computing resources.

[0036] Other advantages, objectives and features of the present invention will be described in the following description to some extent, and to some extent, will be obvious to those skilled in the art based on the following examination and study, or can be taught from the practice of the present invention. The objectives and other advantages of the present invention can be realized and obtained through the following description. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] In order to make the purpose, technical solutions and advantages of the present invention more clear, the present invention will be described in detail below in conjunction with the accompanying drawings, wherein:

[0038] Figure 1 This is a schematic diagram of the improved overall network architecture;

[0039] Figure 2 Schematic diagram of 3D efficient deformable attention module;

[0040] Figure 3 Schematic diagram of the three-way Sparse-MLP module;

[0041] Figure 4 This is a schematic diagram of a partial convolution feature fusion module;

[0042] Figure 5 A comparison chart of the ACDC dataset labels and network segmentation results. DETAILED DESCRIPTION

[0043] The following is further explained in detail through specific implementation methods.

[0044] like Figure 1 As shown, the 3D cardiac MRI segmentation algorithm provided by the present invention includes a 3D efficient deformable attention module, a sparse MLP module and a partial convolution feature fusion module.

[0045] The model consists of eight layers. Taking the voxel input size of 1024*1024*1024 as an example, the first layer contains two convolutional layers to convert the input into small-sized voxels with a size of 256*256*256; the second to fifth layers are feature extraction modules, each of which contains a 3D efficient deformable attention module, a sparse MLP module and a downsampling module. Each feature extraction layer reduces the input size by half, namely 256*256*256, 128*128*128, 64*64*64 and 32*32*32 respectively. In addition to being used for the output of the next layer, the input of each layer also passes the output to the splicing module of the decoder; the sixth layer contains an upsampling module and splicing, which upsamples the input of each stage to an output size of 256*256*256 and splices it; the seventh layer contains 4 partial convolutional feature recovery modules; the eighth layer is a linear layer, which converts the feature map into an output size of 256*256*256.

[0046] The 3D efficient deformable attention module: Figure 2 , at the beginning, the reference point of the deformable attention is randomly selected, and the query value q is generated by the query matrix, q = xW q , the query value q is passed into the reference point generation network. The reference point generation network first passes q of size H×W×D×C through a depth-wise separable convolution with a step size of r, where r is the scaling factor and the size of the reference area is H G =H / r,W G =W / r, D G =D / r, the processed q shape is H G ×W G ×D G ×C, regularized and activated, then passed through a 1×1×1 convolution, the output size is H G ×W G ×2, which is the distance Δp moved by the reference point, and p+Δp is the distance of one reference point movement. The way to calculate attention in the reference area is to perform bilinear sampling on the reference area image to obtain the actual x used to calculate v and k, recorded as The result Calculate the efficient attention with q, and the efficient attention calculation formula is:

[0047]

[0048] After two dot product operations The similarity between pixels is calculated separately and the value is aggregated using the attention matrix to form the output. The calculation complexity is greatly reduced by changing the order of operations in the two dot products. q and ρ kThey represent q and Normalization function:

[0049]

[0050] ρ q (Y) = σ row (Y),ρ k (Y) = σ col (Y)

[0051] The three-way sparse MLP module shown in Figure 3 , for the input feature F(H×W×D×C), in the horizontal mixing path, the data tensor is reshaped into HC×W×D, and a linear layer is applied to each HC row to mix the information. In the same way, the three dimensions H, W, and D of the input feature are reconstructed and linearly projected in turn, and finally the reconstruction results of the three dimensions are fused to obtain the cross-shaped receptive field within the voxel. The specific calculation method is:

[0052] X H =FW H , X W =FW W , X D =FW D

[0053] X out =FC(concat(X H ,X W ,X D ))

[0054] This approach allows each voxel to aggregate information across the current row and across columns and depths. By passing the module three times, each voxel can aggregate information in the entire 3D space. Despite the use of sparse connections, multiple passes allow the entire sparse MLP module to effectively obtain a global receptive field.

[0055] Parameter analysis of three-way sparse MLP: The number of parameters of a sparse MLP is H 2 +W 2 +D 2 +3C 2 , the original MLP module parameter is 2α (HWD) 2 , where α is the scaling ratio of the MLP. Taking the first layer input of 256*256*256, α=4, and the number of channels of 256 as an example, the number of parameters of the sparse MLP is 1 / 393216 of that of the ordinary MLP.

[0056] The partial convolution feature fusion module, such as Figure 4, for the input feature F, one copy is used to construct the residual connection, and one copy of the input feature map is divided into multiple sub-regions by channel, and only a part of them is convolved. By selectively applying convolution to some channels and then putting the results back to the original position, the amount of calculation can be reduced:

[0057] F1, F2 = Slicing (F),

[0058] F1=Conv(F1).

[0059] Then, a linear layer is used to integrate the number of input channels into C', and the input features are divided into two parts with the same number of channels, and one part remains unchanged for residual connection. Reshape separates the height, width and depth for depthwise convolution, flattens the features, and performs matrix dot product calculation with the input before channel separation. The output is restored to the number of output channels through a linear layer and added point by point with the original input to obtain the final feature map.

[0060] The output of the partial convolutional feature fusion module is an abstract feature map, and finally a 1×1 convolution is performed to obtain a result with a segmentation mask.

[0061] The cross entropy loss function is one of the commonly used loss functions in medical image segmentation. This loss function is used to measure the difference between two probability distributions and is usually used for classification problems. In image segmentation, each pixel can be regarded as a category, and then the cross entropy loss function is used to measure the difference between the model's classification prediction of each voxel and the actual label. The calculation formula of the cross entropy loss function is as follows:

[0062]

[0063] Where N is the total number of pixels, C is the number of categories, which is 4 in the present invention, and the 4 categories are left ventricle, right ventricle, myocardium and background area, y ij is the probability that the i-th voxel of the actual label belongs to the j-th category, is the probability that the model predicts that the i-th voxel belongs to the j-th category.

[0064] The Dice loss function is a commonly used loss function in medical image segmentation tasks. It is used to measure the degree of spatial overlap between the model's predicted segmentation results and the true labels. The mathematical definition of the Dice loss function is as follows:

[0065]

[0066] Where Intersection is the intersection of the model prediction result and the true label, Union is their union, and ∈ is a smooth term used to avoid the denominator being zero.

[0067] The combination of the two is used as the final loss function, which is defined as follows:

[0068] Loss=λCE Loss+(1-λ)Dice Loss

[0069] Wherein λ is a factor controlling the weight of the two, and is set to 0.65 in the present invention.

[0070] The evaluation index of the present invention uses Dice to measure the overlap between the segmented area and the real segmented area of ​​the label, and the range is [0,1]. The larger the value, the closer the tumor segmentation result is to the labeling result, and the better the segmentation effect. The specific calculation method is as described above.

[0071] The present invention is compared on the ACDC dataset. The experimental results are shown in the table. Avg represents the average segmentation results of the three regions, RV represents the segmentation results of the right ventricle, MYO represents the segmentation results of the myocardium, LV represents the segmentation results of the right ventricle, Params represents the parameter quantity of the model, and the maximum value is bold. It can be seen from the table that the algorithm we proposed has a great advantage in the parameter quantity, and the segmentation effect in each segmentation area is also excellent. Figure 5 The segmentation performance of several segmentation algorithms on ACDC.

[0072] Table 1 Performance of some algorithms on the ACDC dataset

[0073]

Claims

1. A lightweight 3D cardiac MRI segmentation algorithm based on three-way sparse MLP and 3D efficient deformable attention, which includes a continuous small kernel convolution module for position encoding, a 3D efficient deformable attention and sparse MLP encoder and a partial convolution feature fusion decoder. The continuous small-size convolution kernel transforms the input large-size 3D cardiac MRI image data by position encoding, divides the input image into small blocks of fixed size, and then maps it to a high-dimensional space, embedding the position information between the blocks. The 3D efficient deformable attention module includes: efficient attention calculation of reference points, offset embedding and relative position deviation integration. The input of this module is the cardiac MRI voxel block with embedded position information, and the output is the MRI voxel block with segmented region enhancement. EDAT previously processed the original input F (H×W×D×C) with continuous small-size convolution patch embedding to obtain a low-resolution block F L , the specific calculation is: F L =Conv(Conv(Conv(F))) Voxel Block F P (H P ×W P ×D P ×C), in F P Random sampling F S (H P / r×W P / r×D P / r×C), where r is a hyperparameter controlling the scaling, and the reference voxel block is in a 3D normal grid The reference points are linearly arranged in 3D space coordinates, and the 3D normal grid used to store the voxel block coordinates is: {(p x ,p y ,p d )|(0,0,0),…,(W G -1,H G -1,D G -1)}。 2. The feature fusion 3D cardiac MRI segmentation algorithm based on attention and improved MLP according to claim 1, characterized in that: The 3D normal grid sampling converts the sampling point coordinates into the input ratio (H P ×W P ×D P ) is scaled to [-1,1], where [-1,-1,-1] represents the point in the eighth quadrant with the farthest Euclidean distance from the origin, and [1,1,1] represents the point in the first quadrant with the farthest Euclidean distance from the origin. The offset embedding network: The network input is the original input feature map, which is mapped through a three-dimensional query matrix to obtain the global q vector used to calculate the attention, and the output is multiple sets of three-dimensional vectors Δp, representing the offset direction and distance of the offset point. The offset calculation is performed by a simple neural network θ offset The composition mainly includes feature learning module and feature dimension reduction module; The area near the offset reference point is sampled to obtain the area x for actual attention calculation, which is mapped by the three-dimensional matrix. and The specific calculation is: Where W k and W v is a three-dimensional matrix used for mapping, and the resulting Compute efficient self-attention with q. The Efficient attention module, in the input Before calculating q, change the order of calculation, that is, In order to make the result obtained after the efficient attention changes the calculation method consistent with the original attention, it is necessary to And q normalization, that is, the softmax function is applied along each row, column and depth of the voxel, and the specific calculation method is:

3. The feature fusion 3D cardiac MRI segmentation algorithm based on attention and improved MLP according to claim 1, characterized in that: The sparse MLP module includes reconstruction fusion in three different dimensions. The sparse MLP module takes the attention feature map output by 3DEDAT as input, and through the concatenation of three dimensions and the superposition of three sparse MLP modules, it realizes the diffusion of the receptive field and obtains a high-dimensional feature image with edge semantic features. The reconstruction fusion in the sparse MLP is to transform the input feature F (H×W×D×C) into HC×W×D in the height dimension, i.e., the H dimension. A linear layer acts on the information of the H and C dimensions of the mixed input feature F, and mixes the HC tensor of each row. In the same way, the three dimensions H, W, and D of the input feature are reconstructed and linearly projected in turn. Finally, the reconstruction results of the three dimensions are fused to obtain the cross-shaped receptive field within the voxel. The specific calculation method is: X H =FW H ,X W =FW W ,X D =FW D X out =FC(concat(X H ,X E ,X D )) Furthermore, a single sparse MLP module provides a cross-shaped receptive field centered on a voxel in the voxel block, and the receptive field is diffused globally by stacking multiple sparse MLP modules.

4. The feature fusion 3D cardiac magnetic resonance segmentation algorithm based on attention and improved MLP according to claim 1, characterized in that: The partial convolution feature fusion decoder includes an upsampling channel, a partial convolution module, a deep convolution module and a residual connection. For the encoder's multi-layer input F i (i=1,2,3,4), restored to the output size through upsampling channels Concatenate to get the decoder input F D . Furthermore, for the input feature map F D (H D ×W D ×D D ×C), through partial convolution, and through a linear layer to expand the number of channels to Will Divided into two parts by channel and After separating the height, width and depth, a depth convolution is performed to merge the height, width and depth into and Matrix multiplication to obtain The channel is restored through a linear layer Finally, the residual module converts the original input F D and The final high-quality output is obtained by stitching. The specific calculation method is: Where W1 and W2 represent linear projection; R() and F() represent Reshape and Flatten operations; PConv() and DWConv() refer to partial convolution and depth convolution operations respectively; Represents matrix multiplication.

Images