Image segmentation method, device, electronic device and computer storage medium

Abstract

The present application relates to the field of image processing technology, and provides an image segmentation method, device, electronic device and computer storage medium. It includes: extracting features from an image to be segmented based on an attention-based residual structure to obtain an initial feature map; reconstructing a feature map based on multi-head self-attention based on the initial feature map to obtain an intermediate feature map; performing an attention-enhanced upsampling process based on the intermediate feature map to obtain a segmented image. The present application can alleviate the problem of gradient vanishing and improve the ability to extract fine-grained features, quickly and effectively obtain an initial feature map, and at the same time can strengthen the semantic representation of each feature point on the initial feature map, improve the refined segmentation of pixels between different categories, and enhance the ability to recognize targets, strengthen the distinction between targets of different scales and refine the features between targets of different categories, thereby effectively improving the accuracy of image segmentation.

Description

Technical Field

[0001] The present application relates to the field of image processing technology, and in particular to an image segmentation method, device, electronic equipment and computer storage medium. Background Art

[0002] In recent years, with the development of computer technology and artificial intelligence, image segmentation tasks have gradually become an important part of computer vision. Among them, semantic segmentation of remote sensing images is a popular branch of image segmentation tasks. As a special semantic segmentation task, segmentation of geospatial objects in remote sensing images can be widely used in typical task scenarios such as land and sea segmentation, old city reconstruction, building surveying and mapping, road extraction, and vegetation coverage assessment. Semantic segmentation of images is to divide each pixel in the image into corresponding categories, that is, to achieve pixel-level classification. Together with typical tasks such as motion estimation, image classification, target detection, and target tracking, it constitutes a basic problem in the field of computer vision. However, it is still a difficult task to let machines learn to segment targets in remote sensing images. With the advancement of technology and the innovation of researchers, a series of classic methods have gradually emerged. Traditional image segmentation methods divide images into regions according to features such as grayscale, color, texture, and shape, so that differences between regions are obvious and similarities are present within regions, so as to segment image regions. However, the segmentation results obtained by this method have low accuracy, resulting in low accuracy of current image segmentation. Summary of the invention

[0003] The embodiments of the present application provide an image segmentation method, device, electronic device and computer storage medium to solve the technical problem of low image segmentation accuracy caused by low accuracy of segmentation results obtained by using traditional image segmentation methods.

[0004] In a first aspect, an embodiment of the present application provides an image segmentation method, comprising:

[0005] According to the attention-based residual structure, the feature extraction of the image to be segmented is performed to obtain the initial feature map;

[0006] Performing multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map;

[0007] An upsampling process based on attention enhancement is performed according to the intermediate feature map to obtain a segmented image.

[0008] In one embodiment, the step of extracting features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map includes:

[0009] The image to be segmented is processed by convolution, batch normalization and neural activation to obtain the first feature map;

[0010] Performing convolution and feature extraction based on a parameter-free attention mechanism on the first feature map to obtain a second feature map;

[0011] Superimposing the image to be segmented and the second feature map to obtain a third feature map;

[0012] An initial feature map is determined based on the third feature map.

[0013] In one embodiment, the step of performing convolution on the first feature map and feature extraction based on a parameter-free attention mechanism to obtain a second feature map includes:

[0014] Performing convolution processing on the first feature map to obtain a convolution feature map;

[0015] Determining an energy distribution function according to the convolution feature map;

[0016] A second feature map is generated according to the energy distribution function and the convolution feature map.

[0017] In one embodiment, the step of determining the energy distribution function according to the convolution feature map comprises:

[0018] Determining a feature space mean according to the convolution feature map;

[0019] Calculate the variance of the width and height of the convolution feature map in the channel direction according to the feature space mean;

[0020] An energy distribution function of the convolution feature map is determined according to the feature space mean, the variance and the energy coefficient.

[0021] In one embodiment, the step of reconstructing a feature map based on multi-head self-attention based on the initial feature map to obtain an intermediate feature map includes:

[0022] Performing layer normalization processing on the initial feature map to obtain a layer normalized feature map;

[0023] Determining a query vector, a key vector, and a value vector for each vector in the layer normalized feature map, respectively;

[0024] The layer normalized feature map is reconstructed according to the query vector, the key vector and the value vector of each of the vectors to obtain an intermediate feature map.

[0025] In one embodiment, the step of reconstructing the layer normalized feature map according to the query vector, the key vector and the value vector of each of the vectors to obtain the intermediate feature map comprises:

[0026] Determine the dot product result of the query vector and the corresponding key vector for each of the vectors;

[0027] Normalizing each of the dot product results respectively to obtain multiple normalized results;

[0028] Each of the normalized results is weighted respectively and added to the value vector of the corresponding vector to obtain a summed result;

[0029] The summation results are summed to obtain an intermediate feature map.

[0030] In one embodiment, the step of performing attention-enhanced upsampling processing according to the intermediate feature map to obtain a segmented image includes:

[0031] Performing convolution, batch normalization, and neural activation processing on the intermediate feature map to obtain a fourth feature map;

[0032] Performing transposed convolution, batch normalization, and neural activation processing on the fourth feature map to obtain a fifth feature map;

[0033] The segmented image is obtained by performing upsampling processing in combination with the fifth feature map through spatial attention and / or channel attention.

[0034] In a second aspect, an embodiment of the present application provides an image segmentation device, comprising:

[0035] An extraction module is used to extract features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map;

[0036] A reconstruction module, used for performing multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map;

[0037] An upsampling module is used to perform attention-enhanced upsampling processing according to the intermediate feature map to obtain a segmented image.

[0038] In a third aspect, an embodiment of the present application provides an electronic device, comprising a processor and a memory storing a computer program, wherein when the processor executes the program, the steps of the image segmentation method described in the first aspect or the second aspect are implemented.

[0039] In a fourth aspect, an embodiment of the present application provides a computer storage medium, which is a computer-readable storage medium and includes a computer program. When the computer program is executed by a processor, the steps of the image segmentation method described in the first aspect or the second aspect are implemented.

[0040] The image segmentation method, device, electronic device and computer storage medium provided in the embodiments of the present application can alleviate the problem of gradient vanishing and improve the ability to extract fine-grained features by extracting features of the image to be segmented based on an attention-based residual structure, and can quickly and effectively obtain an initial feature map; based on the obtained initial feature map, a feature map reconstruction based on multi-head self-attention is performed to strengthen the semantic representation of each feature point on the initial feature map, improve the refined segmentation of pixels between different categories, and obtain an intermediate feature map with high segmentation accuracy; further, an upsampling processing based on attention enhancement is performed according to the intermediate feature map, which can enhance the ability to recognize targets, strengthen the distinction between targets of different scales, and at the same time refine the features between targets of different categories to obtain a segmented image with high precision, thereby effectively improving the accuracy of image segmentation. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] In order to more clearly illustrate the technical solutions in the present application or the prior art, a brief introduction will be given below to the drawings required for use in the embodiments or the description of the prior art. Obviously, the drawings described below are some embodiments of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0042] Figure 1 This is one of the flowcharts of the image segmentation method provided in the embodiment of the present application;

[0043] Figure 2 This is the second flowchart of the image segmentation method provided in the embodiment of the present application;

[0044] Figure 3 This is the third flow chart of the image segmentation method provided in the embodiment of the present application;

[0045] Figure 4 This is the fourth flowchart of the image segmentation method provided in the embodiment of the present application;

[0046] Figure 5 It is a schematic diagram of functional modules of an embodiment of the image segmentation device of the present application;

[0047] Figure 6 It is a schematic diagram of the structure of an electronic device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0048] In order to make the purpose, technical solutions and advantages of this application clearer, the technical solutions in this application will be clearly and completely described below in conjunction with the drawings in the embodiments of this application. Obviously, the described embodiments are part of the embodiments of this application, not all of them. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of this application.

[0049] Figure 1 This is one of the flowcharts of the image segmentation method provided in the embodiment of the present application. Figure 1 , the present application embodiment provides an image segmentation method, which may include:

[0050] Step S100, extracting features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map;

[0051] The image segmentation method in this embodiment can be applied to electronic devices such as computers, tablet computers, servers, and smart phones, and the electronic device includes an improved UNet model framework. The improved UNet model in this embodiment is called a MA-UNet (Multiattention UNet) model. Compared with the traditional UNet model framework, this embodiment adds a residual structure to the encoder of the original UNet model, and adds four attention mechanisms to the original UNet model, wherein the four attention mechanisms include a simAM (A Simple, Parameter-Free Attention Module, parameter-free attention model) attention mechanism, a multi-head self-attention mechanism, a spatial attention mechanism, and a channel attention mechanism. At the same time, compared with the original UNet model, this embodiment uses a 3x3 convolution with a step size of 2 and a padding of 1 instead of a maximum pooling for downsampling. Among them, the MA-UNet model includes an encoder and a decoder like the traditional UNet model, wherein the encoder can compress the input image, and the decoder decompresses and restores the compressed image accordingly.

[0052] In one example, when a user has a need to segment an image, the image to be segmented can be input into the MA-UNet model of the electronic device, for example, an image of size 256x256x3 is input, and the image is segmented by the MA-UNet model. On the other hand, after receiving the image to be segmented input by the user, the electronic device determines the image as the image to be segmented. It should be noted that the encoder of the MA-UNet model of this embodiment may include multiple simAM-based residual blocks and multiple convolution kernels, wherein the step size of the convolution kernel can be set according to actual needs. For example, the encoder of the MA-UNet model of this embodiment may include 5 simAM-based residual blocks and 5 convolution kernels with a step size of 2, wherein the convolution kernel is used to perform a convolution operation.

[0053] Therefore, in this embodiment, the electronic device can perform convolution, batch normalization and neural activation processing on the segmented image through the simAM-based residual block and convolution kernel in the encoder of the MA-UNet model, and determine the processed feature map as the first feature map. Batch normalization is also called batch normalization, which is a technology used to improve the performance and stability of artificial neural networks. Further, the first feature map is convolved and feature extracted based on a parameter-free attention mechanism, and the processed feature map is determined as the second feature map. The first feature map obtained is superimposed on the second feature map, and the superimposed feature map is determined as the third feature map. It should be noted that the third feature map obtained above is a feature map obtained by performing a feature extraction based on a parameter-free attention mechanism with a residual block and a convolution kernel, and the encoder of this embodiment may include 5 residual blocks based on simAM and 5 convolution kernels, so it is also necessary to perform multiple layers or multiple processing steps according to the third feature map, that is, convolution, batch normalization and neural activation processing are performed on the third feature map, convolution and feature extraction based on the parameter-free attention mechanism are performed on the processed feature map, and the processed feature map is superimposed with the feature map obtained in the previous step to obtain the output feature map of the layer, and the output feature map of the layer is used as the input of the next layer for the above processing until the processing of the 5 layers in this embodiment is completed, and the feature map of the final output is determined as the initial feature map. So as to facilitate the subsequent reconstruction of the feature map based on multi-head self-attention based on the initial feature map to obtain the intermediate feature map. By extracting features of the segmented image based on the residual structure based on attention, the problem of gradient disappearance can be alleviated and the ability to extract fine-grained features can be improved, and the initial feature map can be obtained quickly and effectively.

[0054] Step S200, performing a feature map reconstruction based on multi-head self-attention based on the initial feature map to obtain an intermediate feature map;

[0055] After the feature extraction is completed to obtain the initial feature map, the initial feature map is layer-normalized, and the obtained feature map is determined as a layer-normalized feature map. It should be noted that the layer-normalized feature map includes multiple vectors, so in this embodiment, the query vector, key vector and value vector (i.e., Query, Keys, Values ​​vector) of each vector in the layer-normalized feature map can be determined respectively, that is, each vector has a corresponding query vector, key vector and value vector. Furthermore, according to the query vector, key vector and value vector corresponding to each vector, the layer-normalized feature map is reconstructed, and after the feature map reconstruction is completed, the obtained feature map is determined as an intermediate feature map. So as to facilitate the subsequent upsampling processing based on attention enhancement according to the intermediate feature map to obtain a segmented image. Based on the obtained initial feature map, feature map reconstruction based on multi-head self-attention is performed, which can strengthen the semantic representation of each feature point on the initial feature map, improve the refined segmentation of pixels between different categories, and obtain an intermediate feature map with higher segmentation accuracy.

[0056] Step S300, performing attention-enhanced upsampling processing on the intermediate feature map to obtain a segmented image.

[0057] After the intermediate feature map is obtained through feature map reconstruction, the intermediate feature map is convolved, batch normalized and neural activated, and the obtained feature map is determined as the fourth feature map. Further, the fourth feature map is transposed convolution, batch normalized and neural activated, and the obtained feature map is determined as the fifth feature map. Further, upsampling processing is performed in combination with the fifth feature map through spatial attention and / or channel attention, and a segmented image is obtained after the upsampling processing is completed, and the segmented image can be output for user viewing, wherein the upsampling processing includes feature fusion of the fifth feature map with the feature map output by the last layer of the encoder, and the corresponding spatial attention or channel attention is selected according to the scale of the feature map output by the last layer. Combined with the convolution kernel for processing, obtain the upsampled feature map of this layer. Furthermore, the upsampled feature map of this layer is subjected to transposed convolution, batch normalization and neural activation processing, and the processed feature map is fused with the feature map outputted by the previous layer. The corresponding spatial attention or channel attention is selected according to the scale of the feature map outputted by the previous layer in combination with the convolution kernel for processing to obtain the upsampled feature map of the previous layer, and so on, until the upsampling processing of all layers is completed, and the final output image is the segmented image. The upsampling processing based on attention enhancement according to the intermediate feature map can enhance the recognition ability of the target, strengthen the distinction between targets of different scales, and refine the features between targets of different categories to obtain a segmented image with high accuracy.

[0058] The image segmentation method provided in the embodiment of the present application can alleviate the problem of gradient vanishing and improve the ability to extract fine-grained features by extracting features of the image to be segmented based on an attention-based residual structure, and can quickly and effectively obtain an initial feature map; based on the obtained initial feature map, a feature map reconstruction based on multi-head self-attention is performed to strengthen the semantic representation of each feature point on the initial feature map, improve the refined segmentation of pixels between different categories, and obtain an intermediate feature map with high segmentation accuracy; further, based on the intermediate feature map, an upsampling processing based on attention enhancement is performed to enhance the recognition ability of the target, strengthen the distinction between targets of different scales, and at the same time refine the features between targets of different categories to obtain a segmented image with high precision, thereby effectively improving the accuracy of image segmentation.

[0059] Figure 2 The second flowchart of the image segmentation method provided in the embodiment of the present application. Figure 2 In one embodiment, the step of extracting features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map includes:

[0060] Step S1001, performing convolution, batch normalization and neural activation processing on the image to be segmented to obtain a first feature map;

[0061] Step S1002, performing convolution and feature extraction based on a parameter-free attention mechanism on the first feature map to obtain a second feature map;

[0062] Step S1003, superimposing the image to be segmented and the second feature map to obtain a third feature map;

[0063] Step S1004: determining an initial feature map based on the third feature map.

[0064] After obtaining the image to be segmented, the electronic device in this embodiment can perform convolution operation, batch normalization and neural activation processing on the image to be segmented through the residual block based on simAM and the convolution kernel in the encoder of the MA-UNet model. Specifically, it can be Conv3x3+BN+Relu processing. More specifically, the image to be segmented can be subjected to Conv3x3 operation, the feature map obtained by the convolution operation is subjected to batch normalization processing, and the feature map obtained by the batch normalization processing is activated by Relu, and the activated feature map is determined as the first feature map. Among them, Relu is a linear rectification function (Linear rectification function), also known as a corrected linear unit, which is an activation function commonly used in artificial neural networks. Further, the first feature map is subjected to convolution operation, and the feature map obtained by the convolution operation is subjected to feature extraction based on a parameter-free attention mechanism (i.e., simAM attention processing), and the feature map obtained by the processing is determined as the second feature map. The first feature map obtained is superimposed with the second feature map, and the superimposed feature map is determined as the third feature map.

[0065] It should be noted that the third feature map obtained above is a feature map obtained by performing a feature extraction based on a parameter-free attention mechanism through a residual block and a convolution kernel, which can be regarded as a feature extraction of a layer of residual structure based on attention. The encoder of this embodiment may include 5 residual blocks based on simAM and 5 convolution kernels, that is, it can be regarded as having five layers. Therefore, it is necessary to perform feature extraction of four layers of residual structure based on attention according to the third feature map, that is, perform convolution, batch normalization and neural activation processing on the third feature map, perform convolution and feature extraction based on the parameter-free attention mechanism on the processed feature map, superimpose the processed feature map with the feature map obtained in the previous step to obtain the output feature map of this layer, and use the output feature map of this layer as the input of the next layer for the above processing until the processing of 5 layers in this embodiment is completed, and the feature map of the final output is determined as the initial feature map.

[0066] Furthermore, the step of performing convolution on the first feature map and feature extraction based on a parameter-free attention mechanism to obtain a second feature map includes:

[0067] Step S10021, performing convolution processing on the first feature map to obtain a convolution feature map;

[0068] Step S10022, determining an energy distribution function according to the convolution feature map;

[0069] Step S10023, generating a second feature map according to the energy distribution function and the convolution feature map.

[0070] After obtaining the first feature map, a convolution operation is performed on the first feature map, and after the convolution operation is completed, the obtained feature map is determined as the convolution feature map. After the first feature map is convolved, it is necessary to evaluate the importance of each neuron in the obtained convolution feature map. In this embodiment, the feature space mean is determined according to the convolution feature map, and the variance of the width and height of the convolution feature map in its channel direction is calculated according to the feature space mean, and the preset energy coefficient is obtained. According to the feature space mean, variance and energy coefficient, the energy distribution function of the convolution feature map is determined. After obtaining the energy distribution function, the energy distribution function is input into formula (1) for calculation, and the feature map after the calculation is determined as the second feature map. Wherein, formula (1) is as follows:

[0071] X * =sigmoid(1 / E)·X (1)

[0072] Among them, X * is the enhanced feature map, i.e., the second feature map, E is the energy distribution function, X is the convolution feature map, and sigmoid (1 / E) is the three-dimensional attention weight.

[0073] Furthermore, the step of determining the energy distribution function according to the convolution feature map includes:

[0074] Step S100221, determining a feature space mean according to the convolution feature map;

[0075] Step S100222, calculating the variance of the width and height of the convolution feature map in the channel direction according to the feature space mean;

[0076] Step S100223, determining the energy distribution function of the convolution feature map according to the feature space mean, the variance and the energy coefficient.

[0077] After obtaining the convolution feature map, the convolution feature map is input into formula (2) for operation, and the feature space mean is obtained after the operation is completed. Formula (2) is as follows:

[0078] d=(XX.mean(dim=[2,3])) 2 (2)

[0079] Among them, d is the mean of the feature space, X is the convolution feature map, mean is the average value, and dim is the dimension.

[0080] After obtaining the feature space mean, the width and height of the convolution feature map are obtained, and the feature space mean, width and height data are input into formula (3) for calculation. After the calculation is completed, the variance of the width and height of the convolution feature map in its channel direction is obtained. Formula (3) is as follows:

[0081]

[0082] Among them, v is the variance of the width and height of the convolution feature map in its channel direction, H is the height of the convolution feature map, W is the width of the convolution feature map, d is the mean of the feature space, dim is the dimension, and sum is the sum.

[0083] After obtaining the variance of the width and height of the convolution feature map in the channel direction, the preset energy coefficient is obtained, and the feature space mean, variance and energy coefficient are input into formula (4) for calculation. After the calculation is completed, the energy distribution function of the convolution feature map is obtained. Formula (4) is as follows:

[0084]

[0085] Among them, E is the energy distribution function, d is the mean of the feature space, v is the variance of the width and height of the convolution feature map in its channel direction, and ρ is the energy coefficient.

[0086] By extracting features from the segmented image based on the attention-based residual structure, the problem of gradient vanishing can be alleviated and the ability to extract fine-grained features can be improved, so that the initial feature map can be obtained quickly and effectively.

[0087] Figure 3 The third flowchart of the image segmentation method provided in the embodiment of the present application. Figure 3 In one embodiment, the step of reconstructing a feature map based on multi-head self-attention based on the initial feature map to obtain an intermediate feature map includes:

[0088] Step S2001, performing layer normalization processing on the initial feature map to obtain a layer normalized feature map;

[0089] Step S2002, respectively determining a query vector, a key vector and a value vector of each vector in the layer normalized feature map;

[0090] Step S2003, reconstructing the layer normalized feature map according to the query vector, key vector and value vector of each of the vectors to obtain an intermediate feature map.

[0091] It should be noted that after 5 convolutions and 5 downsamplings, for an image with an input size of 256×256×3, the size of the initial feature map obtained is 8×8×512, that is, the feature map of each dimension has only 64 pixels, and the receptive field corresponding to the high-dimensional feature is relatively large. In order to further improve the correlation and globality between each feature point, this embodiment reconstructs the feature map through a multi-head self-attention mechanism. Specifically, after obtaining the initial feature map, this embodiment performs layer normalization on the initial feature map, and after the processing is completed, the obtained feature map is determined as a layer normalized feature map, where layer normalization is similar to ordinary normalization and is used to normalize the layers in the network.

[0092] After obtaining the layer-normalized feature map, the query vector, key vector and value vector are generated for each vector of the layer-normalized feature map input through the fully connected layer, that is, the corresponding query vector, key vector and value vector are generated for each vector. Further, the layer-normalized feature map is reconstructed according to the query vector, key vector and value vector of each vector to obtain an intermediate feature map. Specifically, all candidate positions are queried through the query vector to find vectors that are closely related to themselves, where each candidate position will have a pair of key vectors and value vectors. The query process is to find the process of doing dot products with the key vectors of all candidate positions. The dot product results are weighted to their respective value vectors after softmax (normalization), and the feature map obtained by summing up the weighted value vectors is determined as the intermediate feature map.

[0093] Furthermore, the step of reconstructing the layer normalized feature map according to the query vector, the key vector and the value vector of each of the vectors to obtain the intermediate feature map comprises:

[0094] Step S20031, respectively determining the dot product result of the query vector and the corresponding key vector of each of the vectors;

[0095] Step S20032, normalizing each of the dot product results to obtain a plurality of normalized results;

[0096] Step S20033, weighting each of the normalized results respectively and adding them to the value vector of the corresponding vector to obtain a summed result;

[0097] Step S20034, summing up the summation results to obtain an intermediate feature map.

[0098] After obtaining the query vector, key vector and value vector corresponding to each vector, dot product operation is performed on the query vector of each vector with the key vector corresponding to the vector, and multiple dot product results are obtained after completing the dot product operation of all vectors. Further, each dot product result is normalized, and multiple normalized results are obtained after the operation is completed. Obtain the weight value set for each vector in advance (the default is 1, but the weight value of each vector can be adjusted accordingly according to actual needs), perform weighted operation on each normalized result and its corresponding weight value, add the result of the weighted operation and the value vector corresponding to the vector, and obtain the sum result of the vector. After completing the calculation of all normalized results, multiple sum results are obtained. And sum each sum result, that is, add each combined result, and determine the feature map obtained after the operation as the intermediate feature map. Multi-head self-attention uses parallelization on the basis of self-attention to improve the operation efficiency. Among them, the calculation formula for the sum result of each vector is shown in formula (5):

[0099]

[0100] Among them, Q is the query vector, K is the key vector, V is the value vector, softmax is the normalization operation, and d K is the number of key vectors.

[0101] When the combined results are added together through multi-head self-attention, it can be achieved by the following formula (6):

[0102] MultiHead(Q,K,V)=Concat(head1...head2)W O (6)

[0103] Among them, Q is the query vector, K is the key vector, V is the value vector, and W o To calculate the weight, head is the sum of the vectors. For the head of the i-th vector, we have:

[0104] head i =Attention(QW i Q ,KW i K ,VW i V )

[0105] Where i represents the i-th vector, Q is the query vector, K is the key vector, V is the value vector, and W Q i , W K i , W V iare the weights of the query vector, key vector, and value vector in each vector respectively.

[0106] It should be noted that the input and output of multi-head self-attention are both two-dimensional matrix vectors. When applying segmentation tasks, the feature map needs to be flattened into a two-dimensional feature matrix. The specific method is to retain the channel dimension of the initial feature map of 8×8×512, flatten it into a two-dimensional matrix of 512×64, and after layer normalization and attention enhancement for feature reconstruction, output a two-dimensional matrix of 512×64 and restore it to the original feature map size.

[0107] Based on the obtained initial feature map, feature map reconstruction based on multi-head self-attention is performed to strengthen the semantic representation of each feature point on the initial feature map, improve the refined segmentation of pixels between different categories, and obtain an intermediate feature map with higher segmentation accuracy.

[0108] Figure 4 This is a fourth flow chart of the image segmentation method provided in the embodiment of the present application. Figure 4 In one embodiment, the step of performing attention-enhanced upsampling processing according to the intermediate feature map to obtain a segmented image includes:

[0109] Step S3001, performing convolution, batch normalization and neural activation processing on the intermediate feature map to obtain a fourth feature map;

[0110] Step S3002, performing transposed convolution, batch normalization and neural activation processing on the fourth feature map to obtain a fifth feature map;

[0111] Step S3003, upsampling is performed by spatial attention and / or channel attention in combination with the fifth feature map to obtain a segmented image.

[0112] After the intermediate feature map is reconstructed through the feature map, the image encoding stage is completed, and the decoding stage can be entered. The intermediate feature map is decoded by the decoder to restore the image to the original image size to obtain a segmented image. It should be noted that the decoder of the traditional UNet model consists of bilinear interpolation, feature splicing and two 3x3 convolutional layers. After repeated upsampling and convolution operations, the number of channels is finally reduced to a specific number through 1×1 convolution to obtain a segmentation result map. In this embodiment, transposed convolution is used instead of bilinear interpolation for upsampling, and an attention mechanism is used to enhance the fused features in the feature fusion stage based on jump connection. Since the low-resolution high-dimensional features are upsampled, simple feature splicing can easily confuse the features of targets of different scales, especially when their boundaries are adjacent, their contours are blurred and irregular. Therefore, this embodiment designs an attention enhancement module to enhance recognition capabilities, strengthen the distinction between targets of different scales, and refine the features between targets of different categories.

[0113] Specifically, the intermediate feature map is subjected to convolution, batch normalization and neural activation processing, and the obtained feature map is determined as the fourth feature map. Further, the fourth feature map is subjected to transposed convolution, batch normalization and neural activation processing, and the obtained feature map is determined as the fifth feature map. Further, upsampling processing is performed by combining the fifth feature map with spatial attention and / or channel attention, and a segmented image is obtained after the upsampling processing is completed, and the segmented image can be output for user viewing, wherein the upsampling processing includes feature fusion of the fifth feature map with the feature map output by the last layer of the encoder, and selecting the corresponding spatial attention or channel attention combined with the convolution kernel for processing according to the scale of the feature map output by the last layer to obtain the upsampled feature map of the layer. Further, the upsampled feature map of the layer is subjected to transposed convolution, batch normalization and neural activation processing, and the feature map obtained by the processing is subjected to feature fusion with the feature map output by the previous layer, and the corresponding spatial attention or channel attention combined with the convolution kernel is selected according to the scale of the feature map output by the previous layer to obtain the upsampled feature map of the previous layer, and so on, until the upsampling processing of all layers is completed, and the image output at the end is the segmented image.

[0114] It should be noted that high-dimensional features are often compressed in the channel. Therefore, for the last two smaller-scale feature maps (i.e., the feature maps output after the feature extraction based on the parameter-free attention mechanism of the last two layers), the feature map of the corresponding layer after transposed convolution (e.g., the fifth feature map) and the feature map output after the feature extraction based on the parameter-free attention mechanism of the layer (e.g., the feature map output after the feature extraction based on the parameter-free attention mechanism of the layer corresponding to the fifth feature map) are feature fused, and channel attention is used to integrate the fused feature maps. Specifically, the fused feature maps are further extracted through convolution operations, and the convolved feature maps are average pooled and maximum pooled in the spatial dimensions, and after linear mapping and addition, the channel weight coefficient is obtained through the activation function. Finally, the channel weight coefficient is multiplied by the original feature map to obtain an enhanced feature map, and the feature map is convolved, batch normalized, and neurally activated, and the obtained feature map is used as the input of the previous layer so that the previous layer can perform transposed convolution and other operations on it.

[0115] Since the feature map resolution of the shallow features is relatively large, the spatial feature distribution has a greater impact on feature fusion. Therefore, in one embodiment, for the first three larger-scale feature maps (i.e., the feature maps output after feature extraction based on the parameter-free attention mechanism of the first three layers), the feature map of the corresponding layer after transposed convolution is fused with the feature map output after feature extraction based on the parameter-free attention mechanism of the layer, and spatial attention is used to integrate the fused feature maps, specifically: further extract features from the fused feature map through convolution operation, perform average pooling and maximum pooling on the channel dimension of the convolved feature map, and splice them together, then, after convolution and activation, obtain the spatial weight coefficient, finally, multiply the spatial weight coefficient with the original feature map to obtain an enhanced feature map, and perform convolution, batch normalization and neural activation on the feature map to obtain the obtained feature map as the input of the previous layer, until the enhanced feature map of the first layer is obtained, and the feature map is convolved twice, batch normalized and neural activated to obtain the final segmented image. By using channel attention and spatial attention mechanisms respectively in the feature fusion stage of different scales, the feature information of objects of different scales can be better fused.

[0116] Attention-enhanced upsampling processing is performed on the intermediate feature map to enhance the recognition ability of the target, strengthen the distinction between targets of different scales, and refine the features between targets of different categories to obtain a segmented image with high precision.

[0117] It should be further explained that, since the total number of pixels in each category is the sample size for the segmentation task, when the number of categories is large, there will often be an imbalance in the sample sizes of different categories. This makes the network of the MA-UNet model more biased towards predicting categories with larger sample sizes, thereby bringing lower prediction losses. Such prediction results will result in poor segmentation of categories with smaller sample sizes. When pixels of categories with smaller sample sizes are mistakenly classified as pixels of other categories, the accuracy of other categories is also reduced in disguise. In response to the problem of imbalance in multi-class target categories, this embodiment uses weighted cross entropy loss. By adjusting the weight benefits of different categories, the loss costs of misclassification of different categories are different. A larger weight is assigned to a category with a smaller sample size, and a smaller weight is assigned to a category with a larger sample size, so that the network pays more attention to categories with smaller sample sizes. The improved cross entropy loss formula is as follows:

[0118]

[0119] Where m is the number of samples; n is the number of categories; k is the weight corresponding to the current category; y is whether the true category corresponding to the current sample is category i, if yes, it is 1, otherwise it is 0; y *The probability that the current sample point predicted by the model is category i.

[0120] To address the problem of imbalanced multi-class target categories, this embodiment uses weighted cross entropy loss to make the model pay more attention to categories with fewer samples, which is beneficial to improving the accuracy of image segmentation based on the model.

[0121] Furthermore, the present application also provides an image segmentation device.

[0122] Reference Figure 5 , Figure 5 Schematic diagram of functional modules of an embodiment of an image segmentation device of the present application.

[0123] The image segmentation device comprises:

[0124] An extraction module 100 is used to extract features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map;

[0125] A reconstruction module 200, configured to perform multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map;

[0126] The upsampling module 300 is used to perform attention-enhanced upsampling processing according to the intermediate feature map to obtain a segmented image.

[0127] The image segmentation device provided in the embodiment of the present application can alleviate the problem of gradient vanishing and improve the ability to extract fine-grained features by extracting features of the image to be segmented based on the attention-based residual structure, and can quickly and effectively obtain the initial feature map; based on the obtained initial feature map, a feature map reconstruction based on multi-head self-attention is performed to strengthen the semantic representation of each feature point on the initial feature map, improve the refined segmentation of pixels between different categories, and obtain an intermediate feature map with high segmentation accuracy; further, based on the intermediate feature map, upsampling processing based on attention enhancement is performed to enhance the recognition ability of the target, strengthen the distinction between targets of different scales, and at the same time refine the features between targets of different categories to obtain a segmented image with high precision, thereby effectively improving the accuracy of image segmentation.

[0128] In one embodiment, the extraction module 100 is specifically used for:

[0129] The image to be segmented is processed by convolution, batch normalization and neural activation to obtain the first feature map;

[0130] Performing convolution and feature extraction based on a parameter-free attention mechanism on the first feature map to obtain a second feature map;

[0131] Superimposing the image to be segmented and the second feature map to obtain a third feature map;

[0132] An initial feature map is determined based on the third feature map.

[0133] In one embodiment, the extraction module 100 includes a first determination module (not shown in the figure), wherein the first determination module is used to:

[0134] Performing convolution processing on the first feature map to obtain a convolution feature map;

[0135] Determining an energy distribution function according to the convolution feature map;

[0136] A second feature map is generated according to the energy distribution function and the convolution feature map.

[0137] In one embodiment, the first determination module 100 includes a second determination module (not shown in the figure), and the second determination module is used to:

[0138] Determining a feature space mean according to the convolution feature map;

[0139] Calculate the variance of the width and height of the convolution feature map in the channel direction according to the feature space mean;

[0140] An energy distribution function of the convolution feature map is determined according to the feature space mean, the variance and the energy coefficient.

[0141] In one embodiment, the reconstruction module 200 is specifically used for:

[0142] Performing layer normalization processing on the initial feature map to obtain a layer normalized feature map;

[0143] Determining a query vector, a key vector, and a value vector for each vector in the layer normalized feature map, respectively;

[0144] The layer normalized feature map is reconstructed according to the query vector, the key vector and the value vector of each of the vectors to obtain an intermediate feature map.

[0145] In one embodiment, the reconstruction module 200 includes a calculation module (not shown in the figure), wherein the calculation module is used to:

[0146] Determine the dot product result of the query vector and the corresponding key vector for each of the vectors;

[0147] Normalizing each of the dot product results respectively to obtain multiple normalized results;

[0148] Each of the normalized results is weighted respectively and added to the value vector of the corresponding vector to obtain a summed result;

[0149] The summation results are summed to obtain an intermediate feature map.

[0150] In one embodiment, the up-sampling module 300 is specifically used for:

[0151] Performing convolution, batch normalization, and neural activation processing on the intermediate feature map to obtain a fourth feature map;

[0152] Performing transposed convolution, batch normalization, and neural activation processing on the fourth feature map to obtain a fifth feature map;

[0153] The segmented image is obtained by performing upsampling processing in combination with the fifth feature map through spatial attention and / or channel attention.

[0154] Figure 6 An example of a physical structure diagram of an electronic device is shown in FIG. Figure 6 As shown, the electronic device may include: a processor 810, a communication interface 820, a memory 830 and a communication bus 840, wherein the processor 810, the communication interface 820 and the memory 830 communicate with each other via the communication bus 840. The processor 810 may call a computer program in the memory 830 to execute the steps of the image segmentation method, for example including:

[0155] According to the attention-based residual structure, the feature extraction of the image to be segmented is performed to obtain the initial feature map;

[0156] Performing multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map;

[0157] An upsampling process based on attention enhancement is performed according to the intermediate feature map to obtain a segmented image.

[0158] In addition, the logic instructions in the above-mentioned memory 830 can be implemented in the form of a software functional unit and can be stored in a computer-readable storage medium when it is sold or used as an independent product. Based on this understanding, the technical solution of the present application can be essentially or partly embodied in the form of a software product that contributes to the prior art, and the computer software product is stored in a storage medium, including several instructions to enable a computer device (which can be a personal computer, a server, or a network device, etc.) to perform all or part of the steps of the method described in each embodiment of the present application. The aforementioned storage medium includes: various media that can store program codes, such as a USB flash drive, a mobile hard disk, a read-only memory (ROM), a random access memory (RAM), a disk or an optical disk.

[0159] On the other hand, an embodiment of the present application further provides a computer storage medium, wherein the computer storage medium is a computer-readable storage medium, wherein the computer-readable storage medium stores a computer program, wherein the computer program is used to enable a processor to execute the steps of the methods provided in the above embodiments, for example, including:

[0160] According to the attention-based residual structure, the feature extraction of the image to be segmented is performed to obtain the initial feature map;

[0161] Performing multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map;

[0162] An upsampling process based on attention enhancement is performed according to the intermediate feature map to obtain a segmented image.

[0163] The computer-readable storage medium can be any available medium or data storage device that can be accessed by the processor, including but not limited to magnetic storage (such as floppy disks, hard disks, magnetic tapes, magneto-optical disks (MO), etc.), optical storage (such as CD, DVD, BD, HVD, etc.), and semiconductor storage (such as ROM, EPROM, EEPROM, non-volatile memory (NANDFLASH), solid-state drive (SSD)), etc.

[0164] The device embodiments described above are merely illustrative, wherein the units described as separate components may or may not be physically separated, and the components displayed as units may or may not be physical units, that is, they may be located in one place, or they may be distributed on multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the scheme of this embodiment. Ordinary technicians in this field can understand and implement it without paying creative labor.

[0165] Through the description of the above implementation methods, those skilled in the art can clearly understand that each implementation method can be implemented by means of software plus a necessary general hardware platform, and of course, it can also be implemented by hardware. Based on this understanding, the above technical solution is essentially or the part that contributes to the prior art can be embodied in the form of a software product, and the computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, a disk, an optical disk, etc., including a number of instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to execute the methods described in each embodiment or some parts of the embodiments.

[0166] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present application, rather than to limit it. Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. An image segmentation method, characterized in that: include: According to the attention-based residual structure, the feature extraction of the image to be segmented is performed to obtain the initial feature map; Performing multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map; Performing an attention-enhanced upsampling process on the intermediate feature map to obtain a segmented image; The step of performing an attention-enhanced upsampling process according to the intermediate feature map to obtain a segmented image comprises: Performing convolution, batch normalization, and neural activation processing on the intermediate feature map to obtain a fourth feature map; Performing transposed convolution, batch normalization, and neural activation processing on the fourth feature map to obtain a fifth feature map; The segmented image is obtained by performing upsampling processing in combination with the fifth feature map through spatial attention and / or channel attention.

2. The image segmentation method according to claim 1, characterized in that: The step of extracting features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map comprises: The image to be segmented is processed by convolution, batch normalization and neural activation to obtain the first feature map; Performing convolution and feature extraction based on a parameter-free attention mechanism on the first feature map to obtain a second feature map; Superimposing the image to be segmented and the second feature map to obtain a third feature map; An initial feature map is determined based on the third feature map.

3. The image segmentation method according to claim 2, characterized in that: The step of performing convolution on the first feature map and feature extraction based on a parameter-free attention mechanism to obtain a second feature map comprises: Performing convolution processing on the first feature map to obtain a convolution feature map; Determining an energy distribution function according to the convolution feature map; A second feature map is generated according to the energy distribution function and the convolution feature map.

4. The image segmentation method according to claim 3, characterized in that: The step of determining the energy distribution function according to the convolution feature map comprises: Determining a feature space mean according to the convolution feature map; Calculate the variance of the width and height of the convolution feature map in the channel direction according to the feature space mean; An energy distribution function of the convolution feature map is determined according to the feature space mean, the variance and the energy coefficient.

5. The image segmentation method according to claim 1, characterized in that: The step of reconstructing a feature map based on multi-head self-attention based on the initial feature map to obtain an intermediate feature map comprises: Performing layer normalization processing on the initial feature map to obtain a layer normalized feature map; Determining a query vector, a key vector, and a value vector for each vector in the layer normalized feature map, respectively; The layer normalized feature map is reconstructed according to the query vector, the key vector and the value vector of each of the vectors to obtain an intermediate feature map.

6. The image segmentation method according to claim 5, characterized in that: The step of reconstructing the layer normalized feature map according to the query vector, the key vector and the value vector of each of the vectors to obtain the intermediate feature map comprises: Determine the dot product result of the query vector and the corresponding key vector for each of the vectors; Normalizing each of the dot product results respectively to obtain multiple normalized results; Each of the normalized results is weighted respectively and added to the value vector of the corresponding vector to obtain a summed result; The summation results are summed to obtain an intermediate feature map.

7. An image segmentation device, characterized in that: include: An extraction module is used to extract features from the image to be segmented according to the attention-based residual structure to obtain an initial feature map; A reconstruction module, used for performing multi-head self-attention-based feature map reconstruction based on the initial feature map to obtain an intermediate feature map; An upsampling module, used for performing an attention-enhanced upsampling process on the intermediate feature map to obtain a segmented image; The upsampling module is specifically used to perform convolution, batch normalization and neural activation processing on the intermediate feature map to obtain a fourth feature map; perform transposed convolution, batch normalization and neural activation processing on the fourth feature map to obtain a fifth feature map; and perform upsampling processing in combination with the fifth feature map through spatial attention and / or channel attention to obtain a segmented image.

8. An electronic device comprising a processor and a memory storing a computer program, characterized in that: When the processor executes the computer program, the steps of the image segmentation method according to any one of claims 1 to 6 are implemented.

9. A computer storage medium, the computer storage medium being a computer-readable storage medium, comprising a computer program, characterized in that: When the computer program is executed by a processor, the steps of the image segmentation method according to any one of claims 1 to 6 are implemented.

Images