A memory-oriented single picture rain removal method based on transformer

Abstract

The application discloses a single picture rain removing method based on a memory-oriented Transformer, characterized by an encoder-decoder structure with a self-supervised memory module, which can well process input pictures and better extract required features, wherein the self-supervised memory module is a neural network with memory, which can record various forms of rainfall, wherein each entry in the memory corresponds to a prototype feature of a rain pattern; a self-training mechanism is added to the self-supervised memory module, enhancing the adaptability of the algorithm to natural rain pictures. Compared with the prior art, the application can remove more rain streaks with different appearances, restore clearer background scenes, and better retain the structure and details of the background. The addition of the self-training mechanism makes the algorithm more adaptive to natural rain pictures, and good results can also be achieved on natural rain pictures.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of picture preprocessing, in particular to a single picture rain removal method based on memory-oriented Transformer. BACKGROUND

[0002] The information such as rain marks existing in the rain picture can make the picture background highly blurred, and further affect the accuracy of target recognition. For an automatic driving car, its visual-based perception functions such as object detection, recognition and semantic / instance segmentation need to accurately learn the features of the urban street scene image. Rain, as the most common severe weather condition, can greatly reduce the visual quality of the image and block the background objects. These reductions in visibility have a negative impact on image feature learning and can cause many computer vision systems to malfunction. In addition to automatic driving, many other applications, such as outdoor monitoring systems, can also degrade when presenting images containing rain and haze and other artifacts. These reasons make rain removal a very desirable technology that can remove the blurred visual effects caused by rain from the image.

[0003] At present, the image rain removal algorithm mainly includes video and single image rain removal algorithm. The information and data collected by the above system are mostly video information. The video is composed of multiple images of continuous frames. The pixels affected by rain marks or raindrops can be detected by combining the characteristics of image continuous frames. The single image is static information, which lacks the spatiotemporal characteristics of rain mark or raindrop changes. Therefore, the research on the single image rain removal algorithm is more meaningful. The method based on deep learning shows significant improvement in image rain removal by using large-scale paired data of synthetic data sets. However, since the various manifestations of real rain marks may be different from the synthetic training data, it is challenging to directly extend the existing methods to real scenes.

[0004] The existing image rain removal algorithm has poor background and scene recovery clarity, and cannot better preserve the structure and details of the background. Since the various manifestations of real rain marks may be different from the synthetic training data, the existing methods are all trained and tested based on synthetic rain pictures, and cannot be well migrated to natural rain pictures. In order to solve the above problems, the present application is proposed. SUMMARY

[0005] The purpose of the present application is to provide a Transformer-based memory-oriented single picture rain removal method to overcome the shortcomings of the prior art. The method uses the characteristic of transformer focusing on global features to better extract the information of rain streaks in the rain picture. The self-supervised memory module with memory function is used to record various forms of rain streaks. Finally, the self-training mechanism is used to transfer the constraint information on the synthetic rain picture to the natural rain picture, so that it can be better trained. The method can remove more rain streaks with different appearances and restore clearer background scenes while better preserving the structure and details of the background. The addition of the self-training mechanism makes the algorithm more adaptive to natural rain pictures and achieves good results on natural rain pictures. The method not only improves the rain removal effect of the picture, but also solves the problem that existing algorithms trained on synthetic training data cannot be effectively transferred to real scenes.

[0006] The specific technical solution to achieve the purpose of the present application is: a Transformer-based memory-oriented single picture rain removal method, characterized by using the characteristic of transformer focusing on global features to better extract the information of rain streaks in the rain picture, then using a self-supervised memory module with memory function to record various forms of rain streaks, and finally using a self-training mechanism to transfer the constraint information on the synthetic rain picture to the natural rain picture, so that it can be better trained, comprising the following steps:

[0007] 1) In the initialization stage, a pre-trained model swin_large_patch4_window12_384_22k.pth is used, the input size is 384*384, the number of training samples is 4 each time, and the maximum number of training is 300.

[0008] 2) The 384*384 picture is sent to the encoder of the transformer. Patch Partition performs down-sampling on the input image, reducing the width and height by 1 / 4 and adjusting the channel to 48. Linear Embedding adjusts the output of Patch Partition in the channel dimension to C, and then processes the feature map through Layer Normalization.

[0009] 3) The Layer Normalization processed feature map is passed through the encoder unit composed of Swin Transformer Blocker and Patch Merging, then the global relationship is modeled and hierarchical feature transformation is performed, while the width and height of the feature map are reduced and the receptive field is expanded. Patch Merging is a 4x4 single-channel feature map, and Patch Merging will divide each 2x2 adjacent pixel into a patch, then the same position (same color) pixels in each patch are spliced together to get four feature maps. Then the four feature maps are concatenated in the depth direction, and then passed through a LayerNorm layer. Finally, a fully connected layer is used to make linear changes in the depth direction of the feature map, changing the depth of the feature map from C to C / 2.

[0010] 4) Map the feature value in step 3 to z(x) of the inner supervised memory module, z(x) is equivalent to a query to find the most relevant items in memory, and then use soft-attention to combine them, and then update the memory module using self-supervised update and soft attention reading update. Among them, the memory module M ∈ R m×c consists of m memory items, each item has a dimension of e i ∈ R c and the encoding z(x) ∈ R c×h×w is defined as follows (a):

[0011]

[0012] Calculate the cosine similarity s ij (x) between the ith memory item and the jth column vector z j (x) of z(x), which is defined as follows (b):

[0013]

[0014] Then, use the following (c) to retrieve the most relevant memory item e j (x) of z(x): k(j) (x) with the largest s ij (x) is assigned to k (j) (x):

[0015] k (j) (x) = argm i axs ij (x) (c).

[0016] Finally, update the memory item e i Based on the query z j (x) with the most relevant item where e i is defined by the following (d) equation:

[0017]

[0018] Soft attention reading is a way different from traditional hard attention reading, mainly used to solve the problem of gradient back propagation, and the similarity matrix S(x) = {s ij (x) | i = 1,...,m, j = 1,...,n} is calculated again by the above formula (a) on the updated memory item. Then, the attention A = {a ij i = 1,...,m, j = 1,...,n} is obtained by the softmax operation, where a ij is defined by the following (e) equation:

[0019]

[0020] Finally, the memory-based representation is calculated based on the attention-based memory item aggregation by the following (f) equation:

[0021]

[0022] 5) Online training of the rain removal network f θ using data with labels, and then using an additional target network f ξ to generate pseudo labels for unlabeled data. It includes two processes, one is a supervised process and the other is an unsupervised process.

[0023] Supervised rain removal is to use data with labels to online train the rain removal network f θ in the network, where the optimization target is the pixel-level L1 loss function defined by the following (g) equation:

[0024] L SU = f θ (x l )- y l1 (g).

[0025] Unsupervised rain removal adopts a momentum encoder for self-supervised representation learning, and uses an additional target network f ξ to generate pseudo labels for unlabeled data, f ξ is an online network f θ updated by exponential moving average. After each training step, ξ is updated as follows by the following (h) equation:

[0026] ξ←vξ+(1-v)θ(h).

[0027] where v ∈ [0, 1] is the decay rate.

[0028] For the rainy image set X U , any one of the unlabeled images x u , the target network f ξ produces its corresponding pseudo-label f ξ (x u ), which constitutes the rainy image set X U and the corresponding pseudo-label set Y P , and a set of rainy residual networks defined by the following (i) formula is obtained:

[0029] R = {x-y(x, y) ∈ (X L , Y L )∪(X U , Y U )} (i).

[0030] where (X L , Y L ) and (X U , Y U ) are the synthetic pair set and the pseudo pair set, respectively.

[0031] Finally, by data augmentation on X L , Y L , X U , Y U , etc. image sets, a noisy data set X N and a rainy residual set r are obtained. More accurately, a picture x is randomly extracted, and its corresponding label is (the label of the clean image is itself), and a residual image r ∈ R. Where the noise picture x n is calculated by the following (j) formula:

[0032]

[0033] where α is a random value sampled from a uniform distribution U(a, b) (specifically, a and b here are 0.5 and 1.1). T(·) is a clip function to ensure that x n and have the same range. Thus, the paired noisy data is obtained. The pixel-level L1 loss is used for augmented data, and is defined by the following (k) formula:

[0034]

[0035] Self-training of enhanced noisy data can enrich the rainfall pattern during training and improve the robustness of real rain removal.

[0036] W adopts a total variation regularization term to the recovered background image f θ (x n ) is smoothed by the following (l) formula:

[0037] L TV = f θ (x n ) TV (l).

[0038] The total target f of the online network θ is given by the following (m) formula:

[0039] L total = λ1L SU + λ2L UN + λ3L TV (m).

[0040] Where λ 1,2,3 is the hyperparameter balancing each item.

[0041] 6) Decode the feature values generated in step 5) and realize dimension transformation using 3 decoding blocks and Dim Expanding, gradually recovering the image.

[0042] 7) Continuously repeat steps 1) to 6) to continuously optimize the accuracy of the caching decision and finally obtain accurate results.

[0043] Compared with the prior art, the present application has the advantages of removing more different appearance rain streaks, recovering clearer background scenes, better preserving the structure and details of the background, and the self-training mechanism makes the algorithm more adaptive to natural rain pictures, and good results can also be achieved on natural rain pictures. BRIEF DESCRIPTION OF DRAWINGS

[0044] Figure 1 The algorithm framework of the present application. DETAILED DESCRIPTION

[0045] The present application will be further described and explained in detail with specific implementation as follows:

[0046] Example 1

[0047] Referring to Figure 1 , the present application performs single-picture rain removal according to the following steps:

[0048] 1) Initialization stage, using a pre-trained model: swin_large_patch4_window12_384_22k.pth, input size is 384*384, the number of training samples is 4 each time, and the maximum number of training is 300.

[0049] 2) Send the 384*384 picture into the encoder of the transformer, and the Patch Partition performs down-sampling on the input image, reducing the original input image H, W, C by 1 / 4 in width and height, and adjusting the channel to 48. After adjusting the output of the Patch Partition in the channel dimension to C, the feature map is processed through the Layer Normalization.

[0050] 3) The Layer Normalization processed feature map is passed through the Swin Transformer Blocker and the Patch Merging composed encoder unit, then the global relationship is modeled and the hierarchical feature transformation is performed, while the width and height of the feature map are reduced and the receptive field is expanded. Patch Merging is a 4x4 single-channel feature map, and Patch Merging will divide each 2x2 adjacent pixel into a patch, then the same position (same color) pixels in each patch are spliced together to obtain four feature maps. Then the four feature maps are concatenated in the depth direction, and then passed through a LayerNorm layer. Finally, a fully connected layer is used to make linear changes in the depth direction of the feature map, changing the depth of the feature map from C to C / 2.

[0051] 4) Map the feature values in step 3 into z(x) of the self-supervised memory module, z(x) is equivalent to a query to find the most relevant items in memory, and then combine them using soft-attention, and then update the memory module using self-supervised update and soft-attention reading update. Where M∈R m×c consists of m memory items, where the dimension of each item is e i ∈R c and the encoding z(x)∈R c×h×w z(x) is defined as follows (a):

[0052]

[0053] Calculate the cosine similarity s ij(x) the ith memory item and the jth column vector of z(x) j (x) is defined by the following (b) formula:

[0054]

[0055] Then, the z j (x) most relevant memory item Set s ij (x) the largest number to k (j) (x):

[0056]

[0057] Finally, update the memory item e i Based on the query z j (x) has the most relevant item Where e i Defined by the following (d) formula:

[0058]

[0059] Soft attention reading is different from traditional hard attention reading, mainly used to solve the problem of gradient back propagation, and the similarity matrix S(x) is calculated again by the above formula (a) on the updated memory item ij (x) | i = 1,...,m, j = 1,...,n}. Then, the attention A = {a ij i = 1,...,m, j = 1,...,n} is obtained by the softmax operation, where a ij Defined by the following (e) formula:

[0060]

[0061] Finally, the memory-based representation is obtained based on the attention-based memory item aggregation calculation by the following (f) formula:

[0062]

[0063] 6) Online training of rain removal network f using data with labels θ Then, an additional target network f ξ Generates pseudo-labels for unlabeled data. It includes two processes, one is a supervised process, and the other is an unsupervised process.

[0064] Supervised rain removal is to use data with labels to train the rain removal network f θ In the network, where the optimization target is the pixel-level L1 loss function, and is defined by the following (g) formula:

[0065] L SU =f θ (x l )-y l1 (g)

[0066] Unsupervised rain removal employs a momentum encoder for self-supervised representation learning, using an additional target network f. ξ Generate pseudo-labels for unlabeled data, f ξ The online network f is updated using exponential moving averages. θ After each training step, ξ is updated by the following equation (h):

[0067] ξ←vξ+(1-v)θ(h).

[0068] Where v∈[0,1] is the decay rate.

[0069] For the rainy day image set X U Any unlabeled image x u Using the target network f ξ Generate its corresponding pseudo-label f ξ (x u ), forming a rainy day image set X U The corresponding pseudo-label set Y P This yields a set of rain-inclusive residual networks represented by the following equation (i):

[0070] R = {xy(x,y)∈(X)} L ,Y L )∪(X U ,Y U )}(i).

[0071] Among them, (X) L ,Y L ) and (X U ,Y U ) are the synthetic pairing set and the pseudo pairing set, respectively.

[0072] Finally, through X L ,Y L ,X U ,Y P Data augmentation was performed on the image set to obtain a noisy dataset X. N And the residual set r of the rain. More precisely, a randomly selected image. Its corresponding tag is (The label of the clean image is itself), and a residual image r∈R. Wherein, the noisy image x... n The following formula (j) is used to calculate:

[0073]

[0074] Here, α is a random value sampled from a uniform distribution U(a,b) (specifically, a and b are 0.5 and 1.1 here). T(·) is a clamping function to ensure that x n and They have the same range. Therefore, pairs of noise data were obtained. The augmented data uses pixel-level L1 loss, defined by the following equation (k):

[0075]

[0076] Enhancing the self-training of noise data can enrich rainfall patterns during training and improve robustness to real rainwater removal.

[0077] W applies a total variational regularization term to the recovered background image f. θ (x n Smoothing is performed using the following formula (l):

[0078] L TV =f θ (x n ) TV (l).

[0079] The overall goal of online networks f θ Defined by the following formula (m):

[0080] L total =λ1L SU +λ2L UN +λ3L TV (m).

[0081] Where, λ 1,2,3 It is about balancing the hyperparameters of each project.

[0082] 6) Use 3 decoding blocks and Dim Expanding to decode the feature values ​​generated in step 5) and perform dimensional transformation to gradually restore the image.

[0083] 7) Repeat steps 1) through 6) continuously to optimize the accuracy of caching decisions and eventually obtain accurate results.

[0084] This invention employs a swim transformer as its basic framework, leveraging its global characteristics to collect rain streaks. A self-supervised memory module is used to better learn and memorize rain streak information. Therefore, the algorithm processes rain streak information more comprehensively, resulting in clearer images. It can remove rain streaks with more varied appearances, restore a clearer background scene, and better preserve the background structure and details. To address the difficulty of transferring existing algorithms to real rain images, this invention proposes a self-training mechanism. First, labeled data is used for training. Then, the constraint information is transferred to an unlabeled dataset, generating a pseudo-label set. This allows for better training on real rain images, thus enabling more effective processing of real rain images.

[0085] The above description is merely a specific embodiment of the present invention and is not intended to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. A memory-oriented single picture rain removal method based on Transformer, characterized in that, The method utilizes the characteristic of the Transformer focusing on global features, uses a self-supervised memory module with memory function to record various forms of rain streaks, and transfers the constraint information on the rain picture to the natural rain picture, so that it can be better trained. The single picture rain removal specifically includes the following steps: 1) initialization stage The swin_large_patch4_window12_384_22k.pth is used as the pre-training model, and the network weight of the swin transformer is pre-trained and initialized by using the given parameters; 2) the training set is sent into the encoder of the Transformer, the Patch Partition performs down-sampling on the input image, the original input image H, W, C is reduced by 1 / 4 in width and height, and the channel is adjusted to 48, the output of the Patch Partition is adjusted to C in the channel dimension by the Linear Embedding, and the feature map is processed by the Layer Normalization; 3) the Layer Normalization processed feature map is sent into the encoder unit composed of the Swin Transformer Blocker and the Patch Merging, then the global relationship is modeled and the hierarchical feature transformation is performed, and the width and height of the feature map are reduced, the receptive field is expanded, and new feature values are obtained; 4) Map the feature value in step 3) into z(x) of the inner supervised memory module, then update the memory memory module using self-supervised update and soft attention reading update ; 5) Online training of the rain removal network using data with labels and then using an additional target network generate pseudo labels and new feature values for unlabelled data; 6) the feature values generated in step 5) are decoded by using 3 decoding blocks and Dim Expanding, and the dimension is transformed, and the image is gradually restored; 7) repeat steps 1) to 6) to continuously update the network, and obtain the optimal network weight; The memory-in-memory module in the step 4) consists of m memory items, where each item has a dimension of with encoding The encoding z(x) is defined by the following (a) formula: （a）； wherein is the nth child of z(x); The update memory module Specifically comprises: 1) the i-th storage item and the j-th column vector of z(x) the cosine similarity of is calculated by the following (b) formula: （b）； 2) search using the following (c) formula the most relevant memory item put assign the largest number to : （c）； 3) based on query having the most relevant memory item to update the memory item , the updated memory item is defined by the following equation (d): （d）； wherein is 1, not equal to 0, ; The soft attention reading of the step 4) is to calculate the similarity matrix of the updated memory item with the above (a) formula, and then obtain the attention by the softmax operation Wherein the specific calculation method of a ij is shown in formula (e):​ （e）； Finally, the attention-based memory item aggregation computation defined by the following (f) formula is obtained based on the memory-based representation : （f）； The online training in step 5) adopts a self-training mechanism of supervised rain removal and unsupervised rain removal, and specifically includes the following steps: 1) Supervised de-raining using labeled data for online training of de-raining network where the optimization objective is a pixel-wise LI loss function, whose loss function L su is defined by the following (g) equation: （g）； 2) Unsupervised rain removal uses momentum encoder for self-supervised representation learning, with an extra target network Pseudo labels are generated for unlabelled data, is an online network updated with exponential moving average , momentum is updated by the following equation (h): （h）； wherein is the decay rate; 3) for the rainy day image set any one of the unlabelled images , using the target network producing its corresponding pseudo-label , constituting the rainy day image set the corresponding pseudo-label set a set of rainy residual networks represented by the following (i) formula: （i）； wherein, and are synthetic pairing sets and pseudo pairing sets, respectively; 4) Through data enhancement on image sets such as , , , , a noisy data set , and a rain residual set r, a random image is extracted, and the corresponding label is , that is, the label of the clean image is itself, and a residual image , wherein the noise picture is calculated by the following formula (j): （j）； wherein is a random value sampled from a uniform distribution over the interval [0, 1]. 5) Use pixel-level L1 loss for enhanced data, whose loss function is defined by the following (k) formula: （k）； 6) total variation regularization term is used for the recovered background image The smoothing process is performed by the following (l) formula, and the loss function is defined as : （l）； 7) Overall goal of the online network defined by (m) below: （m）； Wherein, λ1, λ2, λ3 are the hyperparameters for balancing each item.

2. The Transformer-based memory-oriented single-image rain removal method according to claim 1, wherein, The size of the pre-training model input image in step 1) is 384*384, the number of training samples is 4 each time, and the maximum number of training is 300.

3. The Transformer-based memory-oriented single-image deraining method according to claim 1, wherein, The Patch Merging in step 3) is a single-channel feature map with a size of 4x4, the Patch Merging divides each 2x2 adjacent pixel into a patch, then splices the pixels at the same position, i.e. the same color, in each patch together to obtain 4 feature maps, then the feature maps are concatenated in the depth direction, then a fully connected layer is used in the LayerNorm layer to perform linear transformation in the depth direction of the feature map, and the depth of the feature map is changed from C to C / 2.

Images