Image segmentation method based on block-level self-attention and block correlation mechanism

Abstract

The application provides an image segmentation method based on block-level self-attention and block correlation mechanism, comprising the following steps: 1, building a ResNet101 backbone network; 2, dividing the output features of the backbone network into blocks, and using 1*1 convolution to generate a query matrix (Q), a key matrix (K) and a value matrix (V) for each feature block, and using a block correlation mechanism (BCM) to calculate the correlation of the Q matrix of each feature block and the global features; 3, building a self-attention network with relative position offset, embedding the BCM into the self-attention calculation process of each feature block, using the output M matrix of the BCM and the K matrix and the V matrix for self-attention calculation, and introducing a relative position offset for the M, K and V matrices of each feature block to jointly form a block-level self-attention (BLSA) module; 4, splicing the output of each feature block in the spatial dimension and sending it into a segmentation head for segmentation, and performing experimental comparison.

Description

[0001] This application claims domestic priority to the invention application filed on June 28, 2022, application number 202210739674.8, entitled "A method for segmenting esophageal precancerous lesion regions based on block-level self-attention and block correlation mechanisms", the entire contents of which are incorporated herein by reference. Technical Field

[0002] This invention relates to the fields of deep learning and medical image segmentation, and more specifically, to an image segmentation method based on block-level self-attention and block correlation mechanisms. Background Technology

[0003] Since the introduction of Fully Convolutional Networks (FCNs), convolutional neural networks have become the primary means of segmentation tasks, achieving high accuracy in natural image segmentation. Image segmentation technology has been rapidly applied to various industries, such as autonomous driving. However, its application in medical image segmentation is relatively limited, and research on the segmentation of esophageal precancerous lesions started even later. Because the inter-class features of the four types of precancerous lesions—esophageal inflammation, low-grade lesions, high-grade lesions, and early-stage cancer—are not obvious and have small appearance differences, the network cannot extract discriminative features for each type of lesion, resulting in relatively low segmentation accuracy.

[0004] To effectively improve the accuracy of a network, methods such as using a larger receptive field and acquiring global information can be employed. However, these methods introduce a large amount of computation during implementation, making network training more difficult and hindering the application of segmentation techniques for esophageal precancerous lesions in real-world scenarios.

[0005] U-Net networks are widely used in medical image segmentation. They increase the receptive field of deeper layers through stacked convolutions and pooling operations, achieving high accuracy in fundus image segmentation. However, this network loses positional information during downsampling, leading to the development of many U-Net variants, such as U-Net++ and Attention UNet (2018), to compensate for its shortcomings. While these encoder-decoder-based networks can achieve high segmentation accuracy, they involve a large number of training parameters and computational overhead, making them challenging in practical applications. Similarly, some networks use feature pyramids to obtain multi-scale features, thereby improving segmentation accuracy. The Atrous Spatial Pyramid Pooling (ASPP) module proposed by Chen et al. and the pyramid pooling model proposed in PSPNet are the most representative. These networks increase the receptive field of deeper layers through dilated convolutions with different dilation rates and pooling at different scales. However, limited by the dilation rate and pooling kernel size, the receptive field obtained from deep features is also limited, and high accuracy cannot yet be achieved on esophageal cancer datasets. Self-attention has a natural advantage in acquiring global and long-range dependency information, but it introduces a significant amount of computation. CCNet and Medt are improvements on self-attention. The former uses pixels on the cross-shaped path to obtain global information, but this results in the loss of the positional information of these pixels. The latter introduces positional information into the attention process, but the introduced positional embedding is multiplied by the Q, K, and V matrices required for self-attention, which also adds considerable computational cost to the network. Based on the current state of research, segmentation networks on esophageal cancer datasets currently struggle to achieve relatively high segmentation accuracy with minimal computation. Summary of the Invention

[0006] In view of this, the present invention provides an image segmentation method based on block-level self-attention and block correlation mechanisms.

[0007] The image segmentation method based on block-level self-attention and block correlation mechanisms provided by this invention includes: processing an endoscope image using a backbone network to obtain a first output feature; dividing the first output feature into blocks to obtain multiple feature blocks; processing each feature block using three 1×1 convolutions to obtain a query matrix, a first key matrix, and a first value matrix; processing the query matrix using a block correlation mechanism to obtain a first block matrix; performing self-attention calculation on the first block matrix, the first key matrix, and the first value matrix based on their relative position offsets to obtain a second output feature; concatenating the multiple second output features corresponding to each of the multiple feature blocks in a spatial dimension to obtain a third output feature; and processing the third output feature using a segmentation head to obtain the segmentation result of the endoscope image.

[0008] According to an embodiment of the present invention, the above-described processing of the query matrix using a block correlation mechanism to obtain a first block matrix includes: performing convolution processing on the first output feature to obtain a block correlation matrix; and obtaining the first block matrix based on the block correlation matrix and the query matrix.

[0009] According to an embodiment of the present invention, obtaining the first block matrix based on the block correlation matrix and the query matrix includes: using the block correlation matrix R to divide the feature blocks and the first output feature X = {X 11 , ..., X ij , ..., X IJ The association is performed, where i∈{1,…,I} and j∈{1,…,J} represent the feature blocks in the i-th row and j-th column of the first output feature X, respectively. H and W are the height and width of the first output feature X mentioned above. and Here, C represents the height and width of the feature block, and C is the number of feature channels. The formula for the first block matrix M is as follows:

[0010]

[0011] In this formula, It is the aforementioned feature block X in the i-th row and j-th column. ij The query matrix generated by 1×1 convolution Let R represent a 1×1 convolution. The block correlation matrix R = Conv(X) means that the block correlation matrix R is obtained by convolution from the first output feature X.

[0012] According to an embodiment of the present invention, the above-described self-attention calculation on the first block matrix, the first key matrix, and the first value matrix based on their respective relative position offsets to obtain a second output feature includes: processing the first block matrix based on the relative position offset between the first block matrix and the first key matrix to obtain a second block matrix; processing the first key matrix based on the relative position offset between the first block matrix and the first key matrix to obtain a second key matrix; processing the first value matrix based on the relative position offset between the first value matrix to obtain a second value matrix; and performing self-attention calculation on the second block matrix, the second key matrix, and the second value matrix to obtain the second output feature.

[0013] According to an embodiment of the present invention, the backbone network is ResNet101. Attached Figure Description

[0014] The above and other objects, features and advantages of the present invention will become more apparent from the following description of embodiments of the invention with reference to the accompanying drawings, in which:

[0015] Figure 1 The diagram illustrates the overall flow of a deep learning model according to an embodiment of the present invention.

[0016] Figure 2 The diagram illustrates the overall network structure of a deep learning model according to an embodiment of the present invention.

[0017] Figure 3 A schematic diagram of a self-attention network structure with relative position offset is shown according to an embodiment of the present invention.

[0018] Figure 4 The diagram illustrates the network structure of a block-level self-attention module according to an embodiment of the present invention. Detailed Implementation

[0019] Hereinafter, embodiments of the present invention will be described with reference to the accompanying drawings. However, it should be understood that these descriptions are exemplary only and are not intended to limit the scope of the invention. In the following detailed description, numerous specific details are set forth to provide a thorough understanding of the embodiments of the invention for ease of explanation. However, it will be apparent that one or more embodiments may be practiced without these specific details. Furthermore, descriptions of well-known structures and techniques are omitted in the following description to avoid unnecessarily obscuring the concept of the invention.

[0020] The terminology used herein is for the purpose of describing particular embodiments only and is not intended to limit the invention. The terms “comprising,” “including,” etc., as used herein indicate the presence of the stated features, steps, operations, and / or components, but do not exclude the presence or addition of one or more other features, steps, operations, or components.

[0021] All terms used herein (including technical and scientific terms) have the meanings commonly understood by those skilled in the art, unless otherwise defined. It should be noted that the terms used herein are to be interpreted in a manner consistent with the context of this specification, and not in an idealized or overly rigid way.

[0022] When using expressions such as "at least one of A, B, and C", they should generally be interpreted in accordance with the meaning that is commonly understood by a person skilled in the art (e.g., "a system having at least one of A, B, and C" should include, but is not limited to, a system having A alone, a system having B alone, a system having C alone, a system having A and B, a system having A and C, a system having B and C, and / or a system having A, B, and C, etc.).

[0023] Self-attention is commonly used to obtain long-range dependency and global information. To reduce the number of network parameters and computational cost, this method improves self-attention to reduce computational cost while still obtaining long-range dependency information. Features are divided into blocks, and self-attention is applied to each block to reduce the number of network parameters and computational cost. To ensure that each feature block contains information related to the entire feature map, this method proposes a block correlation mechanism, modeling the relationship between each feature block and the entire input feature map. This allows each feature block to obtain long-range dependency information after self-attention computation. Furthermore, to compensate for the loss of positional information due to block division, relative position offset is introduced into the self-attention mechanism to improve the accuracy of segmentation boundaries.

[0024] The technical solution of this invention is as follows: An image segmentation method based on block-level self-attention and block correlation mechanisms, mainly comprising the following steps:

[0025] Step 1: Build the ResNet101 backbone network.

[0026] Step 2: Divide the output features of the backbone network into blocks, and use 1×1 convolution to generate a query matrix (Q), a key matrix (K), and a value matrix (V) for each feature block. Use the Block Correlation Mechanism (BCM) to calculate the correlation between the Q matrix of each feature block and the global features.

[0027] Step 3: Construct a self-attention network with relative position offset. Embed the BCM into the self-attention calculation process of each feature block. Use the output M matrix, K matrix and V matrix of the BCM to perform self-attention calculation, and introduce relative position offset for the M, K and V matrices of each feature block to form a block-level self-attention (BLSA) module.

[0028] Step 4: The output of each feature block is concatenated in the spatial dimension and then fed into the segmentation head for segmentation, and experimental comparisons are conducted.

[0029] In the embodiments of this invention, the collection, updating, analysis, processing, use, transmission, provision, disclosure, and storage of data (e.g., including but not limited to user personal information) comply with relevant laws and regulations, are used for legitimate purposes, and do not violate public order and good morals. In particular, necessary measures have been taken to prevent unauthorized access to user personal information data and to safeguard user personal information security, network security, and national security.

[0030] In the embodiments of the present invention, the user's authorization or consent is obtained before acquiring or collecting the user's personal information.

[0031] According to embodiments of the present invention, an image segmentation method based on block-level self-attention and block correlation mechanisms may include the following operations:

[0032] The endoscopic image is processed using a backbone network to obtain the first output feature. The first output feature is then divided into blocks to obtain multiple feature blocks. For each feature block, three 1×1 convolutions are used to process the feature block to obtain a query matrix, a first key matrix, and a first value matrix. The query matrix is ​​processed using a block correlation mechanism to obtain the first block matrix. Based on the relative position offsets of the first block matrix, the first key matrix, and the first value matrix, self-attention is calculated on the first block matrix, the first key matrix, and the first value matrix to obtain the second output feature. The multiple second output features corresponding to the multiple feature blocks are concatenated in the spatial dimension to obtain the third output feature. Finally, the third output feature is processed using a segmentation head to obtain the segmentation result of the endoscopic image.

[0033] According to an embodiment of the present invention, processing a query matrix using a block correlation mechanism to obtain a first block matrix includes: performing convolution processing on a first output feature to obtain a block correlation matrix; and obtaining the first block matrix based on the block correlation matrix and the query matrix.

[0034] According to an embodiment of the present invention, a first block matrix is ​​obtained based on a block correlation matrix and a query matrix, including: using the block correlation matrix R to divide feature blocks and the first output feature X = {X11 , ..., X ij , ..., X IJ The association is performed, where i∈{1,…,P} and j∈{1,…,J} represent the feature blocks in the i-th row and j-th column of the first output feature X, respectively. H and W are the height and width of the first output feature X. and Here, C represents the height and width of the feature block, and C is the number of feature channels. The formula for the first block matrix M is as follows:

[0035]

[0036] In this formula, It is the feature block X in the i-th row and j-th column. ij The query matrix generated by 1×1 convolution Let R represent a 1×1 convolution, and let R = Conv(X) represent the block correlation matrix obtained by convolution of the first output feature X.

[0037] According to an embodiment of the present invention, a second output feature is obtained by performing self-attention calculation on the first block matrix, the first key matrix, and the first value matrix based on their respective relative position offsets, including: processing the first block matrix based on the relative position offset between the first block matrix and the first key matrix to obtain a second block matrix; processing the first key matrix based on the relative position offset between the first block matrix and the first key matrix to obtain a second key matrix; processing the first value matrix based on the relative position offset of the first value matrix to obtain a second value matrix; and performing self-attention calculation on the second block matrix, the second key matrix, and the second value matrix to obtain the second output feature.

[0038] According to an embodiment of the present invention, the backbone network is ResNet101.

[0039] The following will describe in further detail an image segmentation method based on block-level self-attention and block correlation mechanisms according to the present invention, with reference to the accompanying drawings.

[0040] Figure 1 The diagram illustrates the overall flow of a deep learning model according to an embodiment of the present invention.

[0041] Figure 2 The diagram illustrates the overall network structure of a deep learning model according to an embodiment of the present invention.

[0042] like Figure 1 and Figure 2As shown, the deep learning model mainly consists of three parts: a backbone network, a block-level self-attention module, and a segmentation head. The block-level self-attention module includes a block correlation mechanism and self-attention with relative position offset. Feature segmentation in the block-level self-attention module reduces the computational cost and parameter count of the network through local attention. The block correlation mechanism correlates the Q-matrix of all feature blocks with the entire feature map, ensuring that each feature block contains globally relevant information, resulting in feature blocks with global long-range dependencies. The introduction of relative position information compensates for the positional information lost in feature segmentation, and this introduced positional information does not significantly increase the computational cost of the network.

[0043] Step 1: Construct the ResNet101 backbone network. The backbone network of this invention uses the ResNet101 network, downsampled by 32 times, and replaces the 7×7 convolution with two 3×3 convolutions. Batch normalization and ReLU activation function are used after the convolutions.

[0044] Step 2: Divide the output features of the backbone network into blocks, and use 1×1 convolution to generate a query matrix (Q), a key matrix (K), and a value matrix (V) for each feature block. Use the Block Correlation Mechanism (BCM) to calculate the correlation between the Q matrix of each feature block and the global features, so that each feature block contains global information.

[0045] In the block-level self-attention module, the output feature X of the backbone network is divided into X... 11 , ..., X ij , ..., X IJ There are N feature blocks in total. Here, i∈{1,…,I} and j∈{1,…,J} represent the feature blocks in the i-th row and j-th column of X, respectively. X is the number of feature blocks, C is the number of feature channels, and H and W are the height and width of X. and These are the height and width of each feature block. Then, applying self-attention with a relative position offset to each feature block reduces the computational cost of self-attention from... Reduce to The computational cost of self-attention is reduced to 1 / N of its original value. The query matrix Q, key matrix K, and value matrix V required for self-attention computation are derived from each feature block X. ij Generated through a convolution operation, it can be represented as:

[0046]

[0047] in, Indicates the use of generating and 1×1 convolution operation of three matrices, C m =C / 8, which can reduce the number of channels.

[0048] After feature segmentation, if the self-attention mechanism is used only in a single feature block, there will be no information interaction between the feature blocks, and the concatenated feature map will not contain global information, resulting in a decrease in segmentation accuracy. To solve this problem, this invention proposes a block correlation mechanism to model the relationship between each feature block and the complete feature map, so that each feature block contains information related to the entire feature map. Only then can feature blocks containing global information be obtained, and the large feature map obtained after concatenation will contain long-distance dependency information related to the global context.

[0049] Specifically, this method uses a block correlation matrix to model the relationship between feature blocks and the entire feature map, enabling information exchange between feature blocks. The block correlation matrix is ​​generated by the input feature X through a convolution operation, and denoted by R. We can obtain:

[0050]

[0051] in Conv(·) represents a convolution operation. This method uses 1×1 convolution, that is, Conv(X) represents performing a convolution operation on the input feature X. In order to ensure that each feature block can be associated with the global feature, a block correlation matrix R is used to associate each feature block with the global feature X. The formula for the block matrix M is as follows:

[0052] M = R·Q (3)

[0053] According to formulas (1) and (3), we can obtain:

[0054]

[0055] As can be seen from formula (4), the M matrix is ​​formed by the block correlation matrix R for each feature block X in X. 11 , ..., X ij , ..., X IJ By aggregation, we obtain M = {M} 11 M 12 M ij M IJ}, That is, M ij It contains all the information in the entire feature map X, and then uses M ij Come with K ij and V ij By performing self-attention computation with relative position, the features calculated for each feature block no longer only contain information about the current block, but also contain long-range dependency information of the entire input features.

[0056] Step 3: Construct a self-attention network with relative position offset. The Block-Level Self-Attention (BLSA) module is embedded into the self-attention calculation process of each feature block. The output matrix M, matrix K, and matrix V of the BCM are used to calculate self-attention, and a relative position offset is introduced into the M, K, and V matrices of each feature block, collectively forming a Block-Level Self-Attention (BLSA) module. As shown in Step 2, M... ij It contains all the information in the entire feature map X, using M ij Come with K ij and V ij Computing self-attention can yield feature blocks with long-range dependency information that are globally relevant to the output of each self-attention function.

[0057] However, block segmentation causes each feature block to lose positional information. Therefore, in this method, M is used in the self-attention mechanism. ij K ij and V ij The matrix introduces a relative position offset.

[0058] Figure 3 A schematic diagram of a self-attention network structure with relative position offset is shown according to an embodiment of the present invention.

[0059] like Figure 3 As shown, the feature Y calculated for each feature block ij The calculation is as follows:

[0060]

[0061] Where p∈|X ij | indicates that pixel p is in the X ij Within the feature block, Pos ij Represents feature block X ij M ij Matrix and K ij Matrix relative position offset, Due to the different dimensions, this method is used alone. To represent V ij The relative position offset, Position offset Pos ij and In the self-attention method of this method, only matrix addition is performed, which does not bring a large amount of computation to the network.

[0062] Figure 4 The diagram illustrates the network structure of a block-level self-attention module according to an embodiment of the present invention.

[0063] like Figure 4As shown, the expression for the output Y of the entire BLSA module can be written from formula (5):

[0064]

[0065] In formula (6) This indicates that features are concatenated along the height dimension. This indicates the concatenation of features along the width dimension.

[0066] Step 4: The output of each feature block is concatenated in the spatial dimension and then fed into the segmentation head for segmentation, followed by experimental comparison. After feeding the output feature Y of the BLSA module into the segmentation head, predictions for four categories—inflammation, low-grade tumors, high-grade tumors, and early-stage cancer—are completed. The segmentation head consists of a 3×3 convolution and a 1×1 convolution. The 3×3 convolution is followed by batch normalization and a ReLU activation function, while the 1×1 convolution reduces the number of feature channels to 5, enabling predictions for four types of precancerous lesions and one background category.

[0067] This experiment used two NVIDIA Tesla A100 40G GPUs for training and validation, and a self-built esophageal cancer dataset. Each image had a dimension of 3×512×512. The parameter count and computational cost statistics were performed in a 1×3×512×512 dimension. The experimental results are shown in Table 1. According to the experimental results, the network proposed in this method achieves the lowest computational cost. Compared with the self-attention mechanism in DANet, the number of parameters is reduced by 1.64M and the computational cost is reduced by 30.51%. Compared with other networks, it achieves the best in both parameter count and computational cost, and the segmentation accuracy is the highest, surpassing the performance of mainstream segmentation networks in recent years on the esophageal cancer dataset.

[0068] Table 1 Comparative trials on the esophageal cancer dataset.

[0069]

[0070] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of the present invention. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of code containing one or more executable instructions for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions indicated in the blocks may occur in a different order than those indicated in the drawings. For example, two consecutively indicated blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in a block diagram or flowchart, and combinations of blocks in a block diagram or flowchart, may be implemented using a dedicated hardware-based system that performs the specified function or operation, or using a combination of dedicated hardware and computer instructions. Those skilled in the art will understand that the features recited in the various embodiments and / or claims of the present invention can be combined and / or combined in various ways, even if such combinations or combinations are not expressly stated in the present invention. In particular, the features described in the various embodiments and / or claims of this invention can be combined and / or combined in various ways without departing from the spirit and teachings of this invention. All such combinations and / or combinations fall within the scope of this invention.

[0071] The embodiments of the present invention have been described above. However, these embodiments are merely illustrative and not intended to limit the scope of the invention. Although various embodiments have been described above, this does not mean that the measures in the various embodiments cannot be used advantageously in combination. The scope of the invention is defined by the appended claims and their equivalents. Various substitutions and modifications can be made by those skilled in the art without departing from the scope of the invention, and all such substitutions and modifications should fall within the scope of the invention.

Claims

1. An image segmentation method based on block-level self-attention and block correlation mechanism, comprising: processing an endoscopic image using a backbone network to obtain a first output feature; performing block processing on the first output feature to obtain a plurality of feature blocks; for each of the feature blocks, processing the feature block using three 1x1 convolutions to obtain a query matrix, a first key matrix, and a first value matrix; performing convolution processing on the first output feature to obtain a block correlation matrix; obtaining a first block matrix based on the block correlation matrix and the query matrix; performing self-attention calculation on the first block matrix, the first key matrix, and the first value matrix based on respective relative position offsets of the first block matrix, the first key matrix, and the first value matrix to obtain a second output feature; spatially concatenating a plurality of second output features corresponding to the plurality of feature blocks to obtain a third output feature; and processing the third output feature using a segmentation head to obtain a segmentation result of the endoscopic image; wherein the obtaining of the first block matrix based on the block correlation matrix and the query matrix comprises: the performing of the self-attention calculation on the first block matrix, the first key matrix, and the first value matrix based on respective relative position offsets of the first block matrix, the first key matrix, and the first value matrix to obtain a second output feature comprises: using the block correlation matrix R to correlate the feature blocks and a first output feature X = {X 11 ,…,X ij ,…,X IJ}, wherein X ij represents an i-th row, j-th column feature block of the first output feature X, , H and W are the height and width of the first output feature X, and are the height and width of the feature block, C is the number of feature channels, and the formula of the first block matrix M is as follows: In this formula, is the feature block X in the i-th row and j-th column ij a query matrix generated by a 1 x 1 convolution, denotes a 1 x 1 convolution, the block correlation matrix R = Conv(X) denotes that the block correlation matrix R is obtained by convolution of the first output feature X.

2. The method of claim 1, wherein, processing the first block matrix based on a relative position offset between the first block matrix and the first key matrix to obtain a second block matrix; processing the first key matrix based on the relative position offset between the first block matrix and the first key matrix to obtain a second key matrix; processing the first value matrix based on a relative position offset of the first value matrix to obtain a second value matrix; and performing self-attention calculation on the second block matrix, the second key matrix, and the second value matrix to obtain the second output feature. The backbone network is ResNet101.

3. The method of claim 1, wherein, ​

Images