Single image deraining method and system based on image prior and gated attention learning

Abstract

The present invention discloses a single-image rain removal method and system based on image prior and gated attention learning. The method comprises: constructing a network model based on image prior and gated attention learning, the network model comprising a sequentially connected feature extractor, a codec, and a gater with an attention mechanism; obtaining training samples; forming a training dataset of rainy-rainy image pairs; preprocessing the training dataset of rainy-rainy image pairs, inputting the preprocessed data into the network model based on image prior and gated attention learning for training, and obtaining a trained network model based on image prior and gated attention learning; inputting the image to be rained into the trained network model based on image prior and gated attention learning, and finally obtaining the image data after rain removal. The present invention effectively extracts detailed information of the image and identifies more rain streaks, thereby enhancing the rain removal effect of the rain image.

Description

Technical Field

[0001] The present invention relates to the field of image processing, and in particular to a single-image rain removal method and system based on image prior and gated attention learning. Background Art

[0002] In daily life, shooting in the rain will affect the image. Rain streaks will appear on the image, which will reduce the visual quality of the image and even affect most visual tasks based on image visual effects, such as video surveillance, autonomous driving, and target recognition. Therefore, it is particularly important to remove rain from the image.

[0003] Currently, single-image deraining methods can be roughly divided into two categories: model-driven and data-driven. Model-driven methods focus on manually designing and utilizing prior knowledge of images, obtaining rain-free images by adding constraints and modeling algorithms and solving optimization models; while data-driven methods learn the mapping relationship between rainy images and rain-free images from data by building neural networks. Due to their powerful expressive ability, neural networks have achieved great results in image deraining tasks.

[0004] However, the current rain removal methods based on data-driven methods still have the problems of extracting less feature information and ignoring the prior information of the image itself, and the rain removal effect needs to be improved. Summary of the Invention

[0005] The main purpose of the present invention is to overcome the shortcomings and deficiencies of the existing technology and provide a single-image rain removal method based on image prior and gated attention learning, which effectively extracts and fully utilizes the feature information and prior information in the image, thereby efficiently identifying more rain streak information and enhancing the rain removal effect of the rain image.

[0006] In order to achieve the above object, the present invention adopts the following technical solutions:

[0007] A single image deraining method based on image prior and gated attention learning includes the following steps:

[0008] S1. Construct a network model based on image prior and gated attention learning, wherein the network model includes a sequentially connected feature extractor, an encoder-decoder, and a gate controller with an attention mechanism;

[0009] S2. Obtain training samples; obtain corresponding rain image samples by adding synthetic rain patterns to a rainless background image. The rainless background image and rain image samples constitute a rain-no rain image pair training dataset;

[0010] S3. Preprocessing the rain-no rain image training data set, inputting the preprocessed data into the network model based on image prior and gated attention learning for training, to obtain a trained network model based on image prior and gated attention learning; the training process specifically includes the following steps:

[0011] S31, the feature extractor extracts rain image features and prior features from the input rain image sample;

[0012] S32 normalizes the rain image features and the prior features of the rain image sample, then concatenates the normalized rain image features and the prior features, and inputs the concatenated rain image features and the prior features into a codec for feature separation to obtain a preliminary rain streak feature map; the rain image features include texture and edge features of the rain image;

[0013] S33: The gate controller with the attention mechanism extracts feature dependency information in the spatial direction and feature dependency information in the channel direction of the preliminary rain streak feature map, combines the feature dependency information in the spatial direction and the feature dependency information in the channel direction to identify rain streaks and residual background layer information on the preliminary rain streak feature map, and filters the residual background layer information while identifying the rain streak information to obtain final rain streak information.

[0014] S34 subtracts the final rain streak information from the input rain image sample to obtain a rain-removed image;

[0015] S4. Input the image to be derained into the trained network model based on image prior and gated attention learning, and finally obtain the image data after deraining.

[0016] Preferably, the feature extractor is composed of two residual blocks in parallel, one of the two residual blocks is used to extract rain image features of the input rain image, and the other residual block is used to extract prior features of the input rain image.

[0017] Preferably, the codec is provided with 5 layers, and each codec layer includes an encoder and a decoder; each layer in the encoder and decoder is composed of 2 residual blocks with Squeeze and Excitation operations.

[0018] Preferably, the codec uses a skip connection to short-circuit corresponding layers of the encoder and decoder.

[0019] Preferably, the gate controller with the attention mechanism is composed of a residual gating unit and a non-local channel block, and the residual gating unit and the non-local channel block are a serial structure. The residual gating unit is used to obtain feature dependency information of the preliminary rain streak feature map in the spatial direction, and the non-local channel block is used to obtain feature dependency information of the preliminary rain streak feature map in the channel direction.

[0020] Preferably, the feature extractor in step S31 extracts rain image features and prior features from the input rain image sample, specifically as follows:

[0021] First, a residual block with Squeeze and Excitation operations is used to extract the rain image feature L of the input rain image. The rain image feature L is defined as follows:

[0022] L=SE(ResBlock(R))

[0023] Where R represents the input rain image sample, SE(·) represents the Squeeze and Excitation operation, and ResBlock(·) represents a residual block that constitutes ResNet;

[0024] Then set an RCP prior, which is defined as follows:

[0025] I P (x)=I M (x)-I m (x)

[0026] Where x represents the position of each pixel, I represents the RGB image, and I M Indicates the maximum pixel value of the RGB image, I m Indicates the minimum pixel value of the RGB image, I P represents a priori;

[0027] Obtaining grayscale image of residual channel map based on RCP prior

[0028] Finally, the prior feature P of the rain image is extracted, and the prior feature P is as follows:

[0029]

[0030] Preferably, the codec in step S1 adopts a PReLU activation function to obtain nonlinearity, and uses bilinear downsampling instead of the maximum pooling layer operation.

[0031] Preferably, the gate controller with the attention mechanism in step S33 extracts feature dependency information in the spatial direction and feature dependency information in the channel direction of the preliminary rain streak feature map, specifically as follows:

[0032] First, the residual gating unit is used to obtain feature dependency information in the spatial direction. The residual gating unit is defined as follows:

[0033]

[0034] where O∈R H×W×C and X∈R H×W×C Represents the output gate feature map and input feature map, Conv 3×3 Represents a standard 3×3 2D convolution operation, which divides the input feature map into N parts along the channel dimension and obtains the first sub-feature map X1∈R H×W×C and the remaining i-th molecular feature map X i ∈R H×W×C , ⊙ represents Hadamard product, σ represents sigmoid activation function to obtain nonlinearity;

[0035] Then, the non-local channel block is used to obtain the feature dependency information in the channel direction. The non-local channel block is defined as follows:

[0036]

[0037] where Z∈R H×W×C and O∈R H×W×C Represent the output attention feature map and input gated feature map, Conv 1×1 Represents a standard 1×1 2D convolution operation, which convolves the input gated feature map and resizes it to obtain g∈R HW×C ,θ∈R C×HW ,φ∈R HW×C , Represents matrix multiplication, softmax represents the softmax activation function operation, and the attention matrix is ​​C×C;

[0038] Finally, the gate controller filters out the residual background layer information in the preliminary rain streak feature map and detects more fine rain streak information through the feature dependency information in the spatial direction and the feature dependency information in the channel direction obtained above.

[0039] Preferably, the loss function of the network model based on image prior and gated attention learning adopts the sum of L1 loss function and edge function;

[0040] The loss function is defined as follows:

[0041]

[0042] Where B represents the de-rained image after processing, G represents the real rain-free image, △ represents the Laplace operation, ||BG||1 represents the L1 loss function, represents the marginal loss function, and λ and ε represent hyperparameters.

[0043] A second aspect of the present invention provides a single-image rain removal system based on image prior and gated attention learning, comprising a processor and a memory, wherein the memory stores non-temporary computer instructions. When the non-temporary computer instructions are executed by the processor, the single-image rain removal method based on image prior and gated attention learning is executed.

[0044] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0045] (1) The present invention proposes a single-image rain removal method based on image prior and gated attention learning. By combining image prior and the gating mechanism of fusion attention learning, the method makes fuller use of the information provided by a single rain image, better preserves the image details and identifies more rain streaks, effectively improving the subjective visual effect.

[0046] (2) The codec scheme adopted by the present invention obtains the global and local features of the image based on the prior features and rain image features, thereby enhancing the recognition ability of rain streaks.

[0047] (3) The gate controller with attention mechanism of the present invention facilitates the separation of rain streaks and background layer by constructing feature dependencies in space and channels, thereby identifying smaller rain streaks in the image, thereby improving the deraining effect; since rain streaks are relatively small, the establishment of spatial feature dependencies is to find more rain streaks on the preliminary rain streak feature map; as the depth of the neural network increases, the background layer will overlap with the feature distribution of rain streaks in the channel. By establishing feature dependencies in the channel, the neural network pays attention to the feature distribution of rain streaks, thereby filtering out the background layer. BRIEF DESCRIPTION OF THE DRAWINGS

[0048] Figure 1 This is a flowchart of a single-image deraining method based on image prior and gated attention learning according to an embodiment of the present invention.

[0049] Figure 2 This is a workflow diagram of a network model based on image prior and gated attention learning according to an embodiment of the present invention. DETAILED DESCRIPTION

[0050] In order to make the purpose, features, and advantages of the present invention more obvious and easy to understand, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the embodiments described below 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 making creative work are within the scope of protection of the present invention.

[0051] The technical solution of the present invention will be further described below with reference to the accompanying drawings and through specific implementation methods.

[0052] Example 1

[0053] like Figure 1 As shown, a single image deraining method based on image prior and gated attention learning in this embodiment includes the following steps:

[0054] S1. Construct a network model based on image prior and gated attention learning, wherein the network model includes a sequentially connected feature extractor, an encoder-decoder, and a gate controller with an attention mechanism;

[0055] Specifically, the feature extractor is composed of two residual blocks in parallel, one of the two residual blocks is used to extract rain image features of the input rain image, and the other residual block is used to extract prior features of the input rain image.

[0056] The codec has five layers, each of which includes an encoder and a decoder; each layer in the encoder and decoder is composed of two residual blocks with squeeze and excitation operations, and the codec uses skip connections to short-circuit the corresponding layers of the encoder and decoder.

[0057] The gate controller with attention mechanism is composed of a residual gating unit and a non-local channel block. The residual gating unit and the non-local channel block are serial structures. The residual gating unit is used to obtain feature dependency information of a preliminary rain streak feature map in the spatial direction, and the non-local channel block is used to obtain feature dependency information of a preliminary rain streak feature map in the channel direction.

[0058] The loss function of the network based on image prior and gated attention learning includes a fidelity term and a detail term, specifically the sum of an L1 loss function and an edge function;

[0059] Specifically, the loss function is defined as follows:

[0060]

[0061] Where B represents the derained image obtained after network processing, G represents the real rain-free image, △ represents the Laplace operation, ||BG||1 represents the L1 loss function, represents the marginal loss function, λ and ε represent hyperparameters, where the experiment shows that λ is set to 0.05 and ε is set to 10 -3 Can achieve better rain removal effect.

[0062] Traditional rain removal methods use the mean square loss function as the loss function. Since the square constraint on high-frequency textures usually produces a blurred and overly smooth rain removal effect, this embodiment uses the L1 loss function and the edge function to form a loss function to ensure the restoration of the edge information of the image and a real and natural effect.

[0063] S2. Obtain training samples. By superimposing synthetic rain streaks of different densities and directions on a rain-free background image, corresponding rain image samples are obtained. The rain-free background image and the rain image samples constitute a training dataset of rain-free image pairs.

[0064] S3. Preprocessing the rain-no rain image training data set, inputting the preprocessed data into the network model based on image prior and gated attention learning for training, to obtain a trained network model based on image prior and gated attention learning; the training process specifically includes the following steps:

[0065] The feature extractor in step S31 extracts rain image features and prior features from the input rain image sample; specifically, as follows:

[0066] (1) First, a residual block with Squeeze and Excitation operations is used to extract the rain image feature L of the input rain image. The rain image feature L is defined as follows:

[0067] L=SE(ResBlock(R))

[0068] Among them, R represents the input rain map, SE(·) represents the Squeeze and Excitation operation, and ResBlock(·) represents a residual block that constitutes ResNet.

[0069] (2) Setting an RCP prior to obtain the image’s prior features represented as P, to supplement the image’s feature information and protect the background details of the restored image from being lost during the training process;

[0070] The RCP is defined a priori as follows:

[0071] I P (x)=I M (x)-I m (x)

[0072] Where x represents the position of each pixel, I represents the RGB image, and I M Indicates the maximum pixel value of the RGB image, I m Indicates the minimum pixel value of the RGB image, I P Represents a priori.

[0073] (3) Obtaining the grayscale image of the residual channel map based on RCP prior

[0074] Since the RCP prior is single-channel, the residual channel map is a three-channel grayscale image connected along the channel direction. The residual channel map is defined as follows:

[0075]

[0076] (4) Based on residual channel map Extract the prior features P of the rain image, which are as follows:

[0077]

[0078] in, Represents the residual channel image, SE(·) represents the Squeeze and Excitation operation, and ResBlock(·) represents a residual block that constitutes ResNet.

[0079] Compared with the traditional method of extracting only rain image features from the rain image and expressing them as L, this embodiment adopts RCP prior to obtain the image prior features and express them as P to supplement the feature information of the image and protect the background details of the restored image from being lost during the training process.

[0080] S32 normalizes the rain image features and the prior features of the rain image sample, then concatenates the normalized rain image features and the prior features, and inputs the concatenated rain image features and the prior features into a codec for feature separation to obtain a preliminary rain streak feature map; the rain image features include texture and edge features of the rain image;

[0081] Specifically, the codec uses a U-Net-based codec to extract preliminary rain image features. The codec has five layers, each consisting of two residual blocks with squeeze and excitation operations. The PReLU activation function is used to obtain nonlinearity, bilinear downsampling is used instead of the maximum pooling layer operation, and skip connections are used to short-circuit the corresponding layers of the encoder and decoder to reduce the loss of image information.

[0082] S33: The gate controller with the attention mechanism extracts feature dependency information in the spatial direction and feature dependency information in the channel direction of the preliminary rain streak feature map, combines the feature dependency information in the spatial direction and the feature dependency information in the channel direction to identify rain streaks and residual background layer information on the preliminary rain streak feature map, and filters the residual background layer information while identifying the rain streak information to obtain final rain streak information; specifically, as follows:

[0083] (1) First, the residual gating unit is used to obtain the feature dependency information in the spatial direction. The residual gating unit of the gater is defined as follows:

[0084]

[0085] where O∈R H×W×C and X∈R H×W×C Represents the output gate feature map and input feature map, Conv 3×3 Represents a standard 3×3 2D convolution operation, which divides the input feature map into N parts along the channel dimension and obtains the first sub-feature map X1∈R H×W×C and the remaining i-th molecular feature map X i ∈R H×W×C ⊙ represents the Hadamard product, and σ represents the sigmoid activation function to obtain nonlinearity. By using the first sub-feature map as a mask, the other sub-feature maps are given a higher degree of nonlinearity in local features, thereby enhancing the network's expressive power. At the same time, a dependency relationship between the first sub-feature map and other feature maps is also established, allowing the network to use rain streaks as regions of interest and enhance learning.

[0086] (2) The non-local channel block is used to obtain the feature dependency information in the channel direction. The local channel block of the gate controller is defined as follows:

[0087]

[0088] where Z∈R H×W×C and O∈R H×W×C Represent the output attention feature map and input gated feature map, Conv 1×1 Represents a standard 1×1 2D convolution operation, which convolves the input gated feature map and resizes it to obtain g∈R HW×C ,θ∈R C×HW ,φ∈R HW×C , Represents matrix multiplication, softmax represents the softmax activation function operation, and the attention matrix is ​​C×C; by modifying the attention matrix from HW×HW to C×C, the non-local channel block improves the rain removal performance without spending a large number of parameters for matrix calculation.

[0089] Finally, the gate controller filters out the residual background layer information in the preliminary rain streak feature map and detects more fine rain streak information through the feature dependency information in the spatial direction and the feature dependency information in the channel direction obtained above.

[0090] S34: subtracting the final rain streak information from the input rain image sample to obtain a rain-removed image;

[0091] Specifically, due to the complexity of the physical modeling of rain images, rain images are usually considered to be the superposition of the background layer and the rain layer. That is, the physical model of rain is simply defined as:

[0092] I=B+R

[0093] Where I represents the physical model of rain, B is the background layer, and R is the rain layer. Therefore, a clear image can be obtained by subtracting the corresponding rain streaks from the rain image.

[0094] It should be noted that the optimization method used in the network training based on image prior and gated attention learning is Adam, the parameter momentum is 0.9, and the initial learning rate is set to 2×10 -4 The batch size is set to 24, the learning rate regulator is a cosine annealing regulator with restarts, and the minimum learning rate is set to 1×10 -5 .

[0095] The network model based on image prior and gated attention learning was trained for 300 rounds, and the results were tested on the test set every 5 rounds. The test results were measured using peak signal-to-noise ratio (PSNR) and structural similarity (SSIM), and the best result was selected to save the model parameters.

[0096] S4. Input the image to be derained into the trained network model based on image prior and gated attention learning, and finally obtain the image data after deraining.

[0097] Specifically, such as Figure 2 As shown in Figure 2, the workflow of the network model based on image prior and gated attention learning is as follows:

[0098] Inputting the image to be derained into the network based on image prior and gated attention learning;

[0099] The feature extractor in the network extracts rain image features and prior features corresponding to the input rain image; the codec uses the two features as input to separate the rain streak layer and distinguish different image contents, thereby obtaining a preliminary rain streak feature map; the gate controller with an attention mechanism obtains the feature dependency information of the preliminary rain streak feature map in the spatial direction and the feature dependency information in the channel direction to identify the final rain streak information; the final rain streak information is removed from the input rain image to obtain a clear image after rain removal.

[0100] Example 2

[0101] This embodiment provides a single-image rain removal system based on image prior and gated attention learning, including a processor and a memory, wherein the memory stores non-transitory computer instructions. When the non-transitory computer instructions are executed by the processor, the single-image rain removal method based on image prior and gated attention learning described in Example 1 is executed.

[0102] The above embodiments are preferred implementation modes of the present invention, but the implementation modes of the present invention are not limited to the above embodiments. Any other changes, modifications, substitutions, combinations, and simplifications that do not deviate from the spirit and principles of the present invention should be considered as equivalent replacement methods and are included in the scope of protection of the present invention.

Claims

1. A single image deraining method based on image prior and gated attention learning, characterized by: The following steps are involved: S1. Construct a network model based on image prior and gated attention learning, wherein the network model includes a sequentially connected feature extractor, an encoder-decoder, and a gate controller with an attention mechanism; S2. Obtain training samples; obtain corresponding rain image samples by adding synthetic rain patterns to a rainless background image. The rainless background image and rain image samples constitute a rain-no rain image pair training dataset; S3. Preprocessing the rain-no rain image training data set, inputting the preprocessed data into the network model based on image prior and gated attention learning for training, to obtain a trained network model based on image prior and gated attention learning; the training process specifically includes the following steps: The feature extractor in step S31 extracts rain image features and prior features from the input rain image sample: First, a residual block with Squeeze and Excitation operations is used to extract the rain image feature L of the input rain image. The rain image feature L is defined as follows: L=SE(ResBlock(R)) Where R represents the input rain image sample, SE(·) represents the Squeeze and Excitation operation, and ResBlock(·) represents a residual block that constitutes ResNet; Then set an RCP prior, which is defined as follows: I P (x)=I M (x)-I m (x) Where x represents the position of each pixel, I represents the RGB image, and I M Indicates the maximum pixel value of the RGB image, I m Indicates the minimum pixel value of the RGB image, I P represents a priori; Obtaining grayscale image of residual channel map based on RCP prior Finally, the prior feature P of the rain image is extracted, and the prior feature P is as follows: S32 normalizes the rain image features and the prior features of the rain image sample, then concatenates the normalized rain image features and the prior features, and inputs the concatenated rain image features and the prior features into a codec for feature separation to obtain a preliminary rain streak feature map; the rain image features include texture and edge features of the rain image; S33: The gate controller with the attention mechanism extracts feature dependency information in the spatial direction and feature dependency information in the channel direction of the preliminary rain streak feature map, combines the feature dependency information in the spatial direction and the feature dependency information in the channel direction to identify rain streaks and residual background layer information on the preliminary rain streak feature map, and filters the residual background layer information while identifying the rain streak information to obtain final rain streak information. S34 subtracts the final rain streak information from the input rain image sample to obtain a rain-removed image; S4. Input the image to be derained into the trained network model based on image prior and gated attention learning, and finally obtain the image data after deraining.

2. The single image deraining method based on image prior and gated attention learning according to claim 1, characterized in that: The feature extractor is composed of two residual blocks in parallel, one of the two residual blocks is used to extract rain image features of an input rain image, and the other residual block is used to extract a priori features of the input rain image.

3. The single image deraining method based on image prior and gated attention learning according to claim 1, characterized in that: The codec has five layers, and each layer of the codec includes an encoder and a decoder; each layer of the encoder and decoder is composed of two residual blocks with Squeeze and Excitation operations.

4. The single image deraining method based on image prior and gated attention learning according to claim 3, characterized in that: The codec employs skip connections to short-circuit the corresponding layers of the encoder and decoder.

5. The single image deraining method based on image prior and gated attention learning according to claim 1, characterized in that The gate controller with attention mechanism is composed of a residual gating unit and a non-local channel block. The residual gating unit and the non-local channel block are serial structures. The residual gating unit is used to obtain feature dependency information of a preliminary rain streak feature map in the spatial direction, and the non-local channel block is used to obtain feature dependency information of a preliminary rain streak feature map in the channel direction.

6. The single image deraining method based on image prior and gated attention learning according to claim 1, characterized in that: The codec in step S1 uses the PReLU activation function to obtain nonlinearity and uses bilinear downsampling instead of the maximum pooling layer operation.

7. The single image deraining method based on image prior and gated attention learning according to claim 1, characterized in that: The gate controller with the attention mechanism in step S33 extracts the feature dependency information in the spatial direction and the feature dependency information in the channel direction of the preliminary rain streak feature map, as follows: First, the residual gating unit is used to obtain feature dependency information in the spatial direction. The residual gating unit is defined as follows: where O∈R H×W×C and X∈R H×W×C Represents the output gate feature map and input feature map, Conv 3×3 Represents a standard 3×3 2D convolution operation, which divides the input feature map into N parts along the channel dimension and obtains the first sub-feature map X1∈R H ×W×C and the remaining i-th molecular feature map X i ∈R H×W×C , ⊙ represents Hadamard product, σ represents sigmoid activation function to obtain nonlinearity; Then, the non-local channel block is used to obtain the feature dependency information in the channel direction. The non-local channel block is defined as follows: where Z∈R H×W×C and O∈R H×W×C Represent the output attention feature map and input gated feature map, Conv 1×1 Represents a standard 1×1 2D convolution operation, which convolves the input gated feature map and resizes it to obtain g∈R HW×C ,θ∈R C×HW ,φ∈R HW×C , Represents matrix multiplication, softmax represents the softmax activation function operation, and the attention matrix is ​​C×C; Finally, the gate controller filters out the residual background layer information in the preliminary rain streak feature map and detects more fine rain streak information through the feature dependency information in the spatial direction and the feature dependency information in the channel direction obtained above.

8. The single image deraining method based on image prior and gated attention learning according to claim 1, characterized in that: The loss function of the network model based on image prior and gated attention learning adopts the sum of L1 loss function and edge function; The loss function is defined as follows: Where B represents the de-rained image after processing, G represents the real rain-free image, △ represents the Laplace operation, ‖BG‖1 represents the L1 loss function, represents the marginal loss function, and λ and ε represent hyperparameters.

9. A single image deraining system based on image prior and gated attention learning, characterized by: The invention comprises a processor and a memory, wherein non-transitory computer instructions are stored in the memory, and when the non-transitory computer instructions are executed by the processor, the single image deraining method based on image prior and gated attention learning as described in any one of claims 1 to 8 is executed.

Images