Image restoration method

Abstract

The invention discloses an image restoration method, solves the problems of large calculation amount and insufficient global modeling capability of an existing transformer-based image restoration method, and belongs to the technical field of image processing. The method comprises the following steps: inputting a to-be-restored image into an image restoration network, and outputting the restored image by the image restoration network; the image restoration network is realized by adopting a U-Net framework and a transformer block; the U-Net architecture comprises four layers of coding and four layers of decoding, the first three layers of coding and the last three layers of decoding respectively adopt a transformer block, compact attention is used for extracting channel domain information, the fourth layer of coding and the first layer of decoding share one transformer block, and cross-window attention is used for extracting spatial domain information. According to the method, self-attention calculation is performed in the channel domain and the space domain, respective advantages are exerted, and the calculation amount is reduced. The cross-window attention method makes up for the insufficiency of the global modeling capability of window segmentation.

Description

Technical Field

[0001] The present invention relates to an image restoration method, belonging to the technical field of image processing. Background Art

[0002] The image restoration task aims to restore clear and detailed results from degraded images caused by various environmental and device conditions. In reality, atmospheric particles such as fog and snow scatter light, resulting in reduced contrast, color deviation, and blurred details in the image, while imaging under low-light conditions has problems such as high noise and missing dark details due to insufficient light. At present, although the transformer-based image restoration method has achieved certain results, it still has limitations in the following three aspects: 1) The computational complexity of the transformer grows quadratically with the increase in resolution. 2) Existing transformer methods use window partitioning to reduce computational overhead, but this weakens the global modeling ability of the transformer. 3) The transformer usually calculates all information, which may include invalid information, and using this invalid information may have a negative effect on image restoration. Summary of the Invention

[0003] Aiming at the problems of large computational complexity and insufficient global modeling ability of the existing transformer-based image restoration method, the present invention provides an image restoration method.

[0004] An image restoration method of the present invention includes: Inputting the image to be restored into an image restoration network, and the image restoration network outputs the restored image; The image restoration network is implemented by using a U-Net architecture and transformer blocks; the U-Net architecture includes four layers of encoding and four layers of decoding. Each of the first three layers of encoding and the last three layers of decoding uses a transformer block, and the fourth layer of encoding and the first layer of decoding share a transformer block. Among them, the transformer blocks of the first three layers of encoding and the last three layers of decoding use compact attention to extract channel domain information, and the fourth layer of encoding and the first layer of decoding use cross-window attention to extract spatial domain information.

[0005] Preferably, the method of using compact attention to extract channel domain information is: Using strided deformable convolution to extract and compress the information of the input image features, that is, in the form where the convolution kernel size and the stride of the strided deformable convolution are equal, generating a position bias of size , and represent the height and width of the image to be restored.

[0006] Preferably, the method of extracting spatial domain information using cross-window attention is: First, the input image features are evenly divided into two parts in the channel dimension ; right Perform uniform window segmentation to obtain multiple small windows, and perform parallel window attention calculation on each small window; right Perform uniform window segmentation to obtain multiple small windows, extract the information of the corresponding position of each small window to form a tensor , computing each tensor in parallel Token attention of The calculated window attention and token attention are fused using ordinary convolution as the extracted spatial domain information.

[0007] Preferably, the image restoration network includes 3 shallow layers and 7 transformer blocks: Treat the restored image Perform scale transformation to obtain three images of different scales of the image to be restored; Represents a tensor collection, and the superscript represents the dimension. is the height of the image to be restored, is the width of the image to be restored, Indicates the number of channels; The image to be restored is input into the first transformer block after 3×3 convolution, and the output of the first transformer block is downsampled; The three scale images are sorted from high to low. The image to be restored at the first scale is input to the first shallow layer, and the first shallow layer outputs The characteristic image of The output downsampled result of the first transformer block is concatenated with the output of the first shallow layer. The concatenated result is input to the second transformer block after 3×3 convolution, and the output of the second transformer block is downsampled. The image to be restored at the second scale is input to the second shallow layer, and the second shallow layer outputs The characteristic image of The output downsampled result of the second transformer block is concatenated with the output of the second shallow layer. The concatenated result is input to the third transformer block after 3×3 convolution, and the output of the third transformer block is downsampled. The image to be restored at the third scale is input to the third shallow layer, and the third shallow layer outputs The characteristic image of The output downsampled result of the third transformer block is concatenated with the output of the third shallow layer. The concatenated result is input to the fourth transformer block after 3×3 convolution, and the output of the fourth transformer block is upsampled. The output upsampling result of the 4th transformer block is concatenated with the output of the 3rd transformer block. The concatenated result is input to the 5th transformer block after 1×1 convolution, and the output of the 5th transformer block is upsampled. The upsampled output of the fifth transformer block is concatenated with the output of the second transformer block. The concatenated result is input to the sixth transformer block after 1×1 convolution, and the output of the sixth transformer block is upsampled. The output upsampling result of the 6th transformer block is concatenated with the output of the 1st transformer block, and the concatenated result is input into the 7th transformer block after 1×1 convolution; The output of the 7th transformer block is added element-by-element to the image to be restored after 3×3 convolution, and the addition result is the restored image.

[0008] Preferably, the transformer block includes a compact deformation attention module, a cross-window attention module, and a gated deep convolutional feed-forward network; The input image features are normalized by layers and then enter the compact deformation attention module or the cross-window attention module. The transformer blocks of the first three encoding layers and the last three decoding layers use a compact deformation attention module to extract channel domain information, and the fourth encoding layer and the first decoding layer use cross-window attention to extract spatial domain information; The output of the compact deformation attention module or the cross-window attention module is added to the input image features element by element. The result of the addition is normalized by the layer and then enters the gated deep convolutional feedforward network. The output of the gated deep convolutional feedforward network and the result of the addition are added element by element, and the result of the addition is the output of the transformer block.

[0009] Preferably, the compact deformation attention module includes a 1×1 convolution, an offset generator, and two depth-wise separable convolutions DWConv; The input image features are , , image features After layer normalization, it enters 1×1 convolution. The result after 1×1 convolution Enter the offset generator, and the offset generator outputs a position offset , The first depthwise separable convolution DWConv is set to have the same convolution kernel size and stride size. The first depthwise separable convolution DWConv uses the position offset to select relevant information of the query Query and the key Key from the result , complete the encoding of the query Query and the key Key, and obtain the encoding matrix and , , . After the encoding matrices and are reshaped respectively, the encoding matrices and are obtained , . After the encoding matrices and are multiplied in matrix and then a softmax operation is performed, is obtained, where represents the convolution kernel size of the first depthwise separable convolution DWConv; The second depthwise separable convolution DWConv is a convolution. Select the relevant information of the value from the result Value , complete the encoding of the value Value , and obtain the encoding matrix , . After the encoding matrix is reshaped, the encoding matrix is obtained ; The obtained is multiplied by the encoding matrix in matrix and then reshaped to obtain , , , which is the channel domain information extracted by the compact deformable attention module.

[0010] Advantages of the present invention The present invention performs self-attention calculations in the channel domain and the spatial domain respectively, gives full play to their respective advantages, and reduces the amount of calculation. In this network, 1) The present invention studies the strided deformable convolution module, reduces the size of the position bias generated by the deformable convolution, and improves the efficiency of the deformable convolution. 2) The present invention proposes a cross-window attention method to make up for the lack of global modeling ability of window segmentation. Description of the drawings

[0011] Figure 1Schematic diagram of the principle of the image restoration network of the present invention; Figure 2 This is a schematic diagram of the transformer block; Figure 3 Schematic diagram of the compact deformation attention module; Figure 4 Schematic diagram of the principle of the cross-window attention module; Figure 5 This is a schematic diagram of the shallow layer; Figure 6 It is the principle schematic diagram of the offset generator; Figure 7 This is the effect diagram of the cross-window attention module; Figure 8 These are the effects of the four sets of images before and after restoration. DETAILED DESCRIPTION

[0012] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0013] It should be noted that, in the absence of conflict, the embodiments of the present invention and the features in the embodiments may be combined with each other.

[0014] The present invention will be further described below in conjunction with the accompanying drawings and specific embodiments, but they are not intended to limit the present invention.

[0015] The image restoration method of this embodiment includes: Step 1: Establish an image restoration network; The image processed in this embodiment can be used as 3D data express, and Indicates the height and width of the image, Indicates the number of channels of the image. In a preferred embodiment, the image restoration network of this implementation is implemented using a U-Net architecture and a transformer block; the U-Net architecture can capture multi-scale features and contain more underlying information. Different self-attention methods (spatial domain / channel domain) are used in different transformer blocks. When the feature resolution is large, this implementation calculates attention in the channel domain, and when the resolution is small, it calculates attention in the spatial domain, and performs image restoration in the form of a combination of the spatial domain and the channel domain.

[0016] Specifically, the U-Net architecture of this embodiment includes four layers of encoding and four layers of decoding. The first three layers of encoding and the last three layers of decoding each use a transformer block, and all use compact attention to extract channel domain information. The fourth layer of encoding and the first layer of decoding share a transformer block, and use cross-window attention to extract spatial domain information.

[0017] In this implementation, the method of extracting channel domain information using compact attention is: Use strided deformable convolution to extract and compress the information of input image features. The convolution kernel size and step size of strided deformable convolution are equal, and the generated size is The position offset of and Indicates the height and width of the image to be restored; In the channel domain, starting from the query Query and the key Key, this implementation uses strided deformable convolution to share position deviations between them, extract the most relevant information, and achieve efficient extraction of channel domain information.

[0018] In this implementation, the method of extracting spatial domain information using cross-window attention is: First, the input image features are evenly divided into two parts in the channel dimension ;right Perform uniform window segmentation to obtain multiple small windows, and perform parallel window attention calculation on each small window. After the window segmentation, the attention is only calculated within the small window, lacking cross-window connection. Perform uniform window segmentation to obtain multiple small windows, extract the information of the corresponding position of each small window to form a new tensor , the token attention of each tensor is calculated in parallel; then ordinary convolution is used to fuse the calculated window attention and token attention as the extracted spatial domain information.

[0019] The cross-window attention extraction in this embodiment not only captures the relationship between pixels within a single window, but also models cross-window dependencies through window position information, captures the relationship between pixels between different windows, and makes up for the lack of global modeling capabilities of window segmentation.

[0020] Specifically, Figure 1 As shown, the image restoration network of this embodiment includes 3 shallow layers and 7 transformer blocks: the first 4 transformer blocks are encoders, and the last 3 transformer blocks are decoders; Treat the restored image Perform scale transformation to obtain three images of different scales of the image to be restored; Represents a tensor collection, and the superscript represents the dimension. is the height of the image to be restored, is the width of the image to be restored, Indicates the number of channels; The image to be restored is input into the first transformer block after 3×3 convolution, and the output of the first transformer block is downsampled; The three scale images are sorted from high to low. The image to be restored at the first scale is input to the first shallow layer, and the first shallow layer outputs The characteristic image of The output downsampled result of the first transformer block is concatenated with the output of the first shallow layer. The concatenated result is input to the second transformer block after 3×3 convolution, and the output of the second transformer block is downsampled. The image to be restored at the second scale is input to the second shallow layer, and the second shallow layer outputs The characteristic image of The output downsampled result of the second transformer block is concatenated with the output of the second shallow layer. The concatenated result is input to the third transformer block after 3×3 convolution, and the output of the third transformer block is downsampled. The image to be restored at the third scale is input to the third shallow layer, and the third shallow layer outputs The characteristic image of The output downsampled result of the third transformer block is concatenated with the output of the third shallow layer. The concatenated result is input to the fourth transformer block after 3×3 convolution, and the output of the fourth transformer block is upsampled. The output upsampling result of the 4th transformer block is concatenated with the output of the 3rd transformer block. The concatenated result is input to the 5th transformer block after 1×1 convolution, and the output of the 5th transformer block is upsampled. The upsampled output of the fifth transformer block is concatenated with the output of the second transformer block. The concatenated result is input to the sixth transformer block after 1×1 convolution, and the output of the sixth transformer block is upsampled. The output upsampling result of the 6th transformer block is concatenated with the output of the 1st transformer block, and the concatenated result is input into the 7th transformer block after 1×1 convolution; The output of the 7th transformer block is added element-by-element to the image to be restored after 3×3 convolution, and the addition result is the restored image.

[0021] This embodiment uses a four-layer U-Net architecture for deep feature extraction. In the upsampling and downsampling process, inverse pixel reorganization and pixel reorganization operations are applied to aggregate information in the encoder and decoder, and then a 3×3 convolution is used to reduce the number of channels and output the image. This embodiment adopts a multi-input multi-output strategy and constrains the loss function during training to capture the multi-scale information of the image. The shallow layer is used to extract the multi-scale information of the input image. In this embodiment, the shallow layer can be implemented by sequentially connecting 3×3 convolution, 1×1 convolution, 3×3 convolution, and 1×1 convolution.

[0022] Specifically, Figure 2 As shown, the transformer block includes a compact deformation attention module, a cross-window attention module, and a gated deep convolutional feed-forward network; The input image features are normalized by layers and then enter the compact deformation attention module or the cross-window attention module. The transformer blocks of the first three encoding layers and the last three decoding layers use a compact deformation attention module to extract channel domain information, and the fourth encoding layer and the first decoding layer use cross-window attention to extract spatial domain information; The output of the compact deformation attention module or the cross-window attention module is added to the input image features element by element. The result of the addition is normalized by the layer and then enters the gated deep convolutional feedforward network. The output of the gated deep convolutional feedforward network and the result of the addition are added element by element, and the result of the addition is the output of the transformer block.

[0023] In the channel domain, the compact deformation attention module designed in this embodiment has good information extraction capability. The compact deformation attention module can selectively extract effective information of features and compress this information, which can provide a reliable information source for subsequent image restoration. Its structure is as follows Figure 3 As shown. For the feature map Use ordinary convolution to encode the value Value (V). In order to make the information for attention calculation as relevant as possible, use strided deformed convolution to encode Query (Q) and Key (K) and downsample to reduce the amount of calculation. After downsampling, , through matrix multiplication You can get a Transformer attention map of common size, and then pass: Get the output. Specifically, the compact deformable attention module in this embodiment includes a 1×1 convolution, an offset generator, and two depthwise separable convolutions DWConv; The input image feature is , , the image feature After layer normalization, it enters the 1×1 convolution, and the result after the 1×1 convolution Enters the offset generator, and the offset generator outputs the position offset , , the first depthwise separable convolution DWConv is set to have the same convolution kernel size and stride size. The first depthwise separable convolution DWConv selects the relevant information of the query Query and the key Key in the result In the result To complete the encoding of the query Query and the key Key, and obtain the encoding matrix And , , , the encoding matrix And After reshaping respectively, the encoding matrices And are obtained, , , the encoding matrix And After matrix multiplication, perform the softmax operation to obtain , where Represents the convolution kernel size of the first depthwise separable convolution DWConv; The second depthwise separable convolution DWConv is Convolution, select the relevant information of the value In the result Value To complete the encoding of the value Value , and obtain the encoding matrix , , reshape the encoding matrix To obtain the encoding matrix , ; The obtained Performs matrix multiplication with the encoding matrix And then reshapes to obtain , , Is the channel domain information extracted by the compact deformable attention module. The offset generator in this embodiment includes a k × k Convolution and a 1×1 convolution connected in sequence.

[0024] In the spatial domain, the cross-window attention module designed in this embodiment has good spatial information aggregation ability. The cross-window attention module considers both global and local information at the same time, enhancing the global modeling ability of the network. Figure 7 As shown in the figure, the cross-window attention module includes window attention and inter-window attention, which are used to extract local and global information respectively. In window attention, the feature map is partitioned into windows, and then token attention is calculated for each small window. In inter-window attention, the window is also partitioned. In this stage, the information at the same position of each window is extracted to form a new window, which is then subjected to token attention calculation. Then the results of the two stages are sent to a The convolution is combined.

[0025] Step 2: Train the established image restoration network: During the training process of the entire model, this implementation uses two loss functions to constrain the optimal target of the model. They are pixel-by-pixel loss and frequency domain loss. Pixel-by-pixel loss is used to measure the element difference between the clean image and the high-quality image restored by the model. Frequency domain loss is used to calculate the difference in frequency domain information corresponding to the Fourier transform of the clean image and the restored high-quality image. The above two loss functions are formulated as follows: In the formula, represents a clean image, Represents a high-quality image with noise. represents the calculation process of the image restoration network, represents Fourier transform.

[0026] Step 3: Input the image to be restored into the trained image restoration network, and the image restoration network outputs the restored image; The quantitative results on the dehazing and snow removal datasets are shown in Table 1: Table 1 Test results of the method in this embodiment on the defogging and snow removal datasets The quantitative results on the low-light image enhancement dataset are shown in Table 2: Table 2 Test results of the method in this embodiment on the low-light image enhancement dataset The image restoration result of the image restoration network in this embodiment is as follows Figure 8 (a), (b), (c), and (d) are four groups of images before and after restoration.

[0027] Although the present invention is described herein with reference to specific embodiments, it should be understood that these embodiments are merely examples of the principles and applications of the present invention. It should therefore be understood that many modifications may be made to the exemplary embodiments and that other arrangements may be devised without departing from the spirit and scope of the present invention as defined by the appended claims. It should be understood that the features of the various dependent claims and herein may be combined in a manner different from that described in the original claims. It should also be understood that the features described in conjunction with the individual embodiments may be used in other embodiments.

Claims

1. An image restoration method, characterized in that: include: The image to be restored is input into the image restoration network, and the image restoration network outputs the restored image; The image restoration network is implemented using a U-Net architecture and a transformer block; the U-Net architecture includes four layers of encoding and four layers of decoding, the first three layers of encoding and the last three layers of decoding each use a transformer block, the fourth layer of encoding and the first layer of decoding share a transformer block, wherein the transformer blocks of the first three layers of encoding and the last three layers of decoding use compact attention to extract channel domain information, and the fourth layer of encoding and the first layer of decoding use cross-window attention to extract spatial domain information.

2. The image restoration method according to claim 1, characterized in that: The method of using compact attention to extract channel domain information is: Use strided deformable convolution to extract and compress the information of input image features, that is, the convolution kernel size and step size of strided deformable convolution are equal, generating a size of The position offset of and Indicates the height and width of the image to be restored.

3. The image restoration method according to claim 1, characterized in that: The method of extracting spatial domain information using cross-window attention is: First, the input image features are evenly divided into two parts in the channel dimension ; right Perform uniform window segmentation to obtain multiple small windows, and perform parallel window attention calculation on each small window; right Perform uniform window segmentation to obtain multiple small windows, extract the information of the corresponding position of each small window to form a tensor , computing each tensor in parallel Token attention of The calculated window attention and token attention are fused using ordinary convolution as the extracted spatial domain information.

4. The image restoration method according to claim 1, characterized in that: The image restoration network consists of 3 shallow layers and 7 transformer blocks: Treat the restored image Perform scale transformation to obtain three images of different scales of the image to be restored; Represents a tensor collection, and the superscript represents the dimension. is the height of the image to be restored, is the width of the image to be restored, Indicates the number of channels; The image to be restored is input into the first transformer block after 3×3 convolution, and the output of the first transformer block is downsampled; The three scale images are sorted from high to low. The image to be restored at the first scale is input to the first shallow layer, and the first shallow layer outputs The characteristic image of The output downsampled result of the first transformer block is concatenated with the output of the first shallow layer. The concatenated result is input to the second transformer block after 3×3 convolution, and the output of the second transformer block is downsampled. The image to be restored at the second scale is input to the second shallow layer, and the second shallow layer outputs The characteristic image of The output downsampled result of the second transformer block is concatenated with the output of the second shallow layer. The concatenated result is input to the third transformer block after 3×3 convolution, and the output of the third transformer block is downsampled. The image to be restored at the third scale is input to the third shallow layer, and the third shallow layer outputs The characteristic image of The output downsampled result of the third transformer block is concatenated with the output of the third shallow layer. The concatenated result is input to the fourth transformer block after 3×3 convolution, and the output of the fourth transformer block is upsampled. The output upsampling result of the 4th transformer block is concatenated with the output of the 3rd transformer block. The concatenated result is input to the 5th transformer block after 1×1 convolution, and the output of the 5th transformer block is upsampled. The upsampled output of the fifth transformer block is concatenated with the output of the second transformer block. The concatenated result is input to the sixth transformer block after 1×1 convolution, and the output of the sixth transformer block is upsampled. The output upsampling result of the 6th transformer block is concatenated with the output of the 1st transformer block, and the concatenated result is input into the 7th transformer block after 1×1 convolution; The output of the 7th transformer block is added element-by-element to the image to be restored after 3×3 convolution, and the addition result is the restored image.

5. The image restoration method according to claim 4, characterized in that: The transformer block includes a compact deformable attention module, a cross-window attention module, and a gated deep convolutional feed-forward network; The input image features are normalized by layers and then enter the compact deformation attention module or the cross-window attention module. The transformer blocks of the first three encoding layers and the last three decoding layers use a compact deformation attention module to extract channel domain information, and the fourth encoding layer and the first decoding layer use cross-window attention to extract spatial domain information; The output of the compact deformation attention module or the cross-window attention module is added element-by-element to the input image features, and the result of the addition is normalized by the layer and then enters the gated deep convolutional feedforward network. The output of the gated deep convolutional feedforward network and the result of the addition are added element-by-element, and the result of the addition is the output of the transformer block.

6. The image restoration method according to claim 5, characterized in that: The compact deformation attention module includes a 1×1 convolution, an offset generator, and two depth-wise separable convolutions DWConv; The input image feature is , , image features After layer normalization, it enters 1×1 convolution. The result after 1×1 convolution Enter the offset generator, the offset generator outputs the position offset , The first depth-wise separable convolution DWConv is set to have the same convolution kernel size and step size. The first depth-wise separable convolution DWConv is offset by position. In the results Select the relevant information of query Query and key Key, complete the encoding of query Query and key Key, and obtain the encoding matrix and , , , encoding matrix and After reshaping respectively, the encoding matrix is ​​obtained and , , , encoding matrix and After matrix multiplication, softmax operation is performed to obtain ,in Indicates the convolution kernel size of the first depth-wise separable convolution DWConv; The second depth-wise separable convolution DWConv is Convolution, in the result Select the value Value Complete the value of Value The encoding matrix is ​​obtained , , for the encoding matrix After reshaping, the encoding matrix is ​​obtained , ; What you get Matrix with the encoding matrix After multiplication, reshape to get , , Channel domain information extracted for compact deformation attention module.

7. The image restoration method according to claim 6, characterized in that: The offset generator consists of connecting k × k convolution and 1×1 convolution.

8. The image restoration method according to claim 4, characterized in that: The shallow layer includes 3×3 convolution, 1×1 convolution, 3×3 convolution, and 1×1 convolution, which are connected sequentially.

Images