Medical image segmentation model establishment method based on harmonic attention and medical image segmentation method
By introducing the harmonic attention layer and Transformer module into the 3D medical image segmentation model, the problem of insufficient extraction of correlation information between the depth dimension and other dimensions in 3D medical image segmentation is solved, high-precision segmentation of the lesion area is achieved, and the segmentation effect is improved.
Patent Information
- Application Number
- CN202411024514.0
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-07-29
- Publication Date
- 2025-10-03
- Estimated Expiration
- 2044-07-29
AI Technical Summary
Existing three-dimensional medical image segmentation methods cannot effectively extract the correlation information between the depth dimension and other dimensions, resulting in poor segmentation results, especially when the outline of the lesion area is not clear, the segmentation accuracy is insufficient.
A medical image segmentation model based on harmonic attention is adopted. By introducing a harmonic attention layer in the encoder-decoder structure, combining the Transformer module and the modal fusion module, the feature distribution between different planes is harmonized. Hard and soft attention branches and alternating position encoding are used to extract and fuse multimodal information to achieve accurate modeling of the lesion area.
It improves the accuracy of three-dimensional medical image segmentation, can more accurately segment discontinuous and unclear lesion areas, make full use of multimodal information, stabilize feature distribution, and enhance the accuracy and efficiency of segmentation.
Smart Images

Figure CN118967714B_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the field of three-dimensional medical image segmentation, and more specifically, relates to a method for establishing a medical image segmentation model based on harmonic attention and a medical image segmentation method. Background Art
[0002] Medical image segmentation refers to the process of segmenting lesion areas from medical images. Early traditional medical image segmentation methods mainly focused on mathematical methods such as edge detection, template matching technology, region growing, graph cutting, active contours, and machine learning. Traditional segmentation methods have problems such as low segmentation efficiency and low accuracy. In recent years, deep learning technology has gradually become the research focus in fields such as computer vision and image processing. Compared with the high sensitivity of traditional segmentation algorithms to images, convolutional neural networks (CNNs) can extract hierarchical features of images. CNNs are not sensitive to image properties such as noise, blur, and contrast, and they provide good segmentation results for medical images.
[0003] In medical image segmentation, many deep learning-based methods have been widely used, such as residual networks (ResNets), fully convolutional networks (FCNs), and U-Nets based on convolutional neural networks (CNNs). U-Net, among others, employs a symmetric encoder-decoder structure with skip connections, which effectively preserves image details. Although many improved network structures have subsequently emerged, convolution operations struggle to capture global semantic information. Therefore, inspired by the attention mechanism in natural language processing, existing research has overcome this limitation by combining attention mechanisms with CNN models. Recently, Transformer-based network frameworks, exemplified by the Visual Transformer (ViT), have achieved remarkable performance on various computer vision tasks. This approach segments images into blocks and models the correlations between these blocks using a transformer sequence. While this approach can improve segmentation accuracy, it is computationally expensive, resulting in low image segmentation efficiency.
[0004] Compared to two-dimensional medical images, three-dimensional medical images provide depth information simultaneously. Multimodal medical images can provide information and features about the lesion and its surrounding areas from multiple levels and angles, greatly aiding in the diagnosis of lesions. However, existing three-dimensional medical image segmentation methods are unable to effectively extract the correlation between the depth dimension and other dimensions. Furthermore, in actual three-dimensional medical images, the distribution of lesion features varies between planes in different directions. This results in poor segmentation quality when applied to three-dimensional medical images. Summary of the Invention
[0005] In response to the defects of the existing technology and the need for improvement, the present invention provides a method for establishing a medical image segmentation model based on harmonic attention and a medical image segmentation method, the purpose of which is to harmonize the feature distribution between different planes of a three-dimensional medical image to reduce the difference in feature distribution in different planes and improve the segmentation accuracy of the three-dimensional medical image.
[0006] To achieve the above-mentioned object, according to one aspect of the present invention, a method for establishing a medical image segmentation model based on harmonic attention is provided, comprising: building a medical image segmentation model and training it using a medical image dataset with annotated segmentation results to complete the establishment of the medical image segmentation model; the medical image is a three-dimensional image;
[0007] The medical image segmentation model includes: a downsampling module, which is an encoder structure in the U-Net network, used to downsample the medical image N times, compress spatial and depth information, and extract semantic information to obtain multi-scale primary semantic feature maps; a Transformer module, which includes a Transformer model, used to obtain global and long-range semantic information of the minimum-scale primary semantic feature map and perform position encoding to obtain an enhanced semantic feature map; and an upsampling module, which is used to upsample the enhanced semantic feature map N times to obtain a segmentation map.
[0008] Among them, N is a positive integer greater than 1; the upsampling module includes N upsampling units connected in sequence; the upsampling unit includes an upsampling layer, a skip connection layer and a harmonic attention layer connected in sequence; the upsampling layer is used to upsample the input feature map; the skip connection layer is used to splice the feature map output by the upsampling layer with the primary semantic feature map of the corresponding scale; the harmonic attention layer is used to respectively extract the feature information in different planes of the feature map output by the skip connection layer, and adjust the feature information of each plane to obey the Gaussian distribution, and then fuse the adjusted feature information in each plane into the attention feature map output by the corresponding upsampling unit.
[0009] Furthermore, the reconciled attention layer includes a hard attention branch;
[0010] The hard attention branch consists of three sub-attention branches and the first fusion layer;
[0011] The three sub-attention branches are used to extract feature information in different planes of the input feature map of the harmonic attention layer and adjust the feature information in the corresponding plane to obey the Gaussian distribution; the sub-attention branches include: an adaptive average pooling layer and a layer normalization layer connected in sequence;
[0012] The first fusion layer uses the matrix broadcast mechanism to perform dot multiplication on the feature information output by the three sub-attention branches to obtain a hard attention map, which is output as the attention feature map output by the corresponding upsampling unit.
[0013] Furthermore, the reconciled attention layer also includes a soft attention branch and a second fusion layer;
[0014] The soft attention branch includes: spatial attention sub-branch, channel attention sub-branch and the third fusion layer; the spatial attention sub-branch extracts the spatial attention map of the input feature map of the harmonic attention layer based on the spatial attention mechanism, and the channel attention sub-branch extracts the channel attention map of the input feature map of the harmonic attention layer. The third fusion layer uses the matrix broadcast mechanism to perform dot multiplication on the spatial attention map and the channel attention map to obtain the soft attention map;
[0015] The second fusion layer is used to perform dot multiplication on the soft attention map and the hard attention map using the matrix broadcast mechanism before outputting the hard attention map, and output the resulting feature map as the attention feature map output by the corresponding upsampling unit.
[0016] Furthermore, the Transformer module includes a plurality of Transformer models cascaded in sequence, and the plurality of Transformer models alternately use an absolute position encoding method and a relative position encoding method to implement position encoding;
[0017] The absolute position encoding method includes: setting a learnable position vector for each image block in the image block sequence input to the Transformer model at the input layer, which is used to identify the position information of the corresponding image block in the entire image block sequence; splicing the position vector of each image block with the vectorized representation of the corresponding image block as the input of the layer after the input layer;
[0018] The relative position encoding method includes: in the multi-head attention layer of the Transformer model, the attention matrix output by the multi-head attention layer is added to the non-learnable relative distance matrix as the input of the layer after the multi-head attention layer;
[0019] Among them, the image block sequence input to the Transformer model is obtained by dividing the minimum-scale primary semantic feature map or the feature map output by the previous level Transformer model; the relative distance matrix records the distance information between any two image blocks in the image block sequence input to the Transformer model.
[0020] Furthermore, the relative distance matrix is calculated as follows:
[0021] Calculate the Euclidean distance between any two image blocks in the image block sequence to obtain a Euclidean distance matrix;
[0022] Calculate the cosine distance between any two image blocks in the image block sequence to obtain a cosine distance matrix;
[0023] Divide the corresponding elements in the Euclidean distance matrix and the pre-distance matrix to obtain the relative distance matrix.
[0024] Furthermore, the medical image segmentation model further includes: a modality fusion module provided before the downsampling module; the modality fusion module includes an input layer, an output layer, and C convolutional layers; C is the number of modalities of the medical image; wherein:
[0025] The input layer is used to extract C modal data from medical images;
[0026] C convolutional layers are used to perform convolution operations on the data of C modalities with a convolution kernel size of 1, an expansion value of 0, and a step size of 1 to obtain C modal features;
[0027] The output layer is used to merge C modal features to obtain a medical image after modality fusion;
[0028] Moreover, the input image of the downsampling module is a medical image after modality fusion.
[0029] Furthermore, the batch normalization layers in the downsampling modules are replaced by group normalization layers.
[0030] According to another aspect of the present invention, a medical image segmentation method based on harmonic attention is provided, comprising: inputting the medical image to be segmented into the medical image segmentation model established by the above-mentioned medical image segmentation model establishment method based on harmonic attention provided by the present invention to obtain a segmentation map.
[0031] According to another aspect of the present invention, a computer program product is provided, comprising a computer program, which, when processed qualitatively, implements the above-mentioned method for establishing a medical image segmentation model based on harmonic attention provided by the present invention, and / or the above-mentioned medical image segmentation method based on harmonic attention provided by the present invention.
[0032] According to another aspect of the present invention, a computer-readable storage medium is provided, comprising a stored computer program; when the computer program is processed qualitatively, the device where the computer-readable storage medium is located is controlled to execute the above-mentioned method for establishing a medical image segmentation model based on harmonic attention provided by the present invention, and / or the above-mentioned medical image segmentation method based on harmonic attention provided by the present invention.
[0033] In general, the above technical solutions conceived by the present invention can achieve the following beneficial effects:
[0034] (1) The medical image segmentation model based on harmonic attention established by the present invention introduces a harmonic attention layer in the upsampling module to harmonize the feature distribution between planes in different directions of the three-dimensional medical image, which plays a transition role after the jump connection. Specifically, after extracting the feature information of three planes in different directions respectively, the feature information of the three planes is made to obey the Gaussian distribution to reduce the difference in feature distribution between different dimensions and stabilize the feature distribution. Then, the feature information of the three planes is fused to make the feature information between the planes complementary, thereby effectively extracting the correlation between the depth dimension and other dimensions, realizing accurate modeling of the entire segmented area, and facilitating the segmentation of discontinuous lesion areas with unclear contours, thereby improving the segmentation accuracy.
[0035] (2) In the preferred embodiment of the present invention, in the medical image segmentation model based on harmonic attention, the harmonic attention layer introduced includes, in addition to the hard attention branch for harmonizing the feature distribution between different planes, a soft attention branch for extracting the spatial attention map and the channel attention map and fusing the two. The soft attention branch can obtain relatively fine-grained pixel features in advance, thereby realizing modeling of each pixel in the segmented area. Combined with the hard attention branch, it can model the lesion area more accurately, further improving the segmentation accuracy.
[0036] (3) In the preferred embodiment of the present invention, in the medical image segmentation model based on coordinated attention, the Transformer module specifically includes multiple Transforer models, and these multiple Transformer models adopt alternating position coding, that is, the absolute position coding method and the relative position coding method are used alternately to realize position coding. The absolute position coding reflects the globality of the semantic position, and the relative position coding reflects the spatial relationship between semantics. The combination of the two can reflect the spatial changes of semantics from both global and local perspectives at the same time, and at the same time use the explicit long-distance dependency relationship established by the global features and the details contained in the local features to clearly segment the boundaries and lesions, thereby further improving the segmentation accuracy.
[0037] (4) In the preferred embodiment of the present invention, in the relative position coding, the relative distance information between the two image blocks is specifically determined by the Euclidean distance and the cosine distance between the image blocks, which includes both the spatial distance information and the direction information between the image blocks. This is consistent with the characteristics that the internal elements of the lesion area have strong spatial connections and the three-dimensional spatial elements in the segmentation task have directionality, which is conducive to further improving the segmentation accuracy.
[0038] (5) In the preferred embodiment of the present invention, a modality fusion module is introduced before the downsampling module, and features of each modality are extracted and then spliced as the input of the downsampling module, thereby realizing input-level fusion and being able to fully extract and fuse multimodal information in medical images.
[0039] (6) In the preferred embodiment of the present invention, based on the selection of the encoder structure in the U-Net network as the downsampling module, the batch normalization layers therein are modified into group normalization layers, which can reduce the dependence of the data on the batch size and weaken the distribution relationship between features, avoid the influence of diverse data distribution on the training process, and further improve the segmentation performance of the established model. BRIEF DESCRIPTION OF THE DRAWINGS
[0040] Figure 1 A schematic diagram of the medical image segmentation model structure provided by an embodiment of the present invention;
[0041] Figure 2 A schematic diagram of a modality fusion module provided in an embodiment of the present invention;
[0042] Figure 3 A schematic diagram of a downsampling module provided in an embodiment of the present invention;
[0043] Figure 4 A schematic diagram of a Transformer module provided in an embodiment of the present invention;
[0044] Figure 5 A schematic diagram of image segmentation provided by an embodiment of the present invention;
[0045] Figure 6 This is a diagram of the structure of the multi-head attention layer in the existing Transformer model;
[0046] Figure 7 Schematic diagram of the multi-layer perceptron (MLP) layer in the existing Transformer model;
[0047] Figure 8 Schematic diagram of an upsampling module provided in an embodiment of the present invention;
[0048] Figure 9 A schematic diagram of the attention coordination layer provided in an embodiment of the present invention;
[0049] Figure 10 A schematic diagram showing different regions in a three-dimensional MRI image of the brain provided by an embodiment of the present invention;
[0050] Figure 11 A schematic diagram comparing segmentation results of different segmentation methods provided in an embodiment of the present invention. DETAILED DESCRIPTION
[0051] In order to make the objectives, technical solutions and advantages of the present invention more clearly understood, the present invention is further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely for the purpose of explaining the present invention and are not intended to limit the present invention. In addition, the technical features involved in the various embodiments of the present invention described below may be combined with each other as long as they do not conflict with each other.
[0052] In the present invention, the terms "first", "second", etc. (if any) in the present invention and the drawings are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence.
[0053] In order to solve the technical problem that the existing three-dimensional medical image segmentation method cannot accurately extract the correlation information between the depth dimension and other dimensions and cannot accurately model the lesion area, resulting in poor segmentation result quality, the present invention has conducted a thorough analysis and research on the causes of this problem. The present invention finds that traditional segmentation methods are mostly based on the model of the encoder-decoder structure to achieve segmentation. The encoder part mainly uses convolution operation to achieve downsampling, and the decoder part mainly uses deconvolution to achieve upsampling, and the feature maps of corresponding scales in the encoder and decoder are fused through jump connections. In clinical practice, the feature distribution of medical images of different patients is different, and the feature distribution of the lesion area in different planes of the same patient's medical image is also different. This makes the conventional encoder-decoder structure model unable to effectively extract the correlation information between the depth dimension and other dimensions, and thus cannot accurately model the lesion area, and cannot accurately segment the lesion area. When the outline of the lesion area is not clear, the quality of the segmentation result is even worse. Based on this discovery, the present invention accordingly proposes to reconcile the feature distribution of the lesion area in different planes in the medical image to reduce the difference in feature distribution between different dimensions, thereby solving the aforementioned technical problems.
[0054] Based on the above concept, the present invention proposes a new segmentation model structure, which introduces a reconciling attention layer after the jump connection in the encoder-decoder structure to reconcile the feature distribution of the lesion area in different planes in the medical image, so that the model can accurately model the lesion area during the upsampling process, thereby obtaining accurate segmentation results. The reconciling attention layer specifically includes a hard branch for extracting global features under the intersection of two axes, that is, the features in the plane formed by any two coordinate axes in the three-dimensional coordinate axis, and further utilizes the Gaussian distribution to stabilize the distribution of features, reduce the impact of feature distribution differences on the segmentation results, and achieve modeling of the entire segmented area. The reconciling attention layer may further include a soft attention branch for extracting relatively fine-grained pixel features in space and channels, mainly modeling each pixel in the segmented area. In order to fully extract the semantic information of the medical image, the present invention introduces a Transformer module based on the Transformer model between the encoder and the decoder. Therefore, in some optional embodiments, the segmentation model mainly includes: a downsampling module, a Transformer module and an upsampling module, and the reconciling attention layer is introduced in the upsampling module.
[0055] Existing three-dimensional medical images are typically composed of multiple modalities. Making full use of modal information and leveraging the diversity and complexity of modalities for modeling can help improve segmentation accuracy. To fully integrate multimodal information, the present invention further proposes introducing a modality fusion module before the encoder-decoder structure to achieve input-level fusion. Therefore, in some optional embodiments of the present invention, the segmentation model mainly includes: a modality fusion module, a downsampling module, a Transformer module, and an upsampling module. Similarly, a harmonic attention layer is introduced in the upsampling module.
[0056] In both model architectures, the structures and functions of the downsampling module, Transformer module, and upsampling module are the same.
[0057] The medical image segmentation model established by the present invention can be applied to the segmentation of any three-dimensional, multimodal medical image. Without loss of generality, the following embodiments are described using the segmentation of brain tumor areas in three-dimensional MRI images of the brain as an example.
[0058] The following are examples.
[0059] Example 1:
[0060] A method for establishing a medical image segmentation model based on harmonic attention includes: building a medical image segmentation model and training it using a medical image dataset with annotated segmentation results to complete the establishment of the medical image segmentation model; the medical image is a three-dimensional image;
[0061] In this embodiment, the medical image segmentation model is used to segment brain tumor regions in 3D brain MRI images. Accordingly, the medical image dataset is a set of 3D brain MRI images with labeled brain tumor regions. It will be readily understood that if the trained model is used for other types of medical image segmentation, the corresponding labeled image dataset can be used for model training.
[0062] like Figure 1 As shown in the figure, the medical image segmentation model as a whole consists of four modules: modality fusion module, downsampling module, Transformer module and upsampling module.
[0063] The modality fusion module uses a three-dimensional convolutional neural network (CNN) to simultaneously extract spatial and depth feature information from different modalities and perform modality fusion to achieve input-level fusion. The downsampling module uses the encoder structure in the U-Net network to downsample the image through convolution operations, compressing spatial and depth information while extracting semantic information. The Transformer module is implemented based on the Transformer model, using its multi-head attention mechanism to obtain global and long-range semantic information and using position encoding to enhance spatial information. The upsampling module uses deconvolution operations to upsample the image and uses a harmonic attention layer to obtain global features and stabilize the distribution of features to obtain a high-quality segmentation map. The specific implementation methods of each module are described in detail below.
[0064] Regarding the modality fusion module based on the input-level fusion strategy, this module serves as the input end of the model and performs modality fusion on the input brain 3D MRI image data. For the input data X∈R C×H×W×D , H×W is its spatial resolution, D is its depth dimension, and C is its modal dimension and feature dimension.
[0065] like Figure 2 As shown, in this module, for the input data, different modal data are extracted respectively. And for each modality, a convolution operation with kernel size (kernel_size) of 1, padding value (padding) of 0, and stride (stride) of 1 is performed. Based on this convolution operation, the features of each modality can be extracted without changing the dimension of the image. Finally, the output result is z ai (i=1, 2……C) and Afterwards, the output results of each modality are combined (Concat) to obtain the output z of the module b and The convolution operation is uniformly referred to as Conv(kernel_size,padding,stride). The operation process of the modal fusion module is as follows:
[0066] z b =Concat(Conv(z a1 ,z a2 ,...,z ac ))
[0067] There are two main methods for fusing multimodal information: (1) dimensional merging, which suffers from low network learning efficiency as the amount of data increases; and (2) additive fusion, which is simply a superposition and cannot fully extract the correlation information between modalities. In contrast, in this embodiment, the modal fusion module can more fully extract the features of each modality and effectively fuse them.
[0068] Regarding the U-Net-based downsampling module, the input of this module is the output z of the modal fusion module b The downsampling module mainly uses the multi-layer convolution structure in the U-Net network to extract the spatial and depth information of the image, and uses the convolution operation to downsample the image, compress the spatial and depth information, and extract the semantic information at the same time.
[0069] The U-Net network has achieved great success in the field of medical image segmentation due to its high accuracy and high performance. This network is a typical encoder-decoder structure. In order to utilize its superior performance in medical image segmentation, this embodiment implements a downsampling module based on the encoder structure in the U-Net network. The module consists of a convolutional layer, an activation function, a normalization layer, and a downsampling layer.
[0070] Considering the differences in brain tumor images among different patients, and to prevent diverse data distributions from affecting subsequent model training, this embodiment, as a preferred implementation, further improves the traditional U-Net network encoder structure. Specifically, the batch normalization layer is modified into a group normalization layer to improve network performance, reduce the data's dependence on batch size, and weaken the distribution relationship between features.
[0071] like Figure 3 The figure shows the structure of the downsampling module based on U-Net. The module as a whole consists of a four-layer convolution stack structure and three downsampling layers. In the convolution stack layer, the core convolution block is UNetConvBlock (hereinafter referred to as UConv), which consists of a series of convolution layers (Conv(3,1,0)), activation functions, and group normalization layers. In this embodiment, the ReLU activation function is referred to as Re, the Sigmoid activation function is referred to as σ, the batch normalization is referred to as BN, the layer normalization is referred to as LN, the group normalization is referred to as GN, and the linear layer is referred to as L. The specific expression of the convolution block UNetConvBlock is as follows:
[0072] UConv(x)=GN(Re(Conv(GN(Re(Conv(x))))))+x
[0073] The downsampling layer (DownSample) uses the Conv(3,2,1) operation to reduce the image width, height, and depth by half, compressing the spatial and depth information while obtaining more semantic information. The relevant expressions are as follows:
[0074] DownSample(x)=Conv(x)
[0075] For the input z of the downsampling module b , first converted to z through Conv(3,1,1) operation c1 , The output of the module is obtained after four layers of convolution stacking structure and three downsampling layers. The specific operation process is as follows:
[0076]
[0077] The final output of the downsampling module is z d , It will continue to be input into the Transformer module for calculation. It should be noted that in actual applications, the number of downsampling can be determined according to the characteristics of the actual medical image to be segmented.
[0078] In this embodiment, the Transformer module is implemented based on the Transformer model, using the multi-head attention mechanism in the Transformer module to obtain global and long-range semantic information, and using position encoding to enhance spatial information. In this embodiment, the overall structure of the Transformer module is as follows: Figure 4 As shown, in order to simultaneously utilize the explicit long-distance dependencies (areas far away from the tumor) established in the global feature space and the details of local features (areas close to the tumor) to clearly segment boundaries and lesions, in this embodiment, the Transformer module contains multiple Transformer models, and unlike the traditional Transformer module, in this embodiment, the position of the Transformer module adopts alternating position encoding, that is, multiple Transformer models alternately use absolute position encoding and relative position encoding. The absolute position encoding reflects the globality of the semantic position, and the relative position encoding reflects the spatial relationship between semantics. The combination of the two can simultaneously reflect the spatial changes of semantics from both global and local perspectives, thereby improving the accuracy of segmentation.
[0079] The input of this module is the output z of the module in the previous sectiond , drawing on ViT's ideas, d The image is subdivided into multiple patches (patches), which are then indexed into a linear sequence that is fed into the Transformer module. Figure 1 shows the overall structure of the Transformer module based on alternating position coding. The following sections will explain the module's composition and the input data computation process, focusing on image patch processing, alternating position coding, and the Transformer model.
[0080] Before inputting the data into the Transformer block, it needs to be divided into blocks and converted into a linear sequence of image blocks. Convert it into a 2D image block Where (P, P, P) is the spatial resolution of the image block, and N = HWD / P 3 is the number of image blocks. The same hidden layer size H is used in all Transformer blocks. d Therefore, after the data is divided into blocks, a linear layer is needed to convert the feature dimension into H d . Take the input z of this module d For example, the block operation process is as follows Figure 5 The corresponding expressions are as follows:
[0081]
[0082] Regarding alternating positional encoding, in natural language processing tasks, the Transformer adds positional encoding based on the order of the sequence to build a sequence model, considering that the attention mechanism cannot utilize positional information in the sequence. In computer vision tasks, positional encoding can also be used to add positional information to the sequence of image patches, thereby maximizing the semantic relationship between the image patches' context.
[0083] Positional encodings are broadly categorized into two types: absolute positional encoding and relative positional encoding. Absolute positional encoding encodes the input sequence from 1 to the maximum sequence length, with each position assigned a separate positional encoding vector. The positional vector is then combined with the input sequence to allow the model to capture positional information. Relative positional encoding encodes the relative distances between input elements, allowing the model to capture positional information between elements.
[0084] The Transformer modules of this embodiment alternately use learnable absolute position encoding (APE) and non-learnable relative position encoding (RPE). The former is applied to the first and third Transformer modules in the model, and the latter is applied to the second and fourth Transformer modules.
[0085] APE has the same dimension as the input data, and each element can be learned through the back propagation mechanism of the model. At the input layer of the Transformer model, APE sets a learnable position vector for each image block in the image block sequence input to the Transformer model, which is used to identify the position information of the corresponding image block in the entire image block sequence. APE adds position encoding in an additive manner. Specifically, the position vector of each image block is spliced with the vectorized representation of the corresponding image block as the input of the layer after the input layer. The operation process of absolute position encoding is as follows:
[0086]
[0087] RPE also adds position encoding in an additive manner. Specifically, in the multi-head attention layer of the Transformer model, the attention matrix output by the multi-head attention layer is added to the non-learnable relative distance matrix as the input of the layer after the multi-head attention layer; considering the strong spatial connectivity between elements within the brain tumor area, each element does not need to be learned through the backpropagation mechanism, and the Euclidean distance is used to obtain its relative spatial distance; considering the directionality between the three-dimensional spatial elements in the segmentation task, the cosine distance is used to obtain the directional information between the elements. In this embodiment, RPE is determined by the Euclidean distance and cosine distance between image blocks. The specific calculation process is as follows:
[0088] Assume that the three-dimensional absolute position of any two image blocks is p i (x pi ,y pi ,z pi ), p j (x pj ,y pj ,z pj ), first calculate the Euclidean distance between the two:
[0089]
[0090] Then calculate the cosine distance between the two:
[0091]
[0092] By calculating the Euclidean distance and cosine distance between any two image blocks, the Euclidean distance matrix can be obtained. and cosine distance matrix Finally, the two matrices are divided element by element to obtain the final relative distance matrix (The relative distance between an element and its own element is 0.) The timing of adding relative position encoding is different from that of absolute position encoding. Since each element has a relative position encoding with respect to other elements, the position encoding is added to the attention matrix and becomes part of the attention. Let x be the input. The specific calculation process is as follows:
[0093]
[0094] Regarding the Transformer model, this module uses the Transformer encoder structure, which consists of two normalization layers, a multi-head attention layer, and an MLP layer. The following will introduce the components of this module in turn.
[0095] The multi-head attention layer consists of multiple self-attention blocks, such as Figure 6 As shown. The attention function can be described as mapping a query (Query) and a set of key (Key) value (Value) pairs to an output result, where the query, key, value and output result are all vectors. The output result is the weighted sum of the Value, where the weight assigned to each Value is calculated by the compatibility function between the Query and the corresponding Key. In the multi-head attention layer, the compatibility function uses a scaled point attention function, and the input consists of a dimension of d k Query and Key, and dimension d v Calculate the dot product of all Keys and Query, and divide each Key by And apply a Softmax function to obtain the weights of these Values.
[0096] In actual calculations, the attention function is calculated on a set of queries at the same time, and all queries are placed in the matrix Q. All keys and values are placed in the matrices K and V, and the output matrix is calculated as follows:
[0097]
[0098] The attention function used by the multi-head attention layer is different from that of the single-head layer. First, it uses a linear layer to triple the feature dimension of the input data, and then takes the query, key, and value in sequence. It then divides the feature dimension according to the number of attention heads, performs single-head attention operations on each, merges the different feature dimensions, and finally uses a linear layer to output the result. The specific calculation method of the multi-head attention layer is shown in the following formula:
[0099] head i =Attention(QWi i Q ,KW i K,VW i V )
[0100] MultiHead(Q,K,V)=Concat(head1,…,head h )W O
[0101] in and are all learnable parameters of the linear layer, d model is the feature dimension of the input data, h is the number of attention heads, d q =d k =d v =d model / h. All attention functions in this module use self-attention, so Q, K, and V are all themselves. Hereinafter, MultiHead(Q,K,V) will be referred to as MSA(x).
[0102] The MLP layer consists of two linear layers, such as Figure 7 As shown in the figure, the first linear layer uses the GeLU activation function and the Dropout mechanism, while the second linear layer only uses the Dropout mechanism. Therefore, the operation flow of the MLP layer can be expressed by the following expression:
[0103] MLP(x)=L(ReLU(L(x)))
[0104] Based on the above operation mechanism, input z d1 As an example, the operation and change process of the data are described in detail. Since the operation process of the Transformer model using absolute position encoding and relative position encoding is different, they are described separately here. First, the Transformer model using absolute position encoding (hereinafter referred to as APET) is introduced. The operation process is as follows:
[0105] x′=MSA(LN(APE(x)))+x
[0106] APET(x)=MLP(LN(x′))+x′
[0107] In the Transformer model using relative position encoding (hereinafter referred to as RPET), the position encoding is added during the attention calculation of the multi-head attention, as shown in the following formula. The remaining steps are similar to APET, and the calculation process is as follows:
[0108]
[0109] x′=MSA(LN(x))+x
[0110] RPET(x)=MLP(LN(x′))+x′
[0111] The number of Transformer models in the Transformer module of this embodiment is 4. d1 In the Transformer module with input alternating position encoding, the data operation and transformation process is shown in the following formula, where z e This is the final output of the module.
[0112]
[0113] In this embodiment, the upsampling module uses skip connections and a harmonized attention layer. The upsampling module is similar to the downsampling module based on U-Net, and both consist of a four-layer network structure, such as Figure 8 As shown in the figure, the input of each layer is jump-connected to the output of the corresponding layer of the downsampling module, that is, a Concat operation is performed. Then, the harmonized attention block is used to extract the spatial attention and channel attention of the output result after the Concat operation, and the Gaussian distribution is used to stabilize the distribution of the attention weights. Finally, deconvolution is used for upsampling to gradually restore the image to its original size.
[0114] In this embodiment, the harmonized attention layer (HA) consists of two branches, such as Figure 9 As shown in Figure 2, the attention matrix is divided into soft attention and hard attention. The soft attention branch (A) extracts fine-grained pixel features in space and channels, primarily modeling each pixel in the segmented region. The hard attention branch (H) extracts global features under the intersection of two axes, using a Gaussian distribution to stabilize the distribution of features and primarily modeling the entire segmented region. Finally, the soft and hard attention matrices are dot-multiplied using a matrix broadcast mechanism to obtain the final attention matrix.
[0115] The reconciled attention layer plays a transitional role after the skip connection: it globally utilizes the spatial and channel information of low-level and high-level image features while stabilizing the distribution of both features, facilitating the extraction of higher-level features. The soft attention branch is divided into spatial attention (SA) and channel attention (CA). The goal of soft attention learning is to generate a saliency weight map of the same size as the input. Considering the independence between spatial (i.e., pixel-to-pixel) and channel (i.e., scale-to-scale) attention, the attention map A output by the soft attention branch can be expressed as follows:
[0116] A=SA⊙CA
[0117] in, The two codes are the spatial attention weight map and the channel attention weight map respectively, and ⊙ represents the dot product of the matrix-based broadcast mechanism.
[0118] In the soft attention branch, the spatial attention weight map and the channel attention weight map are extracted by the spatial attention sub-branch and the channel attention sub-branch respectively, where:
[0119] The spatial attention branch is modeled using a four-layer subnetwork: a global cross-channel average pooling layer (GCCAP), an HConv(3,1,1) convolutional layer, a trilinear interpolation upsampling layer, and an HConv(1,1,0) convolutional layer. The GCCAP layer implements the sharing of spatial attention across channels, and a convolutional layer is subsequently used to adaptively learn spatial attention. The operation process can be expressed as follows:
[0120]
[0121] In the spatial attention sub-branch, the data operation and flow process can be expressed as follows:
[0122] HConv(x)=BN(Re(Conv(x)))
[0123]
[0124] The channel attention sub-branch is modeled using a three-layer sub-network: a global pooling layer (GAP) and two HConv(1,1,0) convolutional layers. The global pooling layer aggregates feature information distributed in three-dimensional space into channel information. This operation provides relatively complete information for subsequent modeling of inter-channel dependencies. Its operation process can be expressed as follows:
[0125]
[0126] In the second HConv layer, the number of channels is reduced in dimension to aggregate channel information while reducing the amount of computation.
[0127] In the channel attention sub-branch, the operation and flow of data can be expressed as follows:
[0128]
[0129] Finally, the matrix broadcast mechanism is used to multiply the spatial attention map and the channel attention map to form the soft attention map A(x), which is expressed as follows:
[0130] A(x)=SA(x)⊙CA(x)
[0131] The hard attention branch consists of three non-learnable adaptive average pooling (AAP) layers and three layer normalization layers. The adaptive pooling layer extracts feature information from the biaxial cross plane, while the layer normalization layer ensures that the feature information from the biaxial cross plane (i.e., each directional plane in a 3D medical image) follows a Gaussian distribution, reducing the difference in feature distribution between different dimensions. Considering the correlation and dependency between the three dimensions in a 3D image, three AAP layers are used here to extract features from the WD, HD, and HW planes, respectively. Taking the HW plane as an example, the specific calculation method is as follows:
[0132]
[0133] After calculating the cross features of the three planes, we use the matrix broadcast mechanism to perform dot multiplication, share the coaxial information of the three planes, and complement the feature information between the planes. In the hard attention branch, the data operation and flow process is as follows:
[0134]
[0135] H(x)=x′
[0136] Finally, the soft attention map and the hard attention map are dot-multiplied using the matrix broadcast mechanism, and the dot-multiplication result is normalized using the Sigmoid function to obtain the final attention weight matrix. The overall calculation process is as follows:
[0137]
[0138] The model established in this embodiment is an encoder-decoder structure as a whole. The encoder part uses convolutional layers for step-by-step downsampling, and the decoder part uses deconvolution TConv(2,0,2) for step-by-step upsampling. The high-frequency information (i.e., image details) and low-frequency information (i.e., image contours) in the original input image will be gradually restored. Therefore, the low-resolution feature map will eventually be mapped to a pixel-level segmentation map.
[0139] To further compensate for the information lost during downsampling in the encoding phase, the Concat operation is used between the encoder and decoder in the network to fuse the feature maps at corresponding positions in the two processes in the channel dimension. This allows the decoder to retain more high-resolution detail information contained in high-level feature maps during upsampling, thereby more perfectly restoring the detail information in the original image and improving segmentation accuracy. The operation process of the upsampling block (UpSample) is as follows:
[0140] UpSample(x)=TConv(Conv(x))
[0141] During the downsampling process of the encoder, the output results of each layer are: Combined with the output results of the Transformer module based on alternating position encoding in the previous section Compared with the structure of the upsampling module, the data operation and transformation process in the upsampling module based on skip connection and harmonic attention is expressed as follows:
[0142]
[0143] Perform an upsampling, a skip connection, and calculate the harmonic attention at each layer, and finally get the output of the upsampling module based on skip connection and harmonic attention as z g , The calculation process is as follows:
[0144]
[0145] Finally, the output result z of the upsampling module g , through a Conv(1,0,0) operation and Softmax function, the final segmentation result z is obtained h , the operation process is as follows:
[0146] z h =Softmax(Conv(z g ))
[0147] In general, this embodiment uses an input-level fusion strategy in the modal fusion module to integrate multimodal spatial and depth information; in the downsampling module, the translation invariance and local receptive field of convolution are used to compress spatial and depth information while extracting local semantic information; in the Transformer module, a multi-head self-attention mechanism is used to obtain global features and long-range information dependencies, and alternating relative position encoding and absolute position encoding are added to simultaneously strengthen semantics and spatial information between semantics; in the upsampling module, jump connections and harmonic attention are used to obtain global low-level and high-level features and stabilize the distribution of both features, ultimately effectively improving the accuracy of medical image segmentation.
[0148] As a preferred implementation, the model training process of this embodiment adopts additive cross entropy and Dice loss function to solve the problem of imbalance between positive and negative sample classes in medical images; and adopts warm-up-cosine annealing learning rate to accelerate the convergence speed of the model and prevent overfitting. Figure 1 shows a detailed schematic diagram of the medical image segmentation model after training, including module name, structure name, structure details, output results and output size. In the table, all blocks with * use residual structure; in all UEnBlocks, Conv is Conv(3,1,1) operation, in all DeBlock blocks, all Conv is Conv(3,1,1) operation, in all UpSample blocks, Conv is Conv(1,0,1) operation, all TConv is TConv(2,0,2) operation, and in the Segmentation Head block, Conv is Conv(1,0,1) operation.
[0149] Table 1 Medical image segmentation model details
[0150]
[0151]
[0152] To verify the segmentation performance of the medical image segmentation model established in this example (abbreviated as HATransUNet), we further trained the model using 1,600 training samples from the BraTS2021 training set and tested the model using 400 test set samples. We calculated different evaluation metrics to evaluate the model's performance. We trained and tested the 3D U-Net, 3D V-Net, and Attention U-Net models using the same training and test sets, and compared the test results with those of this example.
[0153] The BRATS2021 dataset consists of images from multiple MRI scans of 2,000 patients, divided into a training set of 1,600 cases and a test set of 400 cases in an 8:2 ratio. The sub-regions considered for evaluation in BraTS2021 are tumor enhancement (ET), tumor core (TC), and whole tumor (WT). ET includes tumor enhancement, and TC includes tumor enhancement, tumor necrosis, and tumor non-enhancement. WT includes all areas except the background area, describing the complete range of the tumor.
[0154] The patient's scan images contain four modalities of co-registered MRI images: T1-weighted (T1), T1-weighted and contrast-enhanced (T1ce), T2-weighted (T2) and FLAIR, with an image size of 240×240×155. All data have been skull-stripped, linearly aligned, and spatially resampled to 1mm×1mm×1mm. The segmentation labels of the training set are divided into four categories, namely label 1 (necrotic area and non-enhanced area), label 2 (edema area), label 3 (enhanced area) and label 0 (background). The four modalities share the segmentation labels. Figure 10 shown.
[0155] Two evaluation metrics are used to evaluate the performance of the model: the Dice coefficient and the Hausdorff distance. Table 2 shows the comparison results of the proposed model method with the other three model methods on the test set.
[0156] Table 2 Segmentation performance indicators of different models on the BraTS2021 test set
[0157]
[0158]
[0159] Three random patients from the BraTS2021 test set were selected for demonstration, such as Figure 11 As shown in the figure. The first row shows the 32nd coronal layer of patient BraTS2021_00014, the second row shows the 82nd coronal layer of patient BraTS2021_00324, and the third row shows the 111th coronal layer of patient BraTS2021_00768. The first column shows the segmentation results of 3D U-Net, the second column shows the segmentation results of 3D V-Net, the third column shows the segmentation results of Attention U-Net, the fourth column shows the segmentation results of HATransUNet, and the fifth column shows the distribution of the true labels.
[0160] In order to more clearly show the segmentation performance of each model inside the tumor region and at the edge of the region, Figure 11 The interior and edge of some areas of each patient were magnified and displayed, and different areas were represented in the same way. Figure 10The true labels reflect problems such as similar features, blurred boundaries, and discontinuous targets between tumor sub-regions. The locally enlarged images of the segmentation results of the 3D U-Net, 3D V-Net, and AttentionU-Net models all show problems such as discontinuous segmentation areas and unclear segmentation boundaries, especially between the edema area and the background, and between the enhanced area and the non-enhanced area. The above evaluation index results and segmentation result visualization results show that the HATransUNet model proposed in this embodiment can effectively solve the above problems.
[0161] Example 2:
[0162] A method for establishing a medical image segmentation model based on harmonious attention. This embodiment is similar to the above-mentioned embodiment 1, except that, in this embodiment, the medical image segmentation model does not include a modality fusion module, but only includes a downsampling module, a Transformer module and an upsampling module.
[0163] In this embodiment, the specific implementation of each module in the model and the specific implementation of the remaining steps can refer to the description in the above embodiment 1.
[0164] Example 3:
[0165] A medical image segmentation method based on harmonic attention comprises: inputting the medical image to be segmented into the medical image segmentation model established by the medical image segmentation model establishment method based on harmonic attention provided by the above-mentioned embodiment 1 or 2 to obtain a segmentation map.
[0166] Example 4:
[0167] A computer program product includes a computer program, which, when processed qualitatively, implements the method for establishing a medical image segmentation model based on harmonious attention provided in the above-mentioned embodiment 1 or 2, and / or the medical image segmentation method based on harmonious attention provided in the above-mentioned embodiment 3.
[0168] Example 5:
[0169] A computer-readable storage medium includes a stored computer program; when the computer program is processed, the device where the computer-readable storage medium is located is controlled to execute the medical image segmentation model establishment method based on harmonic attention provided by the above-mentioned embodiment 1 or 2, and / or the medical image segmentation method based on harmonic attention provided by the above-mentioned embodiment 3.
[0170] It will be easily understood by those skilled in the art that the above description is merely a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, and improvements made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.
Claims
1. A method for establishing a medical image segmentation model based on harmonic attention, characterized in that: include: Building the medical image segmentation model and training it using a medical image dataset with annotated segmentation results to complete the establishment of the medical image segmentation model; The medical image is a three-dimensional image; The medical image segmentation model includes: a downsampling module, which is an encoder structure in a U-Net network, for downsampling the medical image N times, compressing spatial and depth information, and extracting semantic information to obtain a multi-scale primary semantic feature map; a Transformer module, which includes a Transformer model, for obtaining global and long-range semantic information of the minimum-scale primary semantic feature map and performing position encoding to obtain an enhanced semantic feature map; and an upsampling module, which is used to upsample the enhanced semantic feature map N times to obtain a segmentation map; Among them, N is a positive integer greater than 1; the upsampling module includes N upsampling units connected in sequence; the upsampling unit includes an upsampling layer, a jump connection layer and a harmonic attention layer connected in sequence; the upsampling layer is used to upsample the input feature map; the jump connection layer is used to splice the feature map output by the upsampling layer with the primary semantic feature map of the corresponding scale; the harmonic attention layer is used to respectively extract the feature information in different planes of the feature map output by the jump connection layer, and adjust the feature information of each plane to obey the Gaussian distribution, and then fuse the adjusted feature information in each plane into the attention feature map output by the corresponding upsampling unit.
2. The method for establishing a medical image segmentation model based on harmonious attention according to claim 1, wherein: The reconciled attention layer includes a hard attention branch; The hard attention branch includes: three sub-attention branches and a first fusion layer; The three sub-attention branches are respectively used to extract feature information in different planes of the input feature map of the harmonic attention layer, and adjust the feature information in the corresponding plane to obey the Gaussian distribution; the sub-attention branches include: an adaptive average pooling layer and a layer normalization layer connected in sequence; The first fusion layer uses the matrix broadcast mechanism to perform dot multiplication on the feature information output by the three sub-attention branches to obtain a hard attention map, which is output as the attention feature map output by the corresponding upsampling unit.
3. The method for establishing a medical image segmentation model based on harmonious attention according to claim 2, wherein: The reconciled attention layer further includes a soft attention branch and a second fusion layer; The soft attention branch includes: a spatial attention sub-branch, a channel attention sub-branch and a third fusion layer; the spatial attention sub-branch extracts the spatial attention map of the input feature map of the harmonic attention layer based on the spatial attention mechanism, the channel attention sub-branch extracts the channel attention map of the input feature map of the harmonic attention layer, and the third fusion layer uses the matrix broadcast mechanism to perform a dot product on the spatial attention map and the channel attention map to obtain a soft attention map; The second fusion layer is used to perform a dot product on the soft attention map and the hard attention map using a matrix broadcast mechanism before outputting the hard attention map, and output the obtained feature map as the attention feature map output by the corresponding upsampling unit.
4. The method for establishing a medical image segmentation model based on harmonic attention according to any one of claims 1 to 3, wherein: The Transformer module includes a plurality of Transformer models cascaded in sequence, and the plurality of Transformer models alternately use an absolute position encoding method and a relative position encoding method to implement position encoding; The absolute position encoding method includes: setting a learnable position vector for each image block in the image block sequence input to the Transformer model at the input layer of the Transformer model, which is used to identify the position information of the corresponding image block in the entire image block sequence; splicing the position vector of each image block with the vectorized representation of the corresponding image block as the input of the layer after the input layer; The relative position encoding method includes: in the multi-head attention layer of the Transformer model, adding the attention matrix output by the multi-head attention layer to the non-learnable relative distance matrix as the input of the layer after the multi-head attention layer; The image block sequence input to the Transformer model is obtained by dividing the minimum-scale primary semantic feature map or the feature map output by the previous-level Transformer model; the relative distance matrix records the distance information between any two image blocks in the image block sequence input to the Transformer model.
5. The method for establishing a medical image segmentation model based on harmonious attention according to claim 4, wherein: The relative distance matrix is calculated as follows: Calculate the Euclidean distance between any two image blocks in the image block sequence to obtain a Euclidean distance matrix; Calculate the cosine distance between any two image blocks in the image block sequence to obtain a cosine distance matrix; The Euclidean distance matrix is divided by corresponding elements in the pre-distance matrix to obtain the relative distance matrix.
6. The method for establishing a medical image segmentation model based on harmonic attention according to any one of claims 1 to 3, wherein: The medical image segmentation model further includes: a modality fusion module provided before the downsampling module; the modality fusion module includes an input layer, an output layer, and C convolutional layers; C is the number of modalities of the medical image; wherein: The input layer is used to extract C modal data from medical images; C convolutional layers are used to perform convolution operations on the data of C modalities with a convolution kernel size of 1, an expansion value of 0, and a step size of 1 to obtain C modal features; The output layer is used to merge C modal features to obtain a medical image after modality fusion; Furthermore, the input image of the downsampling module is the medical image after the modality fusion.
7. The method for establishing a medical image segmentation model based on harmonic attention according to any one of claims 1 to 3, wherein: The batch normalization layers in the downsampling modules are replaced by group normalization layers.
8. A medical image segmentation method based on harmonic attention, characterized in that: include: The medical image to be segmented is input into the medical image segmentation model established by the method for establishing a medical image segmentation model based on harmonic attention according to any one of claims 1 to 7 to obtain a segmentation map.
9. A computer program product, characterized in that It includes a computer program, which, when processed qualitatively, implements the method for establishing a medical image segmentation model based on harmonious attention as described in any one of claims 1 to 7, and / or the medical image segmentation method based on harmonious attention as described in claim 8.
10. A computer-readable storage medium, characterized in that Including a stored computer program; when the computer program is processed, it controls the device where the computer-readable storage medium is located to execute the method for establishing a medical image segmentation model based on harmonic attention according to any one of claims 1 to 7, and / or the method for medical image segmentation based on harmonic attention according to claim 8.
Citation Information
Patent Citations
Multi-scale bimodal text image generation method based on generative adversarial network
CN113421314A
Ultrasonic image processing method and system based on cloud computing
CN117974738A