Training method of single image rain and snow removal general module and single image rain and snow removal method

Abstract

The invention relates to a training method for a general module for removing rain and snow from a single image and a method for removing rain and snow from a single image, belonging to the field of image processing technology, and in particular to a dual U-Net single image method for removing rain and snow based on GAN; the invention solves the problems of occupying too many vehicle resources, over-smoothing of images after removing rain and snow, and rough processing of image edge details existing in existing rain or snow removal methods; the training method comprises the following steps: using real rainy day images, real snowy day images, and real sunny day images in a training set to train a joint network to obtain a trained joint network; using the sunny day image generator of the trained joint network as a general module for removing rain and snow from a single image. The training method for a general module for removing rain and snow from a single image and the method for removing rain and snow from a single image are suitable for removing rain and snow from a single image and for removing rain and snow from lane line recognition in autonomous driving.

Description

Technical Field

[0001] The present invention relates to the technical field of image processing, and in particular to a GAN-based dual U-Net single-image rain and snow removal method. Background Art

[0002] In the field of autonomous driving, the lane line detection task of a vehicle generally involves first collecting a video or image in front of the vehicle, then preprocessing the video or image, and then identifying the lane line. When preprocessing the video or image, weather factors, especially bad weather such as rainy and snowy days, are important factors affecting the accurate identification of lane lines.

[0003] Severe rain and snow weather often obstruct important target features in the image, resulting in reduced contrast, blurred imaging, and loss of detail information, which interferes with image processing results such as image recognition or classification. Therefore, technicians in this field have developed a large number of rain or snow removal algorithms to process the collected images and videos.

[0004] At the same time, with the development of power battery technology, more and more new energy vehicles have appeared in high-latitude areas such as northern China, which has led to the shift of autonomous driving technology from mainly considering the driving environment in the south to considering the driving environment in both the north and the south. Compared with the driving environment in the south, the driving environment in the north is more severe. For example, in the lane line detection task, the driving environment in the south generally only has bad weather such as rainy days, but in the north, there are not only rainy days and snowy days, but also mixed rain and snow weather.

[0005] As the application environment of autonomous driving technology changes, the existing rain and snow removal methods are gradually unable to meet the needs, exposing the following defects:

[0006] (1) Among the existing rain and snow removal methods, there are few general rain and snow removal algorithms, and most of them are single rain or snow removal algorithms. In order to adapt to the rainy and snowy driving environment in the north, it is necessary to integrate multiple algorithms in the vehicle to complete the simultaneous rain and snow removal of the image. However, the vehicle's hardware resources, such as storage and memory, are limited. Loading multiple algorithms will inevitably lead to excessive occupation of vehicle resources, which is not suitable for unmanned driving scenarios with limited hardware resources.

[0007] In addition, some existing general rain and snow removal methods often have poor performance and often use more parameters or more complex network model structures, which still occupy too many vehicle resources and are not suitable for unmanned driving scenarios with limited hardware resources.

[0008] (2) There are relatively few existing snow removal algorithms, and most of them are for video. Compared with single-image (i.e., single-image) rain and snow removal algorithms, video rain and snow removal methods are more complex and difficult to maintain temporal consistency. They have high storage and memory requirements for autonomous driving vehicles, occupy too many vehicle resources, and are not suitable for autonomous driving scenarios with limited hardware resources.

[0009] (3) The existing methods for removing rain and snow have a large difference in the quality of image detail texture compared with the original image, which causes the image after rain and snow removal to appear overly smooth, resulting in a poor visual experience.

[0010] (4) Existing rain removal methods still have the problem of being unable to retain texture and edge details while removing rain streaks, resulting in rough processing of image edge details. Summary of the invention

[0011] The present invention proposes a training method for a universal module for single-image rain and snow removal and a single-image rain and snow removal method, which solves the problems of occupying too many vehicle resources, over-smoothing of images after rain and snow removal, and rough processing of image edge details in existing rain or snow removal methods.

[0012] The training method of the single-image rain and snow removal universal module of the present invention has the following technical solutions:

[0013] The training method comprises the following steps:

[0014] S1, obtaining a collected original image data set; the original image data set includes rainy day images, snowy day images and sunny day images;

[0015] S2, preprocessing the collected original image data set to obtain a training set;

[0016] S3, obtaining a joint network including three generative adversarial network modules; each generative adversarial network module includes a generator and a discriminator; wherein:

[0017] The first generative adversarial network module: its generator is a rainy day image generator, which is used to generate fake rainy day images; its discriminator is a rainy day image discriminator, which is used to determine whether the input image is a real rainy day image;

[0018] The second generative adversarial network module: its generator is a snow image generator, which is used to generate fake snow images; its discriminator is a snow image discriminator, which is used to determine whether the input image is a real snow image;

[0019] The third generative adversarial network module: its generator is a sunny image generator, which is used to generate fake sunny images; its discriminator is a sunny image discriminator, which is used to determine whether the input image is a real sunny image;

[0020] S4, using the real rainy day images, real snowy day images and real sunny day images in the training set to train the joint network to obtain a trained joint network;

[0021] The trained sunny image generator of the joint network is used as a general module for single-image rain and snow removal.

[0022] Further, a preferred embodiment is provided, wherein step S4 is as follows:

[0023] An iterative training algorithm is used to update the model parameters of the joint network; each iterative training process includes the following steps:

[0024] S4.1, inputting the rainy day images in the training set into the sunny day image generator to obtain a false sunny day image generated from the rainy day image;

[0025] The false sunny image generated by the rainy image is input into the sunny image discriminator for identification, and the identification result is recorded in the historical identification result;

[0026] The false sunny image generated by the rainy image is input into the rainy image generator to obtain a new false rainy image, which is then input into the rainy image discriminator for identification;

[0027] S4.2, inputting the snowy images in the training set into the sunny image generator to obtain a false sunny image generated from the snowy images;

[0028] The false sunny image generated by the snowy image is input into the sunny image discriminator for identification, and the identification result is recorded in the historical identification result;

[0029] The false sunny image generated by the snowy image is input into the snowy image generator to obtain a new false snowy image, which is then input into the snowy image discriminator for identification;

[0030] S4.3, input the sunny images in the training set into the rainy image generator to obtain a false rainy image generated from the sunny image.

[0031] The false rainy day image generated by the sunny day image is input into the rainy day image discriminator for identification.

[0032] The false rainy day image generated by the sunny day image is input into the sunny day image generator to obtain the false sunny day image generated by the false rainy day image generated by the sunny day image.

[0033] The false sunny image generated by the false rainy image generated by the sunny image is input into the sunny image discriminator for identification, and the identification result is recorded into the historical identification result;

[0034] At the same time, the same sunny image in the training set is input into the snow image generator to obtain a fake snow image generated from the sunny image.

[0035] The false snowy image generated from the sunny image is input into the snowy image discriminator for identification.

[0036] The false snowy day image generated from the sunny day image is input into the sunny day image generator to obtain the false sunny day image generated from the false snowy day image generated from the sunny day image.

[0037] The false sunny image generated by the false snowy image generated by the sunny image is input into the sunny image discriminator for identification, and the identification result is recorded into the historical identification result;

[0038] S4.4, determine whether the historical identification results meet the given iterative training end conditions:

[0039] If the given iterative training end condition is met, the iterative training algorithm ends, and a trained joint network is obtained, and the sunny image generator is used as a general module for single image rain and snow removal;

[0040] Otherwise, update the model parameters of the joint network and proceed to the next iterative training process.

[0041] Further, a preferred embodiment is provided, wherein the given iterative training end condition is:

[0042] In the input image identified as a true sunny image by the sunny image discriminator:

[0043] The proportion of false sunny images ranged from 45% to 55%.

[0044] Furthermore, a preferred embodiment is provided, wherein the generator of each generative adversarial network module comprises two cascaded U-Net networks for retaining the details of the input image.

[0045] Furthermore, a preferred embodiment is provided, wherein the generator of each generative adversarial network module also includes an advanced attention module; the advanced attention module is used to prevent the image generated by the generator from being overly smooth.

[0046] Furthermore, a preferred embodiment is provided, in which the loss function adopted by each generative adversarial network module is a joint loss function:

[0047] ;

[0048] in, is the dual U-Net loss function; is the structural similarity loss function; is the cycle consistency loss function.

[0049] Further, a preferred embodiment is provided, wherein step S2 is as follows:

[0050] Each image in the original image dataset is cropped according to a pixel resolution of 480×320 to obtain cropped rainy day images, snowy day images, and sunny day images, and their collection is the training set.

[0051] The present invention also proposes a single image rain and snow removal method, and its technical solution is as follows:

[0052] The method comprises the following steps:

[0053] Step 1, obtaining the original image to be processed;

[0054] Step 2, preprocessing the original image to be processed;

[0055] Step 3, using the preprocessed original image to be processed as the input of the single image rain and snow removal general module trained by the training method of the single image rain and snow removal general module;

[0056] The false sunny image output by the single-image rain and snow removal general module is used as the image after rain and snow removal.

[0057] The present invention also proposes a computer program product, and its technical solution is as follows:

[0058] A computer program product includes a computer program or instructions, which, when executed by a processor, implement the steps of the training method of the above-mentioned single-image rain and snow removal general module.

[0059] The present invention also proposes a computer storage medium, and its technical solution is as follows:

[0060] A computer storage medium stores a computer program, and when the computer program is run, the training method of the above-mentioned single-image rain and snow removal general module is executed.

[0061] The present invention has the following beneficial effects:

[0062] 1. The training method of the single-image rain and snow removal universal module of the present invention trains three GANs so that the generators and discriminators of different GANs continuously exchange model parameters and compete with each other, and finally reach the optimal state of the network. The sunny image generator obtained at this time is the final single-image rain and snow removal universal module. The single-image rain and snow removal universal module has a simple network and a small number of parameters, which is suitable for unmanned driving scenarios with limited vehicle hardware resources. A single module can meet the general rain and snow removal requirements, and there is no need to develop multiple modules for different weather environments, which reduces the resource consumption of system deployment and improves the practicality and robustness of the system.

[0063] 2. The training method of the single-image general module for removing rain and snow described in the present invention, by adopting two cascaded U-Net networks and taking advantage of the U-Net network's strong learning ability and simple structure, can not only accurately remove more rain streaks but also retain the details of the image.

[0064] 3. The training method of the general module for removing rain and snow from a single image described in the present invention introduces an advanced attention module (AAM) into the generator of the generative adversarial network module. The advanced attention module is used to obtain an attention map based on the input image, extract different types of prior information, and make the network model focus on important areas blocked by rain patterns and snow patterns, thereby improving the feature extraction ability of the network model when processing rainy and snowy weather images, avoiding excessive smoothness of the generated rain and snow removal images, and improving the visual experience.

[0065] 4. The training method of the single-image rain and snow removal universal module described in the present invention constrains the algorithm model through a joint loss function, restores details in the image (such as edge details), so as to generate higher quality and more natural images, improve the performance of the algorithm model in image rain and snow removal tasks, and maintain the stability and reliability of the algorithm model.

[0066] 5. The training method of the single-image rain and snow removal universal module described in the present invention reduces the difference between the image generated by the generator and the real rain and snow removal image through the dual U-Net loss function, making the generated image closer to the real image, and is used to constrain the output of the first U-Net network of the generator to make the network model more stable.

[0067] 6. In the training method of the general module for removing rain and snow from a single image described in the present invention, the structural similarity loss function is used to measure the distance between two images, and combined with the comprehensive consideration of brightness, contrast and structure, the structure, texture and other information in the image are maintained to prevent the generated image from being over-smoothed, so that the generated image is visually closer to the real image and improve the visual experience.

[0068] 7. In the training method of the general module for removing rain and snow from a single image described in the present invention, the cycle consistency loss function is used to ensure that the image can be restored to the original image as much as possible after two conversions (i.e., from the source domain to the target domain, and then from the target domain back to the source domain). The cycle consistency loss function maintains the consistency of the image content by forcing the generator to maintain the basic structure and content of the image during the generation and degeneration process, so as to avoid the generator only focusing on the target domain features and ignoring the content of the source domain, resulting in the loss or addition of too much irrelevant information when generating the image, thereby improving the quality and credibility of the generated image, and at the same time, it can also improve the stability of the model and avoid model collapse or the generation of unreasonable images.

[0069] The training method of the single-image rain and snow removal universal module and the single-image rain and snow removal method described in the present invention are suitable for single-image rain and snow removal and lane line recognition in automatic driving. BRIEF DESCRIPTION OF THE DRAWINGS

[0070] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings required for use in the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0071] Figure 1 1 is a schematic diagram of the structure of a generative adversarial network module in one embodiment of the present invention; wherein, Genearator is a generator; Discriminator is a discriminator; AAM is an advanced attention module; Attention map is an attention map; input is an input terminal; output is an output terminal; Down-Sampling is down sampling; Up-Sampling is up sampling; Convolution is a convolution layer; Deconvolution is a deconvolution layer; BatchNorm is a batch normalization layer; leakyReLU is an activation function layer;

[0072] Figure 2 Schematic diagram of the iterative training process of the joint network in one embodiment of the present invention; wherein A is a rainy day image in the training set; B is a snowy day image in the training set; C is a sunny day image in the training set; Generator for rainy day images; For snow image generator; Generator for sunny day images; is the rainy day image discriminator; It is a snow image discriminator; is a sunny image discriminator; Fake CA is a false sunny image generated from a rainy image; Fake CB is a false sunny image generated from a snowy image; New A is a new false rainy image; New B is a new false snowy image; Fake A is a false rainy image generated from a sunny image; Fake B is a false snowy image generated from a sunny image; New CA is a false sunny image generated from a false rainy image generated from a sunny image; New CB is a false sunny image generated from a false snowy image generated from a sunny image. DETAILED DESCRIPTION

[0073] In order to make the technical solutions and advantages of the present invention more clearly described, the specific implementation methods of the present invention will be further described in detail and completely in conjunction with the accompanying drawings. The various implementation methods described below are only part of the preferred solutions of the present invention, rather than all implementation methods; the various implementation methods described below are intended to explain the present invention and cannot be understood as limitations on the present invention; the reasonable combination of technical features defined in the various implementation methods of the present invention, and all other implementation methods obtained by ordinary technicians in the field without creative work based on the implementation methods of the present invention, are within the scope of protection of the present invention.

[0074] Implementation Method 1: Combination Figure 1 and 2 This embodiment describes a training method for a general module for removing rain and snow from a single image. The specific implementation contents are as follows:

[0075] The training method comprises the following steps:

[0076] S1, obtaining a collected original image data set; the original image data set includes rainy day images, snowy day images and sunny day images;

[0077] S2, preprocessing the collected original image data set to obtain a training set;

[0078] S3, obtaining a joint network including three generative adversarial network modules; each generative adversarial network module includes a generator and a discriminator; wherein:

[0079] The first generative adversarial network module: its generator is a rainy day image generator, which is used to generate fake rainy day images; its discriminator is a rainy day image discriminator, which is used to determine whether the input image is a real rainy day image;

[0080] The second generative adversarial network module: its generator is a snow image generator, which is used to generate fake snow images; its discriminator is a snow image discriminator, which is used to determine whether the input image is a real snow image;

[0081] The third generative adversarial network module: its generator is a sunny image generator, which is used to generate fake sunny images; its discriminator is a sunny image discriminator, which is used to determine whether the input image is a real sunny image;

[0082] S4, using the real rainy day images, real snowy day images and real sunny day images in the training set to train the joint network to obtain a trained joint network;

[0083] The trained sunny image generator of the joint network is used as a general module for single-image rain and snow removal.

[0084] In this implementation, the generative adversarial network module is abbreviated as “GAN” in English.

[0085] Implementation Method 2: Combination Figure 1 and 2 This embodiment is described. This embodiment further limits the training method of the single-image rain and snow removal general module described in the first embodiment. The specific implementation content is as follows:

[0086] The step S4 is as follows:

[0087] An iterative training algorithm is used to update the model parameters of the joint network; each iterative training process includes the following steps:

[0088] S4.1, inputting the rainy day images in the training set into the sunny day image generator to obtain a false sunny day image generated from the rainy day image;

[0089] The false sunny image generated by the rainy image is input into the sunny image discriminator for identification, and the identification result is recorded in the historical identification result;

[0090] The false sunny image generated by the rainy image is input into the rainy image generator to obtain a new false rainy image, which is then input into the rainy image discriminator for identification;

[0091] S4.2, inputting the snowy images in the training set into the sunny image generator to obtain a false sunny image generated from the snowy images;

[0092] The false sunny image generated by the snowy image is input into the sunny image discriminator for identification, and the identification result is recorded in the historical identification result;

[0093] The false sunny image generated by the snowy image is input into the snowy image generator to obtain a new false snowy image, which is then input into the snowy image discriminator for identification;

[0094] S4.3, input the sunny images in the training set into the rainy image generator to obtain a false rainy image generated from the sunny image.

[0095] The false rainy day image generated by the sunny day image is input into the rainy day image discriminator for identification.

[0096] The false rainy day image generated by the sunny day image is input into the sunny day image generator to obtain the false sunny day image generated by the false rainy day image generated by the sunny day image.

[0097] The false sunny image generated by the false rainy image generated by the sunny image is input into the sunny image discriminator for identification, and the identification result is recorded into the historical identification result;

[0098] At the same time, the same sunny image in the training set is input into the snow image generator to obtain a fake snow image generated from the sunny image.

[0099] The false snowy image generated from the sunny image is input into the snowy image discriminator for identification.

[0100] The false snowy day image generated from the sunny day image is input into the sunny day image generator to obtain the false sunny day image generated from the false snowy day image generated from the sunny day image.

[0101] The false sunny image generated by the false snowy image generated by the sunny image is input into the sunny image discriminator for identification, and the identification result is recorded into the historical identification result;

[0102] S4.4, determine whether the historical identification results meet the given iterative training end conditions:

[0103] If the given iterative training end condition is met, the iterative training algorithm ends, and a trained joint network is obtained, and the sunny image generator is used as a general module for single image rain and snow removal;

[0104] Otherwise, update the model parameters of the joint network and proceed to the next iterative training process.

[0105] In this implementation manner, after the rainy day image or the snowy day image is input into the single image rain and snow removal universal module, the obtained false sunny day image is the image after the rain and snow are removed.

[0106] In this implementation, there are many images in the training set, and each iterative training process will select corresponding rainy day, snowy day and sunny day images to input into the network for training.

[0107] Implementation Method 3: Combination Figure 1 and 2 This embodiment is described. This embodiment further limits the training method of the single-image rain and snow removal general module described in the second embodiment. The specific implementation content is as follows:

[0108] The given iterative training end condition is:

[0109] In the input image identified as a true sunny image by the sunny image discriminator:

[0110] The proportion of false sunny images ranged from 45% to 55%.

[0111] In this embodiment, when the proportion of false sunny images is 45% to 55%, the corresponding proportion of sunny images in the training set is 55% to 45%. At this time, it can be considered that the sunny image discriminator cannot distinguish the false sunny images from the sunny images in the training set (i.e., real sunny images), that is, it has achieved the effect of not being able to distinguish between the real and the fake, or it can be considered that the degree of similarity between the false sunny images and the real sunny images has reached a nearly consistent level.

[0112] In this embodiment, the false sunny images accounting for 45% to 55% include: false sunny images generated from rainy images (i.e., Fake CA), false sunny images generated from snowy images (i.e., Fake CB), false sunny images generated from false rainy images generated from sunny images (i.e., New CA), and false sunny images generated from false snowy images generated from sunny images (i.e., New CB).

[0113] It should be noted that in each iterative training process of the iterative training algorithm, the discriminator and the generator exchange data, which in turn affects the optimization parameters of the network model. Therefore, although the identification results of the "rainy image discriminator" and the "snowy image discriminator" are not used in the judgment of the end conditions of the iterative training, the "rainy image discriminator" and the "snowy image discriminator" are still required to perform the corresponding identification steps in each iterative training process.

[0114] Implementation Method 4: Combination Figure 1 and 2 This embodiment is described. This embodiment further limits the training method of the single-image rain and snow removal general module described in the first embodiment. The specific implementation content is as follows:

[0115] The generator of each generative adversarial network module includes two cascaded U-Net networks to preserve the details of the input image.

[0116] In this embodiment, the generator uses two cascaded U-Net networks and takes advantage of the U-Net network's strong learning ability and simple structure, so that it can not only accurately remove more rain streaks but also retain the details of the image.

[0117] In this embodiment, the joint network includes three GANs, and the generator of each GAN includes two U-Net networks. By training the three GANs, the generators and discriminators of different GANs continuously exchange model parameters and compete with each other, and finally reach the optimal state of the network. The sunny image generator obtained at this time is the final single-image rain and snow removal universal module. The single-image rain and snow removal universal module has a simple network and a small number of parameters, which is suitable for unmanned driving scenarios with limited vehicle hardware resources. A single module can meet the general needs of rain and snow removal, and there is no need to develop multiple modules for different weather environments, which reduces the resource consumption of system deployment and improves the practicality and robustness of the system.

[0118] Further, a preferred embodiment is provided, wherein the generator of each generative adversarial network module further includes an input end, an output end, two independent convolutional layers and a Tanh function layer;

[0119] A first independent convolutional layer, two cascaded U-Net networks, a second independent convolutional layer and a Tanh function layer are sequentially connected in series between the input end and the output end;

[0120] The structures of the two U-Net networks are the same.

[0121] Further, a preferred embodiment is provided, wherein each U-Net network is composed of a first type of convolutional block group and a second type of convolutional block group connected in series;

[0122] The first type of convolutional block group is composed of 4 first type convolutional blocks connected in series;

[0123] The second type of convolutional block group is composed of 4 second type convolutional blocks connected in series;

[0124] The first type of convolution block in the first type of convolution block group and the second type of convolution block in the second type of convolution block group form a pair in a one-to-one correspondence in a centrally symmetrical manner. Specifically:

[0125] The first first-category convolution block in the first-category convolution block group and the fourth second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence;

[0126] The second first-category convolution block in the first-category convolution block group and the third second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence;

[0127] The third first-category convolution block in the first-category convolution block group and the second second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence;

[0128] The fourth first-category convolution block in the first-category convolution block group and the first second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence;

[0129] Each pair of the first type of convolutional blocks and the second type of convolutional blocks is connected by a skip connection;

[0130] Through the jump connections between four pairs of first-type convolutional blocks and second-type convolutional blocks, upsampling and downsampling operations of the input image are realized.

[0131] In this implementation, in each U-Net network, the first type of convolutional block group is used as an encoder.

[0132] In this implementation, in each U-Net network, the second type of convolutional block group is used as a decoder.

[0133] Furthermore, a preferred embodiment is provided, in which the first type of convolution block is composed of convolution, BatchNorm, leaky ReLU, convolution, BatchNorm and leaky ReLU connected in series.

[0134] Furthermore, a preferred embodiment is provided, in which the second type of convolution block is composed of a series connection of deconvolution, BatchNorm and leaky ReLU.

[0135] Furthermore, a preferred embodiment is provided, wherein the discriminator of each generative adversarial network module is composed of four first-class convolutional blocks, independent convolutional layers and sigmoid function layers connected in series.

[0136] Implementation Method 5: Combination Figure 1 and 2 This embodiment is described. This embodiment further limits the training method of the single-image rain and snow removal general module described in the fourth embodiment. The specific implementation content is as follows:

[0137] The generator of each generative adversarial network module also includes a high-level attention module; the high-level attention module is used to prevent the image generated by the generator from being overly smooth.

[0138] In this embodiment, an advanced attention module (AAM) is introduced in the generator of the generative adversarial network module. The advanced attention module is used to obtain an attention map based on the input image, extract different types of prior information, and make the network model focus on important areas blocked by rain patterns and snow patterns, thereby improving the feature extraction ability of the network model when processing rainy and snowy weather images, avoiding excessive smoothness of the generated rain and snow removal images, and improving the visual experience.

[0139] Furthermore, a preferred embodiment is provided, wherein the advanced attention module is used to obtain an attention map based on an input image; and is also used to sequentially input the attention map into the first first-class convolution block and the first second-class convolution block of the first U-Net network, and the first first-class convolution block, the first second-class convolution block and the fourth second-class convolution block of the second U-Net network.

[0140] Implementation Method 6: Combination Figure 1 and 2 This embodiment is described. This embodiment further limits the training method of the single-image rain and snow removal general module described in the fourth embodiment. The specific implementation content is as follows:

[0141] The loss function used by each generative adversarial network module is a joint loss function:

[0142] ;

[0143] in, is the dual U-Net loss function; is the structural similarity loss function; is the cycle consistency loss function.

[0144] In this implementation, the joint loss function is used to constrain the algorithm model and restore details in the image (such as edge details) in order to generate higher quality and more natural images, improve the performance of the algorithm model in image rain and snow removal tasks, and maintain the stability and reliability of the algorithm model.

[0145] Furthermore, a preferred embodiment is provided, in which a dual U-Net loss function is added between the two U-Net networks of the generator.

[0146] Further, a preferred embodiment is provided, in which the structural similarity loss function and the cycle consistency loss function are added to the 4th second-type convolutional block of the second U-Net network of the generator;

[0147] At the same time, the cycle consistency loss function is also added to the discriminator of the same generative adversarial network module, specifically, to the first first-class convolutional block of the discriminator.

[0148] In this implementation, a dual U-Net loss function is used to reduce the difference between the image generated by the generator and the actual image without rain or snow, making the generated image closer to the actual image. This is used to constrain the output of the first U-Net network of the generator, making the network model more stable.

[0149] In this embodiment, the structural similarity loss function is used to measure the distance between two images, and combined with the comprehensive consideration of brightness, contrast and structure, the structure, texture and other information in the image are maintained to prevent the generated image from being over-smoothed, making the generated image visually closer to the real image and improving the visual experience.

[0150] In this implementation, the cycle consistency loss function is used to ensure that the image can be restored to the original image as much as possible after two transformations (i.e., from the source domain to the target domain, and then from the target domain back to the source domain).

[0151] The cycle consistency loss function maintains the consistency of image content by forcing the generator to maintain the basic structure and content of the image during the generation and degeneration process, so as to avoid the generator only focusing on the target domain features and ignoring the content of the source domain, resulting in the loss or addition of too much irrelevant information when generating images, thereby improving the quality and credibility of the generated images. At the same time, it can also improve the stability of the model and avoid model crash or generation of unreasonable images.

[0152] Further, a preferred embodiment is provided:

[0153] ;

[0154] Where C is the number of channels of the input image; W is the width of the input image; H is the height of the input image; X is the input rainy or snowy image; Y is the ground truth; Indicates the first channel; Represents the first pixel in the width direction of the input image; Represents the first pixel in the height direction of the input image; is the output of the first U-Net network of the generator; for The convolution output of .

[0155] Further, a preferred embodiment is provided:

[0156] ;

[0157] in, ; is the pixel average value of image x; is the pixel average of image y; is the pixel standard deviation of image x; is the pixel standard deviation of image y; is the covariance of images x and y; and is a constant, based on the dynamic range L of the pixel value and a small constant and get, The value of is 0.01, The value of is 0.03.

[0158] Further, a preferred embodiment is provided:

[0159]

[0160] in, Represents the data distribution from source domain A The expected value obtained by sampling from Representation generator, used for translation from target domain to source domain; Representation generator, used for translation from source domain to target domain;

[0161] Represents image i in the source domain after the generator Convert to the target domain and then pass through the generator Convert back to the source domain, the result should be consistent with the original image i; Represents the image C in the target domain after the generator Convert to the source domain and then pass through the generator Translated back to the target domain, the result should be consistent with the original image C.

[0162] Implementation Method 7: Combination Figure 1 and 2 This embodiment is described. This embodiment further limits the training method of the single-image rain and snow removal general module described in the first embodiment. The specific implementation content is as follows:

[0163] The step S2 is as follows:

[0164] Each image in the original image dataset is cropped according to a pixel resolution of 480×320 to obtain cropped rainy day images, snowy day images, and sunny day images, and their collection is the training set.

[0165] Implementation Method 8: Combination Figure 1 and 2 This embodiment describes the present invention. This embodiment also provides a single-image rain and snow removal method. The specific implementation contents are as follows:

[0166] The method comprises the following steps:

[0167] Step 1, obtaining the original image to be processed;

[0168] Step 2, preprocessing the original image to be processed;

[0169] Step 3, using the preprocessed original image to be processed as the input of the single image rain and snow removal general module trained by the training method of the single image rain and snow removal general module described in the above embodiment;

[0170] The false sunny image output by the single-image rain and snow removal general module is used as the image after rain and snow removal.

[0171] Furthermore, a preferred embodiment is provided, in the step 2, the preprocessing refers to cropping the original image to be processed according to a pixel resolution of 480×320.

[0172] The technical solution provided by the present invention is further described in detail above through several specific implementation modes in order to highlight the advantages and benefits of the technical solution provided by the present invention. However, the several specific implementation modes described above are not intended to be used as limitations on the present invention. Any reasonable changes and improvements to the present invention, reasonable combinations of implementation modes and equivalent substitutions based on the spirit and principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A training method for a general module for removing rain and snow from a single image, characterized in that: The training method comprises the following steps: S1, obtaining a collected original image data set; the original image data set includes rainy day images, snowy day images and sunny day images; S2, preprocessing the collected original image data set to obtain a training set; S3, obtaining a joint network including three generative adversarial network modules; each generative adversarial network module includes a generator and a discriminator; wherein: The first generative adversarial network module: its generator is a rainy day image generator, which is used to generate fake rainy day images; its discriminator is a rainy day image discriminator, which is used to determine whether the input image is a real rainy day image; The second generative adversarial network module: its generator is a snow image generator, which is used to generate fake snow images; its discriminator is a snow image discriminator, which is used to determine whether the input image is a real snow image; The third generative adversarial network module: its generator is a sunny image generator, which is used to generate fake sunny images; its discriminator is a sunny image discriminator, which is used to determine whether the input image is a real sunny image; S4, using the real rainy day images, real snowy day images and real sunny day images in the training set to train the joint network to obtain a trained joint network; The trained sunny image generator of the joint network is used as a general module for single-image rain and snow removal. The generator of each generative adversarial network module includes two cascaded U-Net networks for retaining details of the input image; The generator of each generative adversarial network module also includes a high-level attention module; the high-level attention module is used to prevent the image generated by the generator from being too smooth; The loss function used by each generative adversarial network module is a joint loss function: ; in, is the dual U-Net loss function; is the structural similarity loss function; is the cycle consistency loss function; The generator of each generative adversarial network module also includes an input end, an output end, two independent convolutional layers, and a Tanh function layer; A first independent convolutional layer, two cascaded U-Net networks, a second independent convolutional layer and a Tanh function layer are sequentially connected in series between the input end and the output end; The structures of the two U-Net networks are the same; Each U-Net network consists of a first-class convolutional block group and a second-class convolutional block group connected in series; The first type of convolutional block group is composed of 4 first type convolutional blocks connected in series; The second type of convolutional block group is composed of 4 second type convolutional blocks connected in series; The first type of convolution block in the first type of convolution block group and the second type of convolution block in the second type of convolution block group form a pair in a one-to-one correspondence in a centrally symmetrical manner. Specifically: The first first-category convolution block in the first-category convolution block group and the fourth second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence; The second first-category convolution block in the first-category convolution block group and the third second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence; The third first-category convolution block in the first-category convolution block group and the second second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence; The fourth first-category convolution block in the first-category convolution block group and the first second-category convolution block in the second-category convolution block group form a pair in one-to-one correspondence; Each pair of the first type of convolutional blocks and the second type of convolutional blocks is connected by a skip connection; Through the jump connection between 4 pairs of first-class convolution blocks and second-class convolution blocks, upsampling and downsampling operations of the input image are realized; The discriminator of each generative adversarial network module is composed of 4 first-class convolution blocks, independent convolution layers and sigmoid function layers connected in series; The high-level attention module is used to obtain an attention map according to an input image; and is also used to sequentially input the attention map into the first first-class convolution block and the first second-class convolution block of the first U-Net network, and the first first-class convolution block, the first second-class convolution block and the fourth second-class convolution block of the second U-Net network; The dual U-Net loss function is added between the two U-Net networks of the generator; The structural similarity loss function and the cycle consistency loss function are added to the 4th second-class convolutional block of the second U-Net network of the generator; at the same time, the cycle consistency loss function is also added to the discriminator of the same generative adversarial network module, specifically, added to the 1st first-class convolutional block of the discriminator; The first type of convolutional block consists of convolution, BatchNorm, leaky ReLU, convolution, BatchNorm, and leaky ReLU in series; The second type of convolutional block consists of deconvolution, BatchNorm, and leaky ReLU in series.

2. The training method of the single-image rain and snow removal general module according to claim 1 is characterized in that: The step S4 is as follows: An iterative training algorithm is used to update the model parameters of the joint network; each iterative training process includes the following steps: S4.1, inputting the rainy day images in the training set into the sunny day image generator to obtain a false sunny day image generated from the rainy day image; The false sunny image generated by the rainy image is input into the sunny image discriminator for identification, and the identification result is recorded in the historical identification result; The false sunny image generated by the rainy image is input into the rainy image generator to obtain a new false rainy image, which is then input into the rainy image discriminator for identification; S4.2, inputting the snowy images in the training set into the sunny image generator to obtain a false sunny image generated from the snowy images; The false sunny image generated by the snowy image is input into the sunny image discriminator for identification, and the identification result is recorded in the historical identification result; The false sunny image generated by the snowy image is input into the snowy image generator to obtain a new false snowy image, which is then input into the snowy image discriminator for identification; S4.3, input the sunny images in the training set into the rainy image generator to obtain a false rainy image generated from the sunny image. The false rainy day image generated by the sunny day image is input into the rainy day image discriminator for identification. The false rainy day image generated by the sunny day image is input into the sunny day image generator to obtain the false sunny day image generated by the false rainy day image generated by the sunny day image. The false sunny image generated by the false rainy image generated by the sunny image is input into the sunny image discriminator for identification, and the identification result is recorded into the historical identification result; At the same time, the same sunny image in the training set is input into the snow image generator to obtain a fake snow image generated from the sunny image. The false snowy image generated from the sunny image is input into the snowy image discriminator for identification. The false snowy day image generated from the sunny day image is input into the sunny day image generator to obtain the false sunny day image generated from the false snowy day image generated from the sunny day image. The false sunny image generated by the false snowy image generated by the sunny image is input into the sunny image discriminator for identification, and the identification result is recorded into the historical identification result; S4.4, determine whether the historical identification results meet the given iterative training end conditions: If the given iterative training end condition is met, the iterative training algorithm ends, and a trained joint network is obtained, and the sunny image generator is used as a general module for single image rain and snow removal; Otherwise, update the model parameters of the joint network and proceed to the next iterative training process.

3. The training method of the single-image rain and snow removal general module according to claim 2 is characterized in that: The given iterative training end condition is: In the input image identified as a true sunny image by the sunny image discriminator: The proportion of false sunny images ranged from 45% to 55%.

4. The training method of the single-image rain and snow removal general module according to claim 1 is characterized in that: The step S2 is as follows: Each image in the original image dataset is cropped according to a pixel resolution of 480×320 to obtain cropped rainy day images, snowy day images, and sunny day images, and their collection is the training set.

5. A single image rain and snow removal method, characterized in that: The method comprises the following steps: Step 1, obtaining the original image to be processed; Step 2, preprocessing the original image to be processed; Step 3, using the preprocessed original image to be processed as the input of the single image rain and snow removal general module trained by the training method of the single image rain and snow removal general module according to any one of claims 1 to 4; The false sunny image output by the single-image rain and snow removal general module is used as the image after rain and snow removal.

6. A computer program product comprising a computer program or instructions, characterized in that When the computer program or instruction is executed by a processor, the steps of the training method of the general module for removing rain and snow from a single image as described in any one of claims 1 to 4 are implemented.

7. A computer storage medium, characterized in that: The storage medium stores a computer program, and when the computer program is run, the training method of the general module for removing rain and snow from a single image as claimed in any one of claims 1 to 4 is executed.

Images