A lightweight image rain removal algorithm based on deep learning

Abstract

The application discloses a kind of light image rain removal algorithm based on deep learning, mainly used for image rain removal, solve the rain image of the day shooting rain removal problem and background recovery problem.Propose a kind of light image rain removal method based on PReNet model, to PReNet model as foundation, improve its network structure, in guarantee rain removal real-time, while, obtain better rain removal effect.For the background of PReNet existing false deletion problem, introduce dilated convolution (Dilated Convolution), it can be aggregated multi-scale context information without losing any resolution systematically, and does not need additional parameters, can guarantee the speed of network;For the gridding artifacts problem caused by dilated convolution, hybrid dilated convolution HDC (Hybrid Dilated Convolution) is used to solve.The application can guarantee better light rain rain stripe removal effect, while improving the removal effect of heavy rain rain stripe.

Description

Technical Field

[0001] This invention belongs to the fields of image processing and deep learning technology, specifically relating to a lightweight image deraining algorithm based on deep learning. Background Technology

[0002] Rainfall, a common weather phenomenon, causes visual degradation to videos and images captured under natural conditions. Rain streaks significantly reduce image visibility and contrast. Rainfall not only severely impacts the human visual perception system but also adversely affects many high-level visual tasks, such as object detection, image segmentation, and object tracking. These tasks largely assume clear weather, unaffected by rain, and use clean outdoor natural images as input. However, images obtained in harsh environments suffer from color distortion, texture blurring, reduced contrast, and non-uniform lighting, leading to image quality degradation and severely impacting the performance of computer vision algorithms. Therefore, removing rain streaks from individual images is a crucial issue in various visual tasks, and rain removal, as a preprocessing step in these systems, has been extensively studied.

[0003] Various algorithms for removing rain patterns have emerged. Most existing deep learning-based rain removal models improve performance by increasing network depth, which has proven to be an effective method for enhancing network performance. However, simply increasing network depth can lead to a large number of redundant parameters, which limits the application of rain removal models in real-world scenarios with limited storage resources. Therefore, in the field of image rain removal, how to maintain good rain removal results while making the model as lightweight as possible and improving processing speed has always been a research hotspot. Previous researchers have proposed many methods for processing small rain streaks, and the results are very good. However, these methods cannot fully handle images taken under heavy rain conditions. The paper "Clearing the skies: A deep network architecture for single-image rainremoval" first designed a CNN-based DerainNet, which automatically learns a nonlinear mapping function between clean and rainy images from the data. To improve the quality of image restoration, the authors also introduced knowledge from the field of image processing. Inspired by the great success of "Deep residual learning for image recognition", the paper "Removing rain from single images via a deep detail network" further proposed a deep detail network (DNN) to remove rain streaks. This network only takes high-frequency details as input and predicts rain residue and clean images. Studies have shown that removing background information from the network input is beneficial, which makes training easier and more stable. However, this method cannot handle large and sharp rain streaks. Since a single network may not be able to learn all the patterns in the training samples, the paper "Density-aware single image de-raining using a multi-stream dense network" further proposes a densely perceptual image de-training method using a multi-stream dense network (DID-MDN). By integrating the residual-aware classification process, DID-MDN can adaptively determine rain density information (heavy / medium / light). However, this model is too complex and cannot guarantee the real-time performance of deraining, affecting the performance of subsequent vision tasks. To alleviate this problem, the paper "Progressiveimage deraining networks: a better and simpler baseline" proposes a simple and effective progressive periodic deraining network (PReNet) by repeatedly unfolding a shallow ResNet with recurrence layers. However, this method has a major problem: it ignores the contextual information of large regions.

[0004] To design a lightweight rain removal model that can maintain the existing effect of removing small rain stripes and improve the effect of removing large rain stripes, this invention introduces dilated convolution and hybrid dilated convolution (HDC) methods on the basis of the PReNet model to ensure the real-time performance of rain removal and improve the effect of removing large rain stripes. Summary of the Invention

[0005] The purpose of this invention is to provide a lightweight image deraining algorithm based on deep learning, which can maintain the existing effect of removing light rain stripes and improve the effect of removing heavy rain stripes, while ensuring the real-time performance of deraining.

[0006] To achieve the above objectives, the present invention provides the following technical solution: a lightweight image deraining algorithm based on deep learning, comprising the following steps:

[0007] S1: Conduct research on the RainTrainH dataset and perform preprocessing operations on the RainTrainH dataset;

[0008] S1-1: Remove images from the RainTrainH dataset whose background content is the same as that of images in the validation set, and use the remaining rainy images as the training set.

[0009] S1-2: Select the Rain100H heavy rain dataset as the validation set;

[0010] S2: Based on the PyTorch deep learning framework, build the PReNet model and improve upon it;

[0011] S2-1: Replace the two ordinary convolutional layers in the ResBlock of the original model with dilated convolutions;

[0012] S2-2: Set the void ratio for dilated convolution to satisfy the characteristics of hybrid dilated convolution;

[0013] S3: Train the model using the SGD method in multiple rounds on the processed RainTrainH dataset;

[0014] S3-1: On the established network model, use stochastic gradient descent to train the model;

[0015] S3-2: Set different dilation rates to the hybrid dilated convolution to meet the requirements, and train again until the best results are obtained;

[0016] S4: Validate the trained model on the Rain100H heavy rain dataset to obtain the final results;

[0017] S4-1: If the results of this experiment are worse than the results of the previous experiment, continue to adjust the void ratio and train and validate the model.

[0018] S4-2: Repeatedly train and validate the model until a good rain removal effect is obtained;

[0019] S5: After adjusting the model, train it on the RainTrainL dataset and validate it on the Rain100L dataset.

[0020] To further facilitate model training, the images in both the training set and the validation set are resized to 128*128 pixels.

[0021] To further improve the rain removal effect, the PReNet model is improved in the following ways:

[0022] 1) The PReNet model consists of 6 stages, each with the same network structure and sharing the same network parameters to reduce the number of parameters. Furthermore, the stage-by-stage rain removal method is beneficial for removing large rain streaks. Each stage consists of four parts, namely f in layer, f recurrent layer, f res layer, f out layer;

[0023] 2) Input / Output: The input for the first stage is the concatenation of two original rain images. The input for each subsequent stage is the concatenation of the stage result and the original rain image. The output for each stage is the result of the following expression:

[0024] x t-0.5 =f in (x t-1 ,y)

[0025] s t =f recurrent (s t-1 ,x t-0.5 )

[0026] x t =f out (f res (s t ))

[0027] Where, x t-1 This is the output of the previous stage, where y is the original rain image and f is the output of the previous stage. in For the input convolutional layer, x t-0.5 For this stage f in The output of the layer, f recurrentFor LSTM layers, s t-1 For the previous stage f recurrent The output of the layer, s t For this stage f recurrent The output result, f res For the ResBlocks layer, f out For the output convolutional layer, x t This is the output result for this stage;

[0028] 3)f in Layer: f in It is a single-layer convolution with ReLU function, filter size of 3*3, patch size of 1*1, input of 6 channels, and output of 32 channels;

[0029] 4)f recurrent Layer: Long Short-Term Memory (LSTM) is a special type of Recurrent Neural Network (RNN) primarily used to solve the vanishing and exploding gradient problems during long sequence training. It controls the transmission state through gating, remembering information that needs to be stored for a long time and forgetting less important information. The input and output of an LSTM layer are both 32 channels.

[0030] 5)f res Layer: f res This is a crucial part for extracting the depth representation of rain streaks, containing 5 ResBlocks that share all parameters. In the original model, each ResBlock contains two ordinary convolutional layers, each followed by a ReLU function. Research has shown that dilated convolutions can systematically aggregate multi-scale contextual information without losing any resolution, which is beneficial for removing large rain streaks. Furthermore, it requires no additional parameters, ensuring network speed. Therefore, the two ordinary convolutional layers in the original ResBlocks are replaced with two dilated convolutional layers. Introducing dilated convolutions into the model introduces "mesh artifacts." To address this issue, Hybrid Dilated Convolution (HDC) is used. Dilated convolutions must satisfy three conditions:

[0031]

[0032] The design goal is to make M i <= K, where K is the kernel size, which is 3 in this invention; n is the number of dilated convolutions, which is 5 in this invention; r i The dilation rate of the i-th dilated convolution.

[0033] After multiple experiments, it was found that the best rain removal effect can be obtained by setting the dilated convolutions in the 5 ResBlocks to 1, 2, 5, 1, 2 respectively, and the convolution rate of the two dilated convolutions in each ResBlocks is the same. The input and output of each convolutional layer are 32 channels.

[0034] 6)f out Layer: f out with f in The layer structure is the same, it is also a single convolution with a ReLU function, the filter size is 3*3, the patch size is 1*1, the input is 6 channels, and the output is 32 channels.

[0035] To further improve the rain removal effect, the specific method for training the model in S3 is as follows:

[0036] 1) Set the number of training epochs to 200 and the batch size to 20;

[0037] 2) Set the initial learning rate to 0.001, and multiply the learning rate by 0.2 after 30, 50, and 80 rounds respectively;

[0038] 3) Negative SSIM loss is selected as the loss function of the model.

[0039] Beneficial effects:

[0040] This invention primarily addresses image deraining, resolving the issues of rain removal and background restoration in images taken on rainy days. It proposes a lightweight image deraining method based on the PReNet model. By improving the network structure of PReNet, it achieves better deraining results while maintaining real-time performance. To address the background deletion issue inherent in PReNet, dilated convolution is introduced. This method systematically aggregates multi-scale contextual information without losing any resolution and requires no additional parameters, ensuring network speed. To mitigate the gridding artifacts caused by dilated convolution, hybrid dilated convolution (HDC) is employed. This invention improves the removal of rain streaks from heavy rain while maintaining good results in removing light rain streaks. Attached Figure Description

[0041] Figure 1 This is a flowchart of the experimental process of the present invention;

[0042] Figure 2 This is the overall network model diagram of the present invention;

[0043] Figure 3A detailed structural diagram of each stage of the PReNet model;

[0044] Figure 4 An improved overall structure diagram of ResBlock;

[0045] Figure 4a This is a detailed structural diagram of the ResBlock in the original model;

[0046] Figure 4b This is a detailed structural diagram of the improved ResBlock. Detailed Implementation

[0047] The present invention will be further described below with reference to specific embodiments. Examples of the embodiments are shown in the accompanying drawings, wherein the same or similar reference numerals denote the same or similar elements or elements having the same or similar functions throughout.

[0048] Example 1: As Figures 1 to 4 As shown, the purpose of this invention is to provide a lightweight image deraining algorithm based on deep learning, which can maintain the original model's removal effect on light rain stripes and improve the removal effect on heavy rain stripes without adding any additional parameters.

[0049] This experiment was conducted on Colab. The training and validation sets were stored on Google Drive. The lab could access the training and validation sets by mounting the contents of Google Drive to Colab. The experimental code was stored on GitHub, and could be imported and used by git clone in Colab.

[0050] The process of the present invention and the technologies involved will be further described in detail below with reference to the invention content and the above-mentioned drawings:

[0051] Step 1: Based on the objective of this invention, primarily to improve the removal of rain streaks from heavy rain, the publicly available heavy rain dataset RainTrainH was selected for experiments. This dataset contains 1800 pairs of rain and no-rain images, effectively preventing overfitting. To make the experimental process more rigorous, the RainTrainH dataset needs to be preprocessed. The specific steps are as follows:

[0052] The RainTrainH dataset contains a wealth of rainy images, and each rainy image has a corresponding rain-free background image, making it effective for supervised rain removal experiments. However, 546 pairs of images in this dataset have the same background content as the validation set Rain100H. Therefore, to prevent other problems, these 546 pairs of images need to be removed from the training set, and the remaining 1254 pairs of images are selected as our training set. The preprocessed training set is still large enough to meet the experimental requirements.

[0053] The validation set was chosen from the Rain100H dataset, which contains 100 rainy images.

[0054] Resizing the images in the training and validation sets to 128*128 pixels is beneficial for training the model.

[0055] Step 2: Build the improved PReNet network on Colab using the PyTorch deep learning framework. The specific structure to be built is as follows:

[0056] We built a PReNet rain removal model in the PyTorch deep learning framework. PReNet is an efficient and lightweight image rain removal network that has emerged in recent years. It consists of six stages, each with the same network structure and sharing the same network parameters to reduce the number of parameters in the model. Furthermore, the stage-by-stage rain removal method is beneficial for removing large rain streaks. Each stage consists of four parts, namely f in layer, f recurrent layer, f res layer, f out The model's biggest advantage is its fast image processing speed, comparable to other deraining networks. However, experiments revealed that the original model suffers from background deletion errors and cannot fully restore the image background. Therefore, future improvements will be made to achieve better background restoration results.

[0057] Input / Output: The input for the first stage is the concatenation of two original rain images. The input for each subsequent stage is the concatenation of the stage result and the original rain image. The output for each stage is the result of the following expression:

[0058] x t-0.5 =f in (x t-1 ,y)

[0059] s t =f recurrent (s t-1 ,x t-0.5 )

[0060] xt =f out (f res (s t ))

[0061] Where, x t-1 This is the output of the previous stage, where y is the original rain image and f is the output of the previous stage. in For the input convolutional layer, x t-0.5 For this stage f in The output of the layer, f recurrent For LSTM layers, s t-1 For the previous stage f recurrent The output of the layer, s t For this stage f recurrent The output result, f res For the ResBlocks layer, f out For the output convolutional layer, x t This is the output result for this stage.

[0062] f in Layer: f in It is a single-layer convolution with ReLU function, a filter size of 3*3, a patch size of 1*1, a 6-channel input, and a 32-channel output.

[0063] f recurrent Layer: Long Short-Term Memory (LSTM) is a special type of recurrent neural network (RNN) primarily used to solve the vanishing and exploding gradient problems during long sequence training. It controls the transmission state through gating, remembering information that needs to be stored for a long time and forgetting unimportant information. In this invention, both the input and output of the LSTM layer are 32 channels.

[0064] f res Layer: f res This is a crucial part for extracting the depth representation of rain streaks, containing 5 ResBlocks that share all parameters. In the original model, each ResBlock contains two ordinary convolutional layers, each followed by a ReLU function. Research has shown that dilated convolutions can systematically aggregate multi-scale contextual information without losing any resolution, which is beneficial for removing large rain streaks. Furthermore, it requires no additional parameters, ensuring network speed. Therefore, the two ordinary convolutional layers in the original ResBlocks are replaced with two dilated convolutional layers. Introducing dilated convolutions into the model introduces "mesh artifacts." To address this issue, hybrid dilated convolution (HDC) is used. Dilated convolutions must satisfy three conditions:

[0065]

[0066] The design goal is to make Mi <= K, where K is the kernel size, which is 3 in this invention; n is the number of dilated convolutions, which is 5 in this invention; r i The dilation rate of the i-th dilated convolution.

[0067] After multiple experiments, it was found that setting the dilated convolutions in the 5 ResBlocks to 1, 2, 5, 1, and 1 respectively was effective.

[0068] 2. Furthermore, the two dilated convolutions in each ResBlock have the same convolution rate, which yields the best rain removal effect. Each convolutional layer has 32 channels for both input and output.

[0069] f out Layer: f out with f in The layer structure is the same, it is also a single convolution with a ReLU function, the filter size is 3*3, the patch size is 1*1, the input is 6 channels, and the output is 32 channels.

[0070] Step 3: After the improved PReNet network model is built, initialize the model, which mainly includes the following:

[0071] The number of training epochs is set to 200, and the batch size is set to 20.

[0072] The initial learning rate is set to 0.001. After 30, 50, and 80 rounds, the learning rate is multiplied by 0.2.

[0073] Choice of loss function. Experiments with the improved model using MSE loss and negative SSIM loss respectively revealed that the PSNR and SSIM scores obtained using negative SSIM loss were higher than those obtained using MSE loss. Therefore, negative SSIM loss was chosen as the loss function for the model.

[0074] Step 4: After completing the above steps, train the improved model on the RainTrainL dataset and validate it on the Rain100L dataset to obtain the rain-removed images and PSNR and SSIM index values, and save the results.

[0075] The above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can be made to the technical solutions of the present invention without departing from the spirit and scope of the technical solutions of the present invention, and all such modifications or substitutions should be covered within the scope of the claims of the present invention.

Claims

1. A lightweight image deraining algorithm based on deep learning, characterized in that, Comprise the following steps: S1: research on heavy rain data set RainTrainH, and pre-processing operation is carried out on heavy rain data set RainTrainH; S1-1: the picture background in heavy rain data set RainTrainH is removed, and the picture background content in the verification set is removed, and the remaining rain picture is used as the training set; S1-2: select Rain100H heavy rain data set as the verification set; S2: on the basis of PyTorch deep learning framework, PReNet model is built, and it is improved; S2-1: the two ordinary convolution layers contained in ResBlock in the original model are replaced by empty convolution; S2-2: the empty rate meeting the characteristics of hybrid empty convolution is set for the empty convolution; The PReNet model improvement contents are as follows: 1) the PReNet model is composed of 6 stages, the network structure of each stage is the same, and the 6 stages share the same network parameters, so as to reduce the number of parameters of the model, and the rain removal method of each stage is beneficial to the removal of heavy rain stripes, each stage is composed of four parts, which are f in layer, f recurrent layer, f res layer, f out layer; 2) input and output: the input of the first stage is the result of the cascade of two original rain images, and the input of each stage is the cascade result of the stage result and the original rain image; The output of each stage is the result of the following expression: x t-0.5 = f in (x t-1 ,y) s t = f recurrent (s t-1 , x t-0.5 ) x t = f out (f res (s t )) wherein x t-1 is the output result of the previous stage, y is the original rain image, f in is the input convolutional layer, x t-0.5 is the output result of the previous stage f in layer, f recurrent is the LSTM layer, s t-1 is the output result of the previous stage f recurrent layer, s t is the output result of the previous stage f recurrent layer, f res is the ResBlocks layer, f out is the output convolutional layer, x t is the output result of this stage; 3)f in Layer: f in is a layer of convolution with ReLU function, filter size is 3*3, patch size is 1*1, input is 6 channels, and output is 32 channels; 4)f recurrent Layer: long short-term memory (LSTM) is a special recurrent neural network (RNN) that solves the problems of gradient disappearance and gradient explosion in long sequence training process, controls the transmission state through the gating state, remembers the information that needs to be remembered for a long time, and forgets the unimportant information; wherein, the input and output of the LSTM layer are both 32 channels; 5)f res Layer: f res is the key part of extracting rain streak depth representation, which contains 5 ResBlocks, and the 5 ResBlocks share all parameters, and each ResBlock in the original model contains two ordinary convolution layers, and each convolution layer is followed by a ReLU function; replace the two ordinary convolution layers in the ResBlock in the original model with two hollow convolution layers; adopt hybrid hollow convolution HDC to solve the grid artifact problem, wherein the hollow convolution needs to meet three conditions: The design goal is to make M i < = K, where K is the kernel size, n is the number of dilated convolutions, r i dilation rate of the i-th dilated convolution; After a plurality of groups of experiments, it is found that the empty convolution in the five ResBlocks is set to 1, 2, 5, 1 and 2 respectively, and the convolution rates of the two empty convolutions in each ResBlock are the same, and the best rain removal effect can be obtained, and the input and output of each convolution layer are 32 channels; 6)f out Layer: f out With f in The layer structure is the same, which is a layer of convolution with a ReLU function, a filter size of 3*3, a patch size of 1*1, an input of 6 channels, and an output of 32 channels. S3: the model is trained, and the RainTrainH data set is used for training; S3-1: the random gradient descent method is used to train the model on the built network model; S3-2: different empty rates meeting the requirements are set for the hybrid empty convolution, and the training is carried out again until the best effect is obtained; S4: the trained model is verified on Rain100H heavy rain data set, and the final result is obtained; S4-1: if the experimental result is worse than the last experimental result, adjust the empty rate and train and verify the model; S4-2: the model is repeatedly trained and verified until a good rain removal effect is obtained; S5: after adjusting the model, the model is trained on the light rain data set RainTrainL and verified on the light rain data set Rain100L.

2. The deep learning based lightweight image deraining algorithm according to claim 1, characterized in that, The size of the picture in the training set and the picture in the verification set is adjusted, and the size is reset to 128*128 pixels.

3. The deep learning based lightweight image deraining algorithm according to claim 2, characterized in that, In the S3, the specific way of training the model is as follows: 1) the training rounds epoch is set to 200, and the batch size batch_size is set to 20; 2) the initial learning rate is set to 0.001, and the learning rate is multiplied by 0.2 when 30, 50 and 80 rounds are reached respectively; 3) negative SSIM loss is selected as the loss function of the model.

Images