A data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction

Abstract

The present invention discloses a data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction. The implementation method is to calibrate the original image data collected by industrial cameras according to image category, adjust the image size by downsampling, generate grayscale histograms of each category of image, fit the grayscale histograms by maximum likelihood estimation, obtain the probability density function f(x) of the original image pixel value distribution, generate two-dimensional random noise that approximately conforms to the probability density function f(x) according to the Box-Muller algorithm or Markov chain Monte Carlo sampling algorithm, and finally combine WGAN-GP and SRCNN to build a data enhancement model. Through layer-by-layer training, the drawback of the generative adversarial network that is difficult to directly map to high-resolution images is avoided, thereby obtaining high-resolution generated images and completing the enhancement of the original image dataset. The present invention effectively reduces the redundancy of model parameters. Compared with traditional methods, the peak signal-to-noise ratio of the generated image is improved by 13.07%, and the structural similarity is improved by 32.40%.

Description

Technical Field

[0001] The present invention belongs to the field of image processing technology, and specifically relates to a data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction. Background Art

[0002] Image processing is widely used in quality inspection and fault diagnosis in the modern parts industry. High-performance image processing algorithms require massive amounts of high-quality raw industrial data. Algorithmic models driven by only a small amount of data struggle to achieve good generalization capabilities. However, in many cases, acquiring large industrial datasets is difficult or even impossible.

[0003] To address the shortage of industrial image data, traditional methods based on image geometric transformations and random adjustments to brightness and contrast can generate new data but do not fundamentally increase the diversity of training samples. With the development of generative models, image data augmentation methods based on techniques such as generative stochastic networks (GANs), variational autoencoders (VAEs), deep Boltzmann machines (DBMs), and generative adversarial networks (GANs) have emerged. Compared to GANs and DBMs, GANs can generate samples in a single pass, rather than repeatedly computing Markov chains. Compared to variational autoencoders (VAEs), GANs are bias-free and can better fit the distribution of real samples. In the prior art, Martin Arjovsky et al., in their paper "Wasserstein GAN" (arXiv, v 2017-December, 1701:07875, 2017), replaced the Jensen-Shannon or Kullback–Leibler divergence used in traditional GANs with the Wasserstein distance as the equivalent optimized distance measure, thereby stabilizing the training process of traditional GANs. Ishaan Gulrajani et al. proposed WGAN-GP in the paper "Improved Training of Wasserstein GANs" (arXiv, v 2017-December, 1704:00028, 2017). They indirectly implemented the Lipschitz continuity condition of Wasserstein GAN by adding a gradient penalty term to the objective function.

[0004] Image superpixel reconstruction technology based on deep convolutional neural networks has emerged in recent years, enabling the reconstruction of low-resolution images into high-resolution images. Chao Dong et al. proposed the super-resolution convolutional neural network (SRCNN) in their paper "Image Super-Resolution Using Deep Convolutional Networks" (IEEE transactions on pattern analysis and machine intelligence, February 2016, 295:307), laying the foundation for subsequent deep learning-based image super-resolution reconstruction. Summary of the Invention

[0005] In order to overcome the shortcomings of traditional GANs-based image data enhancement methods, such as model parameter redundancy, training difficulties, and difficulty in establishing a mapping relationship between the network and high-resolution images, the purpose of the present invention is to provide a data enhancement method based on deep reinforcement learning generative adversarial neural networks and super-resolution reconstruction.

[0006] The present invention uses maximum likelihood estimation and random sampling algorithms to generate two-dimensional random noise that conforms to the distribution of original image data, effectively reducing model parameter redundancy. The low-resolution image generated by WGAN-GP is used as the input of SRCNN for image super-resolution reconstruction to obtain high-resolution image data, thereby avoiding the generative adversarial network from directly establishing a mapping relationship to the high-resolution image, effectively improving the quality of the generated image.

[0007] The specific technical solutions are as follows:

[0008] A method for image data enhancement based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction, the method comprising the following steps:

[0009] (1) Raw image preprocessing: First, a small amount of raw image data collected by the industrial camera is manually calibrated to obtain labels corresponding to different categories of image data, and then the raw image data is downsampled.

[0010] (2) Image data distribution fitting: The image data obtained in (1) is fitted with distribution by the maximum likelihood estimation method to obtain a probability density function f(x) that approximately obeys its data distribution.

[0011] (3) Generate random noise: Generate random noise with a probability density function of f(x) through the Box-Muller algorithm or the Markov Chain Monte Carlo (MCMC) sampling algorithm.

[0012] (4) Construction and training of data enhancement model: A data enhancement model is established by integrating Wasserstein GAN with gradient penalty (WGAN-GP) and super-resolution convolutional neural network (SRCNN).

[0013] (5) Generated data quality evaluation: The quality of the image data generated by the data augmentation model is evaluated based on the peak signal-to-noise ratio and structural similarity of the generated image.

[0014] Comparative analysis with existing traditional CGANs and WGAN-GP alone for direct mapping to high-resolution images shows that the proposed network structure has less model parameter redundancy and a more stable training process than traditional CGANs. Compared to WGAN-GP, which directly maps to high-resolution images, the proposed method uses SRCNN to establish a mapping relationship to high-resolution images, thus avoiding the difficulty of using generative adversarial networks to directly map to high-resolution images. Quality evaluation metrics for the images generated by each model were calculated to verify that the proposed method achieves optimal generated image quality.

[0015] Preferably, the specific process of step (1) is: downsampling the original image size to the required size through a single-layer convolution layer or a pooling layer to reduce data redundancy of the original image.

[0016] Preferably, the specific process of step (2) is: a. generating a grayscale histogram of the image obtained in step (1); b. preliminarily judging the distribution that the data approximately obeys based on the grayscale histogram, and writing the probability density function g that approximately obeys the distribution with estimated parameters k (x),k∈1,2,…,n; c. According to g k (x) Construct the maximum likelihood function; d. Solve the maximum likelihood function to obtain the parameter value to be estimated; e. According to the gray histogram value and g k The sum of squared errors (SSE) between (x) values evaluates the degree of fit, and the g with the smallest sum of squared errors is k (x) is the probability density function f(x) of the gray value distribution of the original image.

[0017] Preferably, in step (3), if both the probability distribution function F(x) of f(x) and the inverse function of F(x) can be calculated, the Box-Muller algorithm is used to generate random noise, otherwise the MCMC sampling algorithm is used to generate random noise.

[0018] Compared to the Box-Muller algorithm, the MCMC sampling algorithm requires a burn-in period, which results in a longer sampling time. However, the Box-Muller algorithm requires the inverse function of the probability distribution function F(x), which has certain limitations. If the inverse function of F(x) is available, using the Box-Muller algorithm directly can save time in generating random noise.

[0019] Preferably, in step (4), a noise generation module is constructed and connected to the front end of WGAN-GP, and the low-resolution image generated by WGAN-GP is used as the input of SRCNN, and a high-resolution image is obtained through SRCNN.

[0020] When mapping directly to high-resolution images through a generative adversarial network, the network struggles and the resulting image quality is poor. Therefore, a layer-by-layer training approach, where the generative adversarial network first generates low-resolution images and then reconstructs them through superpixel reconstruction, can avoid these drawbacks.

[0021] Preferably, in step (5), the image peak signal-to-noise ratio and structural similarity are used to evaluate the quality of image data generated by each model.

[0022] Peak signal-to-noise ratio is an error-sensitive image quality evaluation method that reflects the pixel value error between corresponding pixels in the image. Structural similarity weighs the brightness, contrast, and structure of the image to evaluate image quality, which is more in line with the intuitive visual perception of the human eye.

[0023] Preferably, the image data enhancement method based on generative adversarial neural network and super-resolution reconstruction is not limited to the WGAN-GP combined with SRCNN model.

[0024] The method proposed in the present invention has strong scalability and is not limited to the WGAN-GP+SRCNN model used in the experiment. For example, a boundary-balanced generative adversarial network based on an autoencoder can be combined with SRCNN or other image super-resolution reconstruction models.

[0025] Preferably, the generative adversarial neural network and super-resolution reconstruction model construction process are completed based on TensorFlow 2.0 or above, the program compilation environment is Jupyter Notebook, and the programming language is Python 3.7.

[0026] The present invention has the beneficial effects:

[0027] (1) The present invention proposes an image data enhancement method based on a generative adversarial neural network and super-resolution reconstruction, which can establish a mapping relationship to a high-resolution image through layer-by-layer training and generate corresponding high-resolution image data according to given label information.

[0028] (2) Two-dimensional random noise that conforms to the distribution of the original data is generated through maximum likelihood estimation and random sampling algorithms, which overcomes the randomness of the noise input of traditional methods and reduces the redundancy of model parameters.

[0029] (3) Compared with the traditional method, the peak signal-to-noise ratio of the image generated by the method proposed in this invention is improved by 13.07%, and the structural similarity is improved by 32.40%. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 This is the flow chart of the Box-Muller algorithm;

[0031] Figure 2 This is the flow chart of the MCMC sampling algorithm;

[0032] Figure 3 This is the overall framework diagram of the WGAN-GP+SRCNN model;

[0033] Figure 4 The side image data of the roller bearing and its distribution fitting results;

[0034] Figure 5 The image data of the chamfered surface of the roller bearing and its distribution fitting results;

[0035] Figure 6 The image data of the roller bearing end face and its distribution fitting results;

[0036] Figure 7 is the random noise that matches the distribution of the side image;

[0037] Figure 8 is the random noise that matches the distribution of the chamfered surface image;

[0038] Figure 9 is the random noise that matches the distribution of the end face image;

[0039] Figure 10 Generate image data for the WGAN-GP+SRCNN model;

[0040] Figure 11 Generate image data for traditional CGAN models;

[0041] Figure 12 Generate image data for the WGAN-GP model;

[0042] Figure 13 is the real image data. DETAILED DESCRIPTION

[0043] The present invention will be further described below with reference to specific embodiments.

[0044] A method for image data enhancement based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction, the method comprising the following steps:

[0045] 1) Raw image preprocessing: First, a small amount of raw image data collected by industrial cameras is manually calibrated to obtain labels corresponding to image data of different categories. Then, the raw image data is downsampled through a single-layer volume or pooling layer method to reduce the image size to the required size.

[0046] 2) Image data distribution fitting: a. Generate the grayscale histogram of the image obtained in step (1); b. Preliminarily determine the distribution that the data approximately obeys based on the grayscale histogram, and write the probability density function g that approximately obeys the distribution with estimated parameters k (x),k∈1,2,…,n; c. According to g k (x) Construct the maximum likelihood function; d. Solve the maximum likelihood function to obtain the parameter value to be estimated; e. According to the gray histogram value and g k The sum of squared errors between (x) values evaluates the degree of fit, and the g k (x) is the probability density function f(x) of the original image data distribution.

[0047] 3) Generate random noise: a. Determine whether the probability distribution function F(x) of f(x) obtained in step (2) and the inverse function of F(x) can be calculated; b. If they can be calculated, use the Box-Muller algorithm to generate random noise; otherwise, use the MCMC sampling algorithm to generate random noise. The specific steps are as follows: Figure 1 and Figure 2 shown.

[0048] 4) Construction and training of data augmentation model: a. Build WGAN-GP model, and use Adam or RMSProp optimizer as optimizer; b. Construct the random noise generation method in step (3) into a noise generation module and connect it to the front end of WGAN-GP to generate random noise that conforms to the distribution of original data for the model; c. Build SRCNN model, and use Adam optimizer as optimizer. Input the low-resolution image generated by WGAN-GP into SRCNN for image super-resolution reconstruction to obtain high-resolution image data. The overall structure of the model is as follows: Figure 3 shown.

[0049] 5) Generated data quality evaluation: Calculate the peak signal-to-noise ratio and structural similarity between the generated image data and the original image data. The larger the value, the higher the quality of the generated image.

[0050] The generative adversarial neural network and super-resolution reconstruction model construction process are completed based on TensorFlow 2.0 or above, the program compilation environment is Jupyter Notebook, and the programming language is Python 3.7.

[0051] Example 1

[0052] 1) Original image preprocessing: The roller bearing side, chamfered surface, and end face images captured by the CCD industrial camera are manually calibrated to 0, 1, and 2, respectively. The original image size is (1920 pixels, 1200 pixels, 1 pixel). First, the image size is cropped to (1920 pixels, 1184 pixels, 1 pixel). Then, a single convolution layer with a convolution kernel of 1 is used to downsample the cropped image into a series of images of size (480 pixels, 296 pixels, 1 pixel). k , k=1,2,...,n and image i of size (240 pixels, 148 pixels, 1 pixel) k ′,k=1,2,...,n;

[0053] 2) Image data distribution fitting: a. In i k Randomly select one of the side, chamfered and end bearing images of ′ and get i k ′0、i k ′1 and i k ′2, generate i k ′0、i k ′1 and i k Grayscale histogram of ′2; b. Preliminary judgment based on the grayscale histogram i k ′0、i k ′1 and i k ′2 pixel values approximately obey a series of distributions, and write out their probability density functions with estimated parameters and c. Construct the maximum likelihood function, solve the maximum likelihood function, and obtain the value of the parameter to be estimated. d. Calculate and The corresponding sum of squared errors is taken as and The function with the smallest sum of squared errors is the probability density function f0(x), f1(x), f2(x) of the side, chamfered and end images. The original image data and its fitting results are shown in Figures 4 to 6 As shown, the probability density functions f0(x), f1(x), and f2(x) shown in the following equations (1), (2), and (3) can be fitted from the original image data on the left.

[0054]

[0055]

[0056]

[0057] 3) Generate random noise: The process of solving the inverse function of the cumulative distribution function of f0(x) and f1(x) is extremely complicated. The MCMC sampling algorithm is directly used to generate random noise, and the combustion period is set to 10000. f2(x) is a Cauchy distribution, and the inverse function of the cumulative distribution function can be solved. The Box-Muller algorithm is directly used to generate random noise. The random noise dimension is (60,37). The generated random noise is as follows Figure 7-9 shown.

[0058] 4) Construction and training of data augmentation model: a. Build WGAN-GP model, which includes generative model and discriminative model, with batch size set to 3. The input dimension of generative model is (3, 60, 37, 1), and the output dimension is (3, 240, 148, 1); the input dimension of discriminative model is 3, 240, 148, 1), and the output is 0 or 1. b. According to the method in step (3), a noise generation module is constructed to generate random noise for WGAN-GP model according to the label information. c. After WGAN-GP is trained in 1500 batches, the generative module is taken out separately to generate a series of low-resolution images with a size of (240, 148). d. Build SRCNN model with an input dimension of (3, 240, 148, 1) and an output dimension of (3, 480, 296, 1). The low-resolution image generated by WGAN-GP is used as the input of SRCNN model. In step (1), i k These are real samples. After 500 batches, the SRCNN model training is completed.

[0059] 5) Evaluation of generated data quality: The noise generation module generates three random noises that conform to the distribution according to the label information. The noise is input into the trained WGAN-GP generative model to obtain a low-resolution image G(z). G(z) is then input into the trained SRCNN model to obtain a high-resolution image SR(G). The average peak signal-to-noise ratio of SR(G) is calculated to be 65.467, and the average structural similarity is 0.822. The two evaluation indicators show good consistency, indicating that the model has high image generation quality. The model generates image data such as Figure 10 shown.

[0060] Example 2

[0061] A method for enhancing image data based on deep reinforcement learning generative adversarial neural networks and super-resolution reconstruction is presented. The advantages and effective benefits of the network architecture proposed in the present invention compared to the existing network architecture are analyzed in combination with the embodiments. In order to eliminate the interference of other factors, the data used in Example 2 is the same as that in Example 1. Steps (1) to (3) in Example 2 are consistent with those in Example 1. In step (4), a mapping relationship is directly established to the high-resolution image through the CGAN model and the WGAN-GP model. The batch size is set to 3. The input dimension of the generative model of CGAN and WGAN-GP is (3, 60, 37, 1), and the output dimension is (3, 480, 296, 1); the input dimension of the discriminant model is (3, 480, 296, 1), and the output is 0 or 1. After each model is trained for 1500 batches, the noise generation module generates three random noises that conform to the distribution according to the label information. The noise is input into the trained generative models of CGAN and WGAN-GP, respectively, to obtain new image data G cgan (z) and G w-gp (z), calculate G cgan The average peak signal-to-noise ratio of (z) is 51.610, the average structural similarity is 0.240, and G w-gp The average peak signal-to-noise ratio of (z) is 57.900, and the average structural similarity is 0.621. Example 2 shows that the average peak signal-to-noise ratio of the image generated by the method proposed in the present invention is 26.85% higher than that of the traditional CGAN and 13.07% higher than that of WGAN-GP; the average structural similarity is 241.94% higher than that of the traditional CGAN and 32.40% higher than that of WGAN-GP. Figure 11-13 As shown in the figure, the traditional CGAN model is extremely difficult to train, the model cannot learn useful information, the convergence effect is very poor, and the generated images have no reference value; WGAN-GP can converge stably, but it is difficult to directly establish a mapping relationship model to high-resolution images, and the generated image quality is average. Compared with the comparison model, the model proposed in this invention has the best generated image quality.

Claims

1. A data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction, characterized by The steps include: 1) Raw image preprocessing: First, a small amount of raw image data collected by industrial cameras is manually calibrated to obtain labels corresponding to different categories of image data, and then the raw image data is downsampled; 2) Image data distribution fitting: The image data obtained in step 1) is subjected to distribution fitting by the maximum likelihood estimation method to obtain a probability density function f(x) that approximately obeys its data distribution; 3) Generate random noise: Generate random noise with a probability density function of f(x) using the Box-Muller algorithm or the Markov chain Monte Carlo sampling algorithm; 4) Building and training a data augmentation model: A data augmentation model is constructed by integrating WGAN-GP with a gradient penalty term and a super-resolution convolutional neural network. 5) Generated data quality evaluation: The quality of the image data generated by the data augmentation model is evaluated based on the peak signal-to-noise ratio and structural similarity of the image; The specific process of step 4) is: a. Build the WGAN-GP model and construct the noise generation method in step 3) into a noise generation module externally connected to the front end of the WGAN-GP model to generate noise that conforms to the original data distribution for the model; b. Train the WGAN-GP model. After training, use the WGAN-GP generation module alone to generate a series of low-resolution images G n (x),n∈1,2,…,n; c. Build a super-resolution convolutional neural network model and transform G n (x) Input SRCNN for super-resolution reconstruction to obtain high-resolution image data.

2. A data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction as claimed in claim 1, characterized in that The specific process of step 1) is: downsampling the original image size to the required size through a single convolutional layer or pooling layer method to reduce data redundancy of the original image.

3. The data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction according to claim 1, characterized in that The specific process of step 2) is: a. Generate a grayscale histogram of the image obtained in step 1); b. Preliminarily judge the distribution that the data approximately obeys based on the grayscale histogram, and write the probability density function g of the estimated parameters of the approximate distribution. n (x),n∈1,2,…,n; c. According to g n (x) constructing a maximum likelihood function; d. Solve the maximum likelihood function to obtain the value of the parameter to be estimated; e. According to the gray histogram value and g n The sum of squared errors between (x) values evaluates the degree of fit, and the g n (x) is the probability density function f(x) of the original image data distribution.

4. The data enhancement method based on deep reinforcement learning generative adversarial neural network and super-resolution reconstruction according to claim 1, characterized in that In step 3), if the probability distribution function F(x) of f(x) and the inverse function of F(x) can both be calculated, the Box-Muller algorithm is used to generate random noise; otherwise, the Markov chain Monte Carlo sampling algorithm is used to generate random noise.

Images