Image reconstruction method based on multi-window cross feature fusion attention mechanism

Abstract

The invention discloses an image reconstruction method based on a multi-window cross feature fusion attention mechanism. According to the method, the resolution of an original image is reduced through bicubic interpolation to obtain model training data, and the original image is used as a sample label. For training data, shallow layer features are extracted through blueprint separable convolution; and improving an efficient separable distillation module to obtain a plurality of attention fusion distillation modules, and extracting deep features. A feature fusion module based on multi-window cross attention is used to carry out operations such as multiple times of segmentation, attention calculation, multi-window cross attention calculation, recombination and the like on deep features, and feature enhancement is realized. And fusing the features with shallow layer features, and outputting a reconstructed high-resolution image. Models with different reconstruction multiplying powers can be obtained according to different sampling multiplying powers under training data.

Description

Technical Field

[0001] The present invention belongs to the field of computer vision, and relates to image feature extraction and image super-resolution technology, and specifically to an image reconstruction method based on a multi-window cross-feature fusion attention mechanism. Background Art

[0002] In traditional image processing, images are often blurred and distorted due to the limitations of device resolution and acquisition conditions. In order to improve image clarity and quality, super-resolution technology was introduced, whose goal is to restore low-resolution, degraded images to corresponding high-resolution images. Super-resolution technology can not only improve image quality and enhance people's visual experience, but also assist the model's target detection task and improve its accuracy. In addition, super-resolution technology can achieve the same effect while reducing costs, breaking through the limitations of optical imaging systems and obtaining high-resolution images. With the continuous innovation of digital image processing technology, super-resolution has become one of the research directions that has attracted much attention in the field of image processing. Image super-resolution technology has been widely used in video surveillance, medical imaging and other fields.

[0003] In the process of image super-resolution, it is necessary to extract and fuse deep features, and generate a reconstructed image close to the original high-resolution image based on the extracted features while maintaining the accuracy of the original image information. A more common image super-resolution method is to use the Transformer model with an attention mechanism, which captures the global information of the image based on the image features and can better understand the relationship between different areas in the image, especially for tasks involving long-distance relationships. However, the defects of these methods are: first, they require a lot of computing resources, and second, they do not pay enough attention to local areas, which reduces the performance of image super-resolution. Summary of the invention

[0004] In view of the shortcomings of the prior art, the present invention proposes an image reconstruction method based on a multi-window cross-feature fusion attention mechanism. The square window of the traditional window attention mechanism is further divided into overlapping small windows of various shapes, which not only fully considers the context information of the local area, but also further reduces the amount of calculation, thereby achieving lightweight and high-precision models.

[0005] The image reconstruction method based on multi-window cross-feature fusion attention mechanism specifically includes the following steps: Step 1: Data collection and processing Step 1.1: Collect high-resolution images of different scenes, including architectural images, face images, cartoon images, etc., reduce the image resolution through the bicubic interpolation method, and form image pairs with the original high-resolution images to obtain a diverse data set.

[0006] Step 1.2: Crop the image pair obtained in step 1.1 into multiple slices of the same size to enhance the data and expand the data set, and then convert the slices into Tensor vectors. The high-resolution image in an image pair is used as a label, and the low-resolution image is used as a training sample.

[0007] Step 2: Shallow feature extraction First, extract the shallow features of the sample image. The specific steps are as follows: Step 2.1: Perform channel expansion on the Tensor vector of the training sample in step 1.2, and expand its channel number from 3 channels to 12 channels to obtain the corresponding feature image tensor.

[0008] Step 2.2: Use Blueprint Separable Convolutions (BSConv) to extract features from the feature image tensor in step 2.1. First, use the deep convolution layer to perform independent convolution operations on 12 channels, and independently calculate the convolution results to obtain a tensor of the same size as the input. Then, use the point-by-point convolution layer to use a 1*1 convolution kernel to integrate the information of each channel of the tensor to obtain the shallow feature F. in ,The shallow features preserve the information of the image mapped from low-dimensional features to high-dimensional features.

[0009] Step 3: Deep feature extraction After steps 1 and 2, the shallow features F in the low-resolution image are initially extracted. in By improving the efficient separable distillation block (ESDB), multiple attention fusion distillation blocks (AFDB) are obtained to extract deep information. The attention fusion distillation module includes a distillation layer, a cohesive layer, and an enhancement layer. The specific steps are as follows:

[0010] Step 3.1: Feature extraction The shallow feature F obtained in step 2 in The distillation layer includes L levels, each of which includes two branches, for extracting information learned by the integration model.

[0011] In the first layer, the first branch learns complex patterns and features through ordinary convolutional layers and Gaussian Error Linear Unit (GELU) activation function layers to obtain F d1The second branch introduces a deeper feature representation while maintaining the original information through the blueprint separable convolution layer, GELU activation layer and residual connection of the blueprint shallow residual block (BSRB), and obtains F r1 :

[0012] (1) (2) (3) (4) in, , when l=1, is the shallow feature F in , is the feature map of the input activation layer.

[0013] Step 3.2: Feature aggregation The L outputs of the distillation layer are concatenated on the channel to obtain the concatenated feature map F ct Then, feature condensation is achieved through the convolution layer to obtain F con : (5) Among them, CONCAT() represents channel concatenation, and Conv1() represents a 1*1 convolutional layer.

[0014] Step 3.3: Feature enhancement The enhancement layer enhances the feature map obtained in step 3.2 through three attention layers. Through the adjusted structure of the three attention layers, the three attention layers can make up for each other's defects and obtain accurate image features.

[0015] First, the feature map F con The input is sent to the ECA (Efficient Channel Attention) attention module, which focuses on the information in the channel dimension, considers the correlation between different channels, and improves the efficiency of the network in learning specific features. First, the global information of each channel is obtained through one-dimensional global average pooling (GAP). , and then obtain the learnable weights of each channel through a fully connected layer , the output F is calculated based on the weight of each channel eca :

[0016] (6) (7) in , , They are feature maps F con length, width and channel, Represents the feature map F con The pixel value at position (i, j) in the cth channel, σ represents the activation function, represents the convolution operation, Represents the output feature map F eca The pixel value at position (i,j) in the cth channel.

[0017] At the same time, the feature map F con Input to the ESA (Efficient Spatial Attention) attention module, which focuses on information in the spatial dimension and improves the perception of local information by learning the relationship between positions, and obtains the output F esa : (8) (9) in Represents the feature map F con The spatial information at position (i, j) in is the weight matrix obtained through learning, Represents the output feature map F esa The pixel value at position (i,j) in the cth channel.

[0018] Finally, F eca and F esa After splicing on the channel, it is input into the CCA (Compact Channel Attention) attention module, and the channel relationship is modeled using low-rank matrix decomposition, which reduces the number of parameters and computational complexity, and the output F is obtained. cca Finally, a 1*1 convolution is performed to obtain the output F after feature enhancement. eo1 .

[0019] Step 3.4: The enhanced output F eo1 As the input of the next attention fusion distillation module, after repeated multiple times, the output features of each attention fusion distillation module are spliced ​​in the channel dimension, and then the integrated aggregate features are output through the convolution layer, and finally the deep features F in the low-resolution image are obtained through GELU activation. ac .

[0020] Step 4: Multi-window cross attention Design a feature fusion module MWT based on multi-window cross attention to further enhance the deep feature F ac .

[0021] Step 4.1: First, the deep feature F ac Input the ESA attention module and divide the feature map output by the ESA attention module into a set of feature maps with c channels, h*w size and no overlap.

[0022] Step 4.2 inputs the segmented feature map into MHA (Multi-Head Attention) to obtain the attention map, and then passes it through the linear layer to generate three tensors: Q (Query), K (Key), and V (Value).

[0023] Step 4.3: Restore the tensors Q, K, and V output by the linear layer to the image shape from the dimension (b, n, h*w, c / / n) to (b, n, h, w, c / / n), where b is the batch size and n is the number of attention heads of the multi-head attention module (MHA). Perform the following three window-cutting processes on the image shape restored from the tensors Q, K, and V:

[0024] ① The image is divided into t blocks along the longitudinal direction, and t blocks (Q h ,K h ,V h ); ② The image is divided into t blocks along the horizontal direction, and t blocks (Q w ,K w ,V w ); ③ Divide the entire image into t blocks evenly, and get t blocks (Q c ,K c ,V c ).

[0025] Perform position encoding on each block after segmentation, add the result of position encoding to the corresponding block, and get 3×t groups (q, k, v). Calculate the attention value O of each block respectively: (10) Among them, d k The dimension size of k after adding the position encoding.

[0026] Step 4.4, then reorganize the attention calculation results according to the window slicing method to obtain the attention feature maps of the three windows. After merging the three attention feature maps on the channel, a convolution is performed to integrate the feature information on the three window channels to obtain an output feature map of size h*w. Reorganize multiple output feature maps at the corresponding positions of the window to obtain the size and deep features Fac The same output can be operated by sliding the window (Shift Window, SWIND).

[0027] Step 4.5: Repeat steps 4.2 to 4.4, and then use an ESA attention module to complete the feature enhancement part and get the output F wo . For output F wo After a blueprint separable convolution, it is combined with the shallow feature F in Perform residual connection and then output high-resolution image after upsampling.

[0028] Step 5: Image super-resolution Set the parameters of bicubic interpolation and train the models at different magnifications. Select the model with the corresponding magnification according to the needs, input the low-resolution image into the model, and obtain the super-resolution image through feature extraction and enhancement in steps 2 to 4.

[0029] The present invention has the following beneficial effects: 1. In the process of deep image feature extraction, in order to capture the image features in space and channels and the correlation features between them, a new feature enhancement method is introduced based on the spatial attention mechanism and the channel attention mechanism. By dispersing the information in the ESA channel dimension and the ECA spatial dimension to each channel, and then using the channel attention mechanism CCA to integrate the enhanced information, by adjusting the module structure position, it can effectively utilize the advantages of each module in efficiently capturing features in space and channels. At the same time, it solves the shortcomings of other structural attention modules in the lack of correlation between image features in channels and space, and further enhances the feature information contained in the feature image after feature extraction.

[0030] 2. In the process of feature fusion, in order to make the features of each part clearly identifiable and obtain complete and important feature information, an improved feature fusion module based on multi-window crossover is proposed. First, the attention to the features on different channels is adjusted through channel attention, and then the feature image is divided into windows, and each window is divided into three parts: vertical, horizontal, and rectangular for attention. The three parts are fused and then the feature fusion part is enhanced through the channel attention module, so that which features are more important can be noticed in the process of upsampling the image. Because the multi-window mechanism is introduced, the calculation amount is dispersed to each window, and the features of multiple windows of different shapes can be captured, so while reducing the amount of calculation, the accuracy of the model super-resolution is improved, providing an important reference for the generation of the final super-resolution image. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] Figure 1 This is a flow chart of the image reconstruction method based on multi-window cross-feature fusion attention mechanism; Figure 2 is the original high-resolution image used in the embodiments; Figure 3 is a low-resolution image obtained by bicubic interpolation in the embodiment; Figure 4 Schematic diagram of the image reconstruction model based on the attention mechanism of multi-window cross-feature fusion; Figure 5 This is a schematic diagram of the attention fusion distillation module structure; Figure 6 Schematic diagram of multi-window cross-feature fusion attention mechanism; Figure 7 It is a schematic diagram of a multi-window cross MWA module; Figure 8 It is the super-resolution image output in the embodiment. DETAILED DESCRIPTION

[0032] The present invention will be further explained below with reference to the accompanying drawings; like Figure 1 As shown in the figure, the image reconstruction method based on multi-window cross-feature fusion attention mechanism has the following specific steps: Step 1: Data collection and processing This example uses 2650 images from Flickr2K and 800 images from DIV2K as training datasets. Figure 2 The high-resolution image of size 256×256 shown in the figure is downsampled by 3 times through the bicubic interpolation method to obtain an image of size 85×85, and then bicubic interpolation is performed to restore the size to 256×256, as shown in the figure. Figure 3 The two 256×256 images in the image pair are cut into 32 64×64 images to form 16 groups of image pairs to increase the generalization of the model training images. The cut low-resolution images are used as model training samples, and the high-resolution images are used as sample labels for .

[0033] Step 2: Shallow feature extraction like Figure 4 As shown in the figure, the training sample is converted from a 64×64 image to a tensor of size (1, 3, 64, 64), and copied to 12 channels. Concat is used to merge and expand on the channel to obtain a tensor of size (1, 12, 64, 64), which can enhance the color information of the low-resolution image and capture the details of the image to reduce the loss of image information. Finally, the blueprint separable convolution layer is used to extract shallow features, increase the model complexity to enhance the feature representation, and obtain a shallow feature map F of size (1, 64, 64, 64). in .

[0034] Step 3: Image deep feature extraction For the shallow feature map F obtained in step 2 in , through 8 consecutive attention fusion distillation modules AFDB, and then splicing the 8 results in the channel dimension to obtain a deep feature image F of size (1, 512, 64, 64) ac .

[0035] like Figure 5 As shown in Figure 1, the attention fusion distillation module AFDB includes three distillation layers, a cohesive layer and an enhancement layer. Each distillation layer includes two branches. The first branch outputs a feature map F through a common convolution layer and a Gaussian error activation function layer. d1 The second branch outputs the feature map F through the blueprint shallow residual module r1 , and input into the two branches of the next layer. Each layer can continuously refine the features of the feature image. The first layer will reduce the number of feature image channels by half, and obtain a feature map of size (1, 32, 64, 64). The size of the feature map output by the next two layers will not change, and is (1, 32, 64, 64). The optimal blueprint shallow residual module outputs the feature map F r3 After a convolutional layer of size 3*3, the feature map F obtained by the first three layers d1 After merging, a feature map tensor of size (1, 128, 64, 64) is obtained, and then the number of channels is restored to a feature image of size (1, 64, 64, 64) through a common convolution layer to complete feature condensation. Finally, in the enhancement layer, three attention layers ECA, ESA, and CCA are used to enhance the features in space and channels to complete the deep feature extraction of the image. These three attention layers will not change the size of the image.

[0036] Step 4: Image feature fusion and reconstruction like Figure 6 As shown, the deep feature image F ac The input convolutional layer and ESA attention layer are processed to obtain a feature image of size (1, 64, 64, 64). The multi-window cross-feature fusion attention mechanism is used for this feature map. First, the feature map is divided into 4 windows with a length and width of 16 according to the window to obtain 4 feature maps of size (1, 64, 16, 16). After each divided window passes through the MHA and linear modules, it is sent to MWA (Multi-Window Attention), as shown in Figure 7As shown in Figure 1, in MWA, q, k, and v are split so that each window is divided into 4 windows of size (8, 8), 4 windows of size (4, 16), and 4 windows of size (16, 4). Finally, the attention in the three directions is fused to obtain a feature image of size (1, 64, 64, 64) after one attention. In order to make the attention pay attention to adjacent features, the window is translated, and then the translated image is processed by the multi-window cross attention mechanism. Then, the final feature image is obtained through an ESA attention layer, and finally, it is residually connected with the initial low-resolution image and up-sampled again to obtain the reconstructed high-resolution image, as shown in Figure 1. Figure 8 shown.

[0037] The computational complexity of using a fixed-size window for segmentation and using the multi-window cross-feature fusion method for segmentation is calculated separately. , : (12) (13) Where h, w, c are the height, width, and number of channels of the window, and n is the number of heads of multi-head attention. Obviously, in this embodiment, the computational complexity of this method is reduced by 1 / 4 compared with the traditional fixed-size window.

[0038] Set the loss function to: (14) Set the learning rate to 1×10 -3 , set betas to 0.9 and 0.999, and set the adjustment parameters to decrease by 1×10 every 100 rounds -5 , train for a total of 500,000 rounds and save the model with the best training effect.

[0039] In order to verify the enhancement effect of the proposed multi-window cross-feature fusion attention mechanism on the model, relevant comparative experiments were carried out. The specific experimental results are shown in the following table:

[0040] The super-resolution ratio is achieved by changing the downsampling multiple of the bicubic interpolation. According to the experimental results, the PSNR (peak signal-to-noise ratio) and SSIM (structural similarity) indicators of this method at different super-resolution ratios (×2, ×3, ×4) are better than those of the baseline models BSRN, RFDN and LAPAR-A. Especially at ×2 and ×3 ratios, the improvement of this method compared with the baseline model is obvious and the performance is more outstanding. In addition, this method shows high PSNR and SSIM values ​​on different datasets (Set5, Set14, BSD100, Urban100, Manga109), indicating that it can achieve significant improvements in various scenarios and image types, especially on the Urban100 and Manga109 datasets.

[0041] On the Urban100 dataset, the triple super-resolution model obtained by this method has a PSNR improvement of 1.48dB. At the same magnification, Manga109 has an improvement of 0.6dB. The two datasets are composed of 100 city images and 109 comic images. Among them, the images in the Manga109 dataset have a certain artistic style, and the color information of comics is richer than that of real-world images, which shows that the method can better handle color relationships and ensure that the image is not too distorted. The performance of this method on the Urban100 dataset shows that this method can effectively capture the relationship between houses in the city, handle the structural information of the house, and has structural advantages.

[0042] In addition, the parameters and computational complexity of this method are relatively low, showing a high balance between performance and efficiency. Compared with RFDN and LAPAR-A, this method has a relatively reasonable demand for resources while improving performance. From the consistency of PSNR and SSIM at various magnifications, this method shows relatively stable performance and has good adaptability to images at different super-resolution magnifications.

[0043] The above experimental results prove that this method has achieved comprehensive advantages in super-resolution tasks, with better image reconstruction effect, higher performance and efficiency balance, wider adaptability and better stability.

Claims

1. Image reconstruction method based on multi-window cross-feature fusion attention mechanism, Features: The specific steps include: Step 1: Data collection and processing Collect high-resolution images from different scenes, reduce the image resolution through bicubic interpolation, and form image pairs with the original high-resolution images; convert the images into Tensor vectors, where the high-resolution images are used as labels and the low-resolution images are used as training samples; Step 2: Shallow feature extraction The Tensor vector of the training sample is expanded to 12 channels, and then input into the blueprint separable convolution to obtain the shallow feature F in ; Step 3: Deep feature extraction Use the attention fusion distillation module to extract deep information; the attention fusion distillation module includes a distillation layer, a cohesive layer and an enhancement layer. The specific steps are as follows: Step 3.1: Feature extraction The shallow feature F obtained in step 2 in The distillation layer includes L levels, each of which includes two branches. In the first layer, the first branch learns complex patterns and features through a common convolutional layer and a Gaussian error activation function layer GELU to obtain F d1 The second branch uses the blueprint separable convolution layer, GELU activation layer and residual connection to introduce a deeper feature representation while maintaining the original information to obtain F r1 : (1) (2) (3) (4) in, , when l=1, =F in , is the feature map of the input activation layer; Step 3.2: Feature aggregation The L outputs of the distillation layer are concatenated on the channel, and then feature condensation is achieved through the convolution layer to obtain the feature map F con Step 3.3: Feature enhancement At the same time, the feature map F con Input into the efficient channel attention module ECA and the efficient spatial attention module ESA, and then the feature map F output by these two attention modules eca and F esa The concatenation is performed on the channel and then input into the compact channel attention module CCA to obtain the output feature map F cca Finally, a 1*1 convolution is performed to obtain the output F after feature enhancement. eo1 ; Step 3.4: Output F after feature enhancement eo1 As the input of the next attention fusion distillation module, after repeated multiple times, the output features of each attention fusion distillation module are spliced ​​in the channel dimension, and then the integrated aggregate features are output through the convolution layer, and finally the deep features F of the low-resolution image are obtained through GELU activation. ac ; Step 4: Multi-window cross attention Design a feature fusion module MWT based on multi-window cross attention, and further enhance the deep feature F according to the following steps ac : Step 4.1: Deep feature F ac Input the efficient spatial attention module ESA, and divide the output feature map into a set of feature maps with a channel number of c, a size of h*w, and no overlap; Step 4.2: Input the segmented feature map into the multi-head attention module MHA and the linear layer to generate three tensors Q, K, and V; Step 4.3: Restore the tensors Q, K, and V to the image shape, and then perform the following three window cutting processes: ① The image is divided into t blocks along the longitudinal direction, and t blocks (Q h ,K h ,V h ); ② The image is divided into t blocks along the horizontal direction, and t blocks (Q w ,K w ,V w ); ③ Divide the entire image into t blocks evenly, and get t blocks (Q c ,K c ,V c ); Perform position encoding on each block after segmentation, add the result of position encoding to the corresponding block, and get 3×t groups (q, k, v). Calculate the attention value O of each block respectively: （5） Among them, d k The dimension size of k after adding the position encoding; Step 4.4: Reorganize the attention calculation results according to the window slicing method to obtain the attention feature maps of the three windows. After merging the three attention feature maps on the channel, integrate the feature information on the three window channels through a convolution to obtain an output feature map of size h*w; reorganize multiple output feature maps at the corresponding positions of the window to obtain the size and deep feature F ac The same output is then operated through window sliding; Step 4.5: Repeat steps 4.2 to 4.4, and then pass through an ESA attention module and a blueprint separable convolution, and then combine it with the shallow feature F in Perform residual connection and then upsample to output high-resolution images; Step 5: Image super-resolution Set the parameters of bicubic interpolation and train models at different magnifications. Select a model with the corresponding magnification as required, input the low-resolution image into the model, and obtain a super-resolution image through feature extraction and enhancement in steps 2 to 4.

2. The image reconstruction method based on multi-window cross-feature fusion attention mechanism as claimed in claim 1, Features: The image pairs are cropped into multiple slices of the same size to enhance the data and expand the dataset.

3. The image reconstruction method based on multi-window cross-feature fusion attention mechanism as claimed in claim 1, Features: The blueprint separable convolution first passes through the deep convolution layer to perform independent convolution operations on the input according to the channel, and independently calculates the convolution result to obtain a tensor of the same size as the input; then passes through the point-by-point convolution layer, using a 1*1 convolution kernel to integrate the information of each channel.

4. The image reconstruction method based on multi-window cross-feature fusion attention mechanism as claimed in claim 1, Features: The efficient channel attention module ECA first obtains the global information of each channel through one-dimensional global average pooling GAP , and then obtain the learnable weights of each channel through a fully connected layer , the output F is calculated based on the weight of each channel eca : (6) (7) in , , They are feature maps F con length, width and channel, Represents the feature map F con The pixel value at position (i, j) in the cth channel, σ represents the activation function, represents the convolution operation, Represents the output feature map F eca The pixel value at position (i,j) in the cth channel.

5. The image reconstruction method based on multi-window cross-feature fusion attention mechanism as claimed in claim 1, Features: The efficient spatial attention module ESA improves the perception of local information by learning the relationship between positions and obtains the output F esa : (8) (9) in Represents the feature map F con The spatial information at position (i, j) in is the weight matrix obtained through learning, Represents the output feature map F esa The pixel value at position (i,j) in the cth channel.

6. A computer-readable storage medium having a computer program stored thereon, which, when executed in a computer, causes the computer to execute the method according to any one of claims 1 to 5.

Images