Image segmentation method based on lightweight recursive non-local self-attention

Abstract

The present invention relates to an image segmentation method based on lightweight recursive non-local self-attention, comprising the following steps: step S1: collecting portrait segmentation data, and pre-processing to obtain a portrait segmentation data set; step S2: constructing an image segmentation model based on lightweight recursive non-local self-attention; step S3: based on the portrait segmentation data set, using a binary cross entropy loss function and back-propagating the loss gradient, and training until the model converges according to a preset training learning rate, optimizer and learning rate strategy; step S4: normalizing the image to be tested and inputting it into the trained image segmentation model to obtain a portrait segmentation map. The present invention can effectively realize portrait segmentation, and the imaging effect is significantly improved.

Description

Technical Field

[0001] The present invention relates to the fields of computer vision and image processing, and in particular to an image segmentation method based on lightweight recursive non-local self-attention. Background Art

[0002] Computer vision methods based on deep learning generally extract local features based on convolution operations, and gradually expand the receptive field by stacking a certain number of convolutional layers to complete the transition from local shape and detail feature extraction to relatively high-level global and semantic information extraction. However, there are two obvious defects in using only convolutional layers for feature extraction: First, high-resolution input images need to pass through a sufficiently deep network to expand their theoretical receptive field, thereby capturing global or long-range feature information, which requires a huge amount of computational overhead; second, deep networks gradually expand the receptive field through convolutional layers, and do not use explicit methods to model long-range features, so they cannot well obtain and process long-range feature interdependencies, and have inherent shortcomings for targets or scenes with structured shapes and distributions. (Long-range feature: long-range feature)

[0003] Inspired by the non-local mean image denoising method in classic image processing, Facebook AI Research proposed the Non-local Neural Networks method, which uses the cross-correlation normalized attention of all valid pixel (feature) values ​​at any position in space to represent the weighted representation. This method can be used as a plug-and-play module to replace the convolution module, significantly improving the effect of video classification and demonstrating its versatility in object detection and instance segmentation. However, the bottleneck of this method lies in two points: (1) calculating the cross-correlation matrix of any pixel to all pixels will consume a huge amount of storage space and computation, and (2) weighting the attention on global pixels also requires a huge amount of storage space and computation.

[0004] Therefore, subsequent work has made trade-offs in reducing computational overhead. For example, the Attentional GAN ​​method uses the process of downsampling the feature map side length to 1 / K during inference, calculating non-local self-attention, and amplifying the weighted feature map. This method can reduce the amount of matrix multiplication calculation to (1 / K)^4, but the disadvantage is that lossy downsampling will cause the recovered detail information to be also lost; for example, the SAN method divides the global feature map into K*K small blocks and stacks them to the batch dimension, so that only non-local self-attention is calculated in each block, and finally the small blocks are reassembled from the batch dimension to the spatial dimension. This method can reduce the amount of matrix multiplication calculation to (1 / K)^2. When the task is high-resolution input such as super-resolution, a larger value of K can reduce the amount of calculation considerably, but the disadvantage is that the block size of the block operation (i.e., SIZE / K) needs to be manually specified, and (in other tasks) the idea of ​​block operation is contradictory to non-local, which is a compromise solution that loses long-range dependencies beyond the block scale.

[0005] Recursive modules have also emerged in recent years, such as the ResidualinResidual Dense Block (RRDB) sub-module proposed by the ESRGAN method, which uses a Residual Dense Block (RDB) as the basic module, stacks three RDBs and uses the residual structure again to form a recursive RRDB module; for example, the U^2Net method proposes to use U-Net directly as the basic module, and then use U-Net modules to combine into a U-Net architecture, forming a UNet-I n-UNet structure. Summary of the invention

[0006] In view of this, the purpose of the present invention is to provide an image segmentation method based on lightweight recursive non-local self-attention, which can effectively realize portrait segmentation and significantly improve the imaging effect.

[0007] To achieve the above object, the present invention adopts the following technical solution:

[0008] A lightweight recursive non-local self-attention based image segmentation method comprises the following steps:

[0009] Step S1: collecting portrait segmentation data and preprocessing it to obtain a portrait segmentation data set;

[0010] Step S2: constructing a lightweight recursive non-local self-attention based image segmentation model;

[0011] Step S3: Based on the portrait segmentation dataset, a binary cross entropy loss function is used and the loss gradient is back-propagated, and training is performed until the model converges according to the preset training learning rate, optimizer and learning rate strategy;

[0012] Step S4: normalize the image to be tested and input it into the trained image segmentation model to obtain a portrait segmentation map.

[0013] Furthermore, the step S1 is specifically as follows:

[0014] Step S11: Collecting a portrait segmentation data set including an RGB three-channel image as input and a {0,1} binary segmentation map as its corresponding output target, where 0 represents the background class and 1 represents the target class;

[0015] Step S12: By performing random cropping, scaling, mirroring, rotation and input value normalization processing on the input and output, the augmented image data is obtained, which is the portrait segmentation data set.

[0016] Furthermore, the image segmentation model based on lightweight recursive non-local self-attention includes an infrastructure and feature extraction network, a separable non-local self-attention module and an output feature mapping module.

[0017] Furthermore, the basic architecture and feature extraction network are specifically as follows: a feature pyramid is used as the basic architecture, including a skeleton network and a multi-scale upsampling fusion layer; the skeleton network adopts a ResNeSt structure, and the multi-scale upsampling fusion layer is a top-down structure, which is implemented by upsampling, feature map point addition and convolution layer, and the feature layer of a certain stage of FPN is used as the output feature layer.

[0018] Furthermore, the separable non-local self-attention module specifically models long-range features in an explicit manner, obtains and processes long-range feature interdependencies, further integrates the whole-image level context features through the infrastructure and feature extraction network, and sends the fused and reorganized features to the output feature mapping module.

[0019] Furthermore, the output feature mapping module is a sequential module composed of basic blocks stacked by convolution layers, batch normalization, and non-linear activation functions, and finally uses a layer of convolution to map the number of channels to 1 to obtain an intermediate result image, and uses bilinear interpolation to enlarge it to the input size, and finally uses the Sigmoid function to map the result to the [0,1] interval.

[0020] Compared with the prior art, the present invention has the following beneficial effects:

[0021] The present invention can effectively realize portrait segmentation and significantly improve the imaging effect. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 It is the network structure of a traditional non-local self-attention module in one embodiment of the present invention;

[0023] Figure 2 is a diagram illustrating a network structure in one embodiment of the present invention;

[0024] Figure 3 is a schematic diagram of expansion and sub-block selection in one embodiment of the present invention;

[0025] Figure 4 and Figure 5 This is a test effect diagram of the portrait segmentation task in the present invention. DETAILED DESCRIPTION

[0026] The present invention will be further described below in conjunction with the accompanying drawings and embodiments.

[0027] This paper proposes a solution to the huge computational overhead of traditional non-local self-attention, whose matrix multiplication computation is

[0028] OP(SA)=(c1+c2)·(H·W) 2 …(1)

[0029] Where c1 is the number of feature channels for cross-correlation calculation, c2 is the number of feature channels for the final output, and HW is the size of the input feature map. It can be seen that a large amount of calculation is generated by the spatial dimension of the feature map.

[0030] Considering that a certain number of convolutional layers can make any pixel (feature) have a certain receptive field, by fixing a specific step size (SH, SW), we can extract (SH, SW) sub-feature maps of size (H / SH, W / SW), stack the sub-maps in the batch dimension for global self-attention calculation, and finally reassemble the sub-maps from the batch dimension to the spatial dimension. We can obtain the incomplete dependence of any pixel on the global. This step is named Dilated-NonLocal-SelfAttn. The matrix multiplication calculation amount of the non-local self-attention generated by this step is

[0031]

[0032] On the reorganized feature map, the adjacent (SH, SW) sub-blocks have obtained the information of all global positions. At this time, according to the adjacent sub-blocks as sub-feature maps, (H / SH, W / SW) (SH, SW) sub-feature maps can be obtained. The sub-maps are stacked in the batch dimension for global self-attention calculation. Finally, the sub-maps are reassembled from the batch dimension to the spatial dimension. The complete global dependency after fusion can be sampled through the neighboring features. This step is named sub-block non-local self-attention (SubBlock-NonLocal-SelfAttn). The matrix multiplication calculation amount of non-local self-attention generated by this step is

[0033]

[0034] If DSBSA (Dilated-SA+SubBlock-SA) is used to replace SA, then

[0035]

[0036] According to the mean inequality, we can get

[0037]

[0038] For example, when HW=64x64, the optimal OP can be reduced to 1 / 32 of the original, and when HW=128x128, the optimal OP can be reduced to 1 / 64 of the original

[0039] Traditional non-local self-attention belongs to additive attention, which uses the residual method to add the attention feature map obtained by the weighted combination of the globally acquired cross-correlation normalized attention to the residual feature map as compensation for the original branch; when the multiplicative attention is used, better performance will be obtained. The attention feature map itself after adding the Sigmoid activation function is multiplied as an attention to the original branch feature map, forming a recursive structure that generates attention by weighted combination of global attention. Furthermore, a channel attention (CA) layer is added to the basis of DSBSA to enhance the capture of the importance of the channel by the entire attention module. Similarly, the residual is changed to Sigmoid multiplication, so that the entire DSBSA+CA module is transformed into a module that uses attention to generate attention.

[0040] Considering that grouped convolution can reduce the amount of calculation while enabling the feature layers of different groups to learn as rich features as possible before finally fusing them, the original convolutional layers are grouped for calculation, and then the cross-correlation features when calculating attention are split along the channel direction and stacked in the batch dimension for calculation. Finally, the channels are split and reorganized by the batch dimension. Although the amount of calculation of the matrix product does not change, multiple groups of attention combination feature maps can be obtained, and the amount of calculation of the convolutional layer inside the module is significantly reduced.

[0041] Taking the portrait segmentation task as an example, this embodiment provides an image segmentation method, comprising the following steps:

[0042] Step S1: collecting portrait segmentation data and preprocessing it to obtain a portrait segmentation data set;

[0043] Step S2: constructing a lightweight recursive non-local self-attention based image segmentation model;

[0044] Step S3: Based on the portrait segmentation dataset, a binary cross entropy loss function is used and the loss gradient is back-propagated, and training is performed until the model converges according to the preset training learning rate, optimizer and learning rate strategy;

[0045] Step S4: normalize the image to be tested and input it into the trained image segmentation model to obtain a portrait segmentation map.

[0046] In this embodiment, a portrait segmentation dataset VIP (Video instance-level Parsing) is selected, which contains portrait videos of 404 different scenes, sampled every 25 frames, totaling 18469 data pairs. Specifically, it includes an RGB three-channel image as input and a {0,1} binary segmentation map as its corresponding output target, where 0 represents the background class and 1 represents the target class (i.e., the human body area). By performing random cropping, scaling, mirroring, rotation, input value normalization and other operations on the input and output, a large amount of augmented image data is obtained, which is a portrait segmentation dataset. The augmented data satisfies the input size of (3,360,480) and the output size of (1,360,480).

[0047] In this embodiment, the image segmentation model based on lightweight recursive non-local self-attention includes an infrastructure and feature extraction network, a separable non-local self-attention module and an output feature mapping module.

[0048] In this embodiment, the infrastructure and feature extraction network are specifically as follows: a feature pyramid is used as the infrastructure, including a skeleton network and a multi-scale upsampling fusion layer; the skeleton network can select any classification pre-trained network, and the ResNeSt structure is used in this embodiment, and the last residual block is replaced by a dilated convolution with a step size of 2. The multi-scale upsampling fusion layer is consistent with the original FPN, using 2x2 bilinear upsampling, feature map point addition, and a convolution kernel of 3 with a step size of 1 from top to bottom. Finally, the feature layer with a downsampling step size of 4 for FPN is used as the output feature layer.

[0049] In this embodiment, the separable non-local self-attention module specifically models long-range features in an explicit manner, obtains and processes long-range feature interdependencies, further integrates the whole-image level context features through the infrastructure and feature extraction network, and sends the fused and reorganized features to the output feature mapping module.

[0050] In this embodiment, the output feature mapping module is based on the separable non-local self-attention module, and adds two basic blocks consisting of 3x3 convolution with a convolution kernel of 3 and a step size of 1, batch normalization, and a nonlinear activation function ReLU. Finally, the number of channels is mapped to 1 using a 3x3 convolution with a convolution kernel of 3 and a step size of 1, and an output size of (1,90,120) is obtained. The result is amplified by 4 times using bilinear interpolation to obtain a (1,360,480) result image that matches the input, and finally the Sigmoid function is used to map the result to the [0,1] interval.

[0051] The final structure can be expressed as the following equation:

[0052] Out=Sigmoid(Up(Fm(DSBSA(FPN(In)))))

[0053] Among them, In / Out represent input and output respectively, FPN represents feature extraction network, DSBSA represents separable non-local self-attention module, Fm (Feature Mapping) represents output feature mapping module, Up represents upsampling the output prediction to the original input size, and Sigmoid is the final activation function used.

[0054] In this embodiment, step S3 specifically compares the [0,1] interval segmentation result output by the output feature mapping module with the corresponding {0,1} binary segmentation map in step S1, uses the binary cross entropy loss function (BCE) and back-propagates the loss gradient. The training learning rate is 2e-4, and the Adam optimizer and cosine annealing strategy are used, and the training is repeated for 50k times until the model converges.

[0055] In this embodiment, step S4 is specifically as follows: inputting the RGB image to be tested, scaling it accordingly according to the input size (without affecting the result generation), and performing numerical normalization, and sending it to the above network to perform feature extraction, output mapping, upsampling and output value normalization according to formula (6), and finally obtaining a continuous prediction map, that is, a portrait segmentation map, which can be further binarized to obtain a {0,1} segmentation map.

[0056] Embodiment 1:

[0057] In this embodiment, after the input feature map size is determined, the module first searches for the optimal sparsification hyperparameters, with the square root of the size as the center and searching toward both sides. When the size is divided by the current search value, if the redundant or supplementary boundaries are less than the selected threshold, the hyperparameter search ends, otherwise the search continues. Figure 3 As shown, the same color represents all sampling positions in the same sampling sub-feature map. For example, a 6x6 feature map will be divided into an expansion and sub-block feature sub-map with SH=SW=2.

[0058] After the search is completed, the module is trained or reasoned, such as Figure 2 As shown, the sub-feature maps are stacked on the batch dimension (split operation), and then the following is performed Figure 1 The traditional non-local self-attention calculation shown in FIG. 1 finally reorganizes the sub-graph from the batch dimension to the spatial dimension (merge operation), that is, completes the calculation of the expansion / sub-block non-local self-attention.

[0059] The specific amount of calculation is compared in Table 1. When the input feature map resolution is 32, 64, 96, and 128, the optimal theoretical value of the matrix multiplication calculation after optimization should be 1 / 16, 1 / 32, 1 / 48, and 1 / 64 of the original version. After optimization, it is 1 / 11, 1 / 32, 1 / 42, and 1 / 58. The difference is caused by the root-root of the feature map size. The difference in convolution and other calculations at the same input resolution lies in the convolution module added between the two steps of separable attention.

[0060] Table 1

[0061]

[0062]

[0063] like Figure 2 As shown in the Sigmoid-multiply (dashed line) in the figure, the feature map that completes the self-attention calculation is activated by the Sigmoid function and then multiplied with the original residual branch to replace the additive structure in the original version. In addition, there is a set of Conv+BN+PReLU before and after each self-attention module, where the Conv kernel size is 1x1 and only the channel fusion is performed. Finally, a set of channel attention (CA) modules (GAP is the global average pooling operation, FC is the fully connected layer) is used to further extract the importance of each channel ignored by SA, and the Sigmoid function is also used to activate it and then multiply it with the residual branch initially input to the module to complete the calculation of the entire module.

[0064] The convolutional layers in the first and second Conv+BN+PReLU groups are replaced by grouped convolution, and the features for calculating cross-correlation are split, calculated, and recombined along the channel direction, thus completing grouped parallel attention.

[0065] The specific amount of computation is compared in Table 2. When the input resolution is the same, the use of grouped parallel attention does not change the amount of computation of the matrix multiplication part, while grouped convolution can significantly reduce the amount of convolution computation. When the number of groups is 2, it is roughly reduced to 3 / 5 of that of the ungrouped one, and when the number of groups is 4, it is roughly reduced to 2 / 5 of that of the ungrouped one.

[0066] Table 2

[0067]

[0068] Figure 5The figure shows the test effect of adding a separable self-attention module before the FPN output layer. It can be seen that adding this module can greatly reduce the false detection of background and difficult areas, indicating that the module can obtain effective information from different areas of the feature map to improve the overall feature representation ability, and it is a lightweight module that can significantly improve the effect without stacking a large number of modules. Table 3 is a comparison chart of the verification indicators of FPN (+DSBSA) in the VIP subset of the portrait segmentation dataset, where PixelAcc represents the pixel-level accuracy of the pixel-by-pixel comparison between the output result and the true annotation, IOU represents the ratio of the set intersection part to the set union part between the output result and the true annotation, and Binary means that the output result is binarized, otherwise it is a continuous output value. It can be seen that in various indicators, only adding a separable self-attention can achieve significant improvement. Figure 4 is a comparison chart of indicators in the verification subset of the human detection dataset CrowdHuman. We used the classic test scheme on the COCO dataset to test the effect of a small model. It can be seen that separable self-attention is also effective in detection tasks and can be further extended to other scenarios and tasks in the field of computer vision.

[0069] Table 3

[0070] Model Methods PixelAcc IOU PixelAcc-Binary IOU-Binary Inseparable attention 98.202％ 92.841％ 98.552％ 94.202％ Separable self-attention 98.418％ 93.663％ 98.715％ 94.825％

[0071] Table 4

[0072]

[0073] The above description is only a preferred embodiment of the present invention. All equivalent changes and modifications made according to the scope of the patent application of the present invention should fall within the scope of the present invention.

Claims

1. A lightweight recursive non-local self-attention based image segmentation method, characterized in that: The following steps are involved: Step S1: collecting portrait segmentation data and preprocessing it to obtain a portrait segmentation data set; Step S2: constructing a lightweight recursive non-local self-attention based image segmentation model; The lightweight recursive non-local self-attention based image segmentation model includes a basic architecture and feature extraction network, a separable non-local self-attention module and an output feature mapping module; The basic architecture and feature extraction network are specifically as follows: a feature pyramid is used as the basic architecture, including a skeleton network and a multi-scale upsampling fusion layer; the skeleton network adopts a ResNeSt structure, and the multi-scale upsampling fusion layer is a top-down structure, which is implemented by upsampling, feature map point addition and convolution layer, and a feature layer of a certain stage of FPN is used as the output feature layer; The separable non-local self-attention module specifically models long-range features in an explicit manner, obtains and processes long-range feature interdependencies, further fuses the whole-image level context features through the infrastructure and feature extraction network, and sends the fused and reorganized features to the output feature mapping module; The output feature mapping module is a sequential module composed of a stack of basic blocks consisting of a convolution layer, batch normalization, and a nonlinear activation function. Finally, a layer of convolution is used to map the number of channels to 1 to obtain an intermediate result map, and bilinear interpolation is used to enlarge it to the input size. Finally, a Sigmoid function is used to map the result to the interval [0,1]; Step S3: Based on the portrait segmentation dataset, a binary cross entropy loss function is used and the loss gradient is back-propagated, and training is performed until the model converges according to the preset training learning rate, optimizer and learning rate strategy; Step S4: normalize the image to be tested and input it into the trained image segmentation model to obtain a portrait segmentation map.

2. The image segmentation method based on lightweight recursive non-local self-attention according to claim 1, characterized in that: The step S1 is specifically as follows: Step S11: Collecting a portrait segmentation data set including an RGB three-channel image as input and a {0,1} binary segmentation map as its corresponding output target, where 0 represents the background class and 1 represents the target class; Step S12: By performing random cropping, scaling, mirroring, rotation and input value normalization processing on the input and output, the augmented image data is obtained, which is the portrait segmentation data set.

Images