A defect detection method and system based on a two-stage generative adversarial network

Abstract

The application discloses a kind of defect detection method and system based on two-stage generative adversarial network.There are following steps: obtaining defect sample and good sample dataset, marking defect sample to make it become defect mask sample;Defect mask sample is used to train BayesGAN defect mask generation network;Non-defect texture extraction network is used to obtain the feature map of good sample and calculate its style conditional gram matrix;The style conditional gram matrix of good sample, defect sample and defect mask sample are used to train mask-driven defect generation network;The output of the generator of the two trained generation networks is used to train the defect detection network based on segmentation decision;The sample to be detected is input into the defect detection network, and the defect detection result is output.The application proposes two-stage generative adversarial network to generate defect sample, provides a large amount of training data containing annotation for defect detection model, without manual post-labeling.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of deep learning and computer vision, and more particularly relates to a defect detection method and system based on a two-stage generative adversarial network. BACKGROUND

[0002] In industry, transportation, national defense and daily life, various equipment serves as an indispensable important component, and its downstream industry almost covers all basic fields of the national economy. The core components of these equipment, while accompanying the continuous development of the related industry, also face growing demand. These core components bear key functions in the production process and are widely used in many aspects such as national defense and military, aviation field and civil facilities. However, the production process is inevitably disturbed by uncontrollable external factors, resulting in various defects of the components. For example, scratch defects may occur on the side of the component due to slight scratches with the machine tool during the manufacturing process, and usually present as strip-shaped and dark color. Such defects can damage the insulation between materials, cause short circuits, and further damage the product. Therefore, for the core components of various equipment, it is particularly important to accurately identify abnormal samples to avoid possible serious consequences. Although deep learning has made significant progress in defect detection, it still faces many challenges, such as lack of labeled sample data, limited feature learning ability, insufficient detection accuracy, and weak model generalization ability.

[0003] The prior art patent with publication number CN112365478A proposes a motor commutator surface defect detection model based on semantic segmentation, including the following steps: divide the data set into three parts, optionally select two parts as the training set, and the other part as the test set, repeat three times, and perform three-fold cross-validation; input the training data into the segmentation network, train the segmentation network using pixel segmentation labels, and output the semantic segmentation result; input the result of the segmentation network into the classification network, first freeze the weight parameters of the trained segmentation network, then train the classification network according to the prepared labels, and the sample label of the defective sample is 1 and the sample label of the non-defective sample is 0; input the test data into the trained segmentation network and classification network, and output a probability value between 0 and 1, the greater the probability, the higher the possibility of defects. Select 0.9 as the threshold for determining positive and negative samples, and greater than or equal to 0.9 is a defective sample (positive sample) and less than 0.9 is a non-defective sample (negative sample). This method requires a large amount of labeled sample data when training the defect detection model, but in reality there is often a lack of labeled sample data, which ultimately leads to weak detection ability of the model for defects. SUMMARY

[0004] The application provides a defect detection method and system based on a two-stage generative adversarial network to overcome the problem in the prior art that a large number of labeled sample data are needed to improve the defect detection capability of a model.

[0005] The primary object of the application is to solve the above technical problem, and the technical solution of the application is as follows:

[0006] The application provides a defect detection method based on a two-stage generative adversarial network, which comprises the following steps:

[0007] A data set is obtained, the data set comprising defect samples and good samples, the defect samples are labeled and marked with a defect mask to become defect mask samples;

[0008] A BayesGAN defect mask generation network is trained by using the defect mask samples, the network comprising a first generator and a first discriminator trained in cooperation;

[0009] A non-defect texture extraction network is used to obtain a feature map of the good samples and calculate a style condition gram matrix thereof;

[0010] A mask-driven defect generation network is trained by using the style condition gram matrix of the good samples, the defect samples and the defect mask samples, the network comprising a second generator and a second discriminator trained in cooperation;

[0011] The trained first generator and second generator are used to generate defect samples with labeled information, the defects are supplemented to the data set, and a defect detection network based on segmentation decision is trained by using the supplemented data set;

[0012] The sample to be detected is input into the defect detection network, and a defect detection result is output.

[0013] Further, the step of training the BayesGAN defect mask generation network is as follows:

[0014] The weight parameters of the first generator are fixed, noise sampled in accordance with a Gaussian distribution is randomly sampled, the noise is converted into a geometric mask image by the first generator G, the geometric mask image G(z) and the defect mask sample are input into the first discriminator, the first discriminator is trained by using a loss function and the weight parameters of the first discriminator are updated, and the loss function D loss is defined as:

[0015] D loss =-log[D(x)]+log[D(G(z))]

[0016] wherein x is a defect mask sample image, z is noise sampled in accordance with a Gaussian distribution, and D and G represent the first discriminator and the first generator, respectively;

[0017] The weight parameters of the first discriminator are fixed, noise sampled from a Gaussian distribution is converted into a geometric mask by the first generator G, the generated geometric mask G(z) is input into the first discriminator, the first generator is trained by using a loss function and the weight parameters thereof are updated, and the loss function G loss is defined as:

[0018] G loss =-log[D(G(z))]

[0019] wherein z is noise sampled from a Gaussian distribution, and D and G represent the first discriminator and the first generator, respectively.

[0020] The above steps are repeated until the BayesGAN defect mask generation network model converges.

[0021] Further, the steps of obtaining the feature map of the good sample and calculating the style conditional Gram matrix thereof are:

[0022] The good sample picture is input into a preset non-defect texture extraction network, and a feature map is output.

[0023] The output feature map is used to calculate the style conditional Gram matrix, and the formula is as follows:

[0024]

[0025] wherein x is the feature map, that is, the matrix [C, WH]. W, H, and C are the width, height, and channel number of the feature map, respectively.

[0026] Further, the preset non-defect texture extraction network comprises the first five blocks of a VGG-16 network.

[0027] Further, the steps of training the mask-driven defect generation network are:

[0028] The weight parameters of the second generator are fixed, a defect mask sample picture x is input into the second generator G t , the style conditional Gram matrix GM(y) of the good sample is input into the second generator G t as a condition, a defect sample G t (x, GM(y)) is generated, the defect sample picture I and the style conditional Gram matrix GM(y) of the good sample are input into the second discriminator, a loss function is calculated and the weight parameters of the second discriminator are updated, and the loss function is defined as:

[0029]

[0030] wherein Dt , G t respectively represent the second discriminator and the second generator;

[0031] fixing the weight parameters of the second discriminator, taking the defect mask sample picture x as the input of the second generator G t , taking the style condition gram matrix GM(y) of the good sample as the condition input of the second generator G t , generating the defect sample G t (x, GM(y)) and inputting the style condition gram matrix GM(y) of the good sample into the second discriminator, calculating the loss function and updating the weight parameters of the second generator, and the loss function is defined as:

[0032]

[0033] repeating the above steps until the mask-driven defect generation network model converges.

[0034] Further, the trained first generator and second generator can generate any number of defect samples with annotation information.

[0035] Further, the step of training the defect detection network based on segmentation decision is:

[0036] randomly sampling a fixed batch of noise subject to Gaussian distribution, generating a defect geometric mask through the trained first generator, and then generating a defect picture through the trained second generator, inputting the generated defect geometric mask, the generated defect picture and the original picture into a preset shared network, and outputting a high-level feature map;

[0037] inputting the extracted feature map into a preset segmentation network to output a defect segmentation heat map, and inputting the high-level feature map into a preset classification network to output a feature map;

[0038] fusing the defect segmentation heat map of the segmentation network after global average pooling with the output feature map of the classification network, outputting a probability representing the existence of defects in the whole picture through a fully connected layer and a preset function, calculating the loss function and updating the network parameters, and the loss function L is defined as:

[0039]

[0040] wherein C is the number of categories, x is a detection sample, p(x i ) is the probability of the detection sample being i category, and q(x i ) is the probability of the detection sample being predicted as i category.

[0041] Further, the preset classification network is a classification network embedded with a residual attention module; the preset shared network comprises three convolutional layers and three pooling layers, and each convolutional layer comprises a BN layer and a RELU layer; and the preset segmentation network comprises four convolutional layers.

[0042] Further, the preset function outputting a probability of presence of defects in the whole picture is a sigmoid function.

[0043] The second aspect of the present application provides a defect detection system based on a two-stage generative adversarial network, which comprises a memory and a processor, wherein the memory comprises a defect detection method based on a two-stage generative adversarial network, and the defect detection method based on the two-stage generative adversarial network is implemented when the processor is executed.

[0044] A data set is obtained, which comprises defect samples and good samples, the defect samples are labeled and marked with defect masks to become defect mask samples;

[0045] The defect mask samples are used to train a BayesGAN defect mask generation network, and the network comprises a first generator and a first discriminator trained cooperatively;

[0046] A preset non-defect texture extraction network is used to obtain a feature map of the good samples and calculate a style condition gram matrix thereof;

[0047] The style condition gram matrix of the good samples, the defect samples and the defect mask samples are used to train a mask-driven defect generation network, and the network comprises a second generator and a second discriminator trained cooperatively;

[0048] The trained first generator and the second generator are used to generate defect samples with labeled information, the defects are supplemented to the data set, and a defect detection network based on segmentation decision is trained by using the supplemented data set;

[0049] The sample to be detected is input into the defect detection network, and a defect detection result is output.

[0050] Compared with the prior art, the technical scheme of the present application has the following beneficial effects:

[0051] The defect detection model is obtained by using the defect detection method based on the two-stage generative adversarial network, the two-stage generative adversarial network is proposed to generate defect samples, a large amount of training data containing labels are provided for the defect detection model, manual post-labeling is not required, the segmentation network and the attention mechanism are fused into the defect detection classification network, and thus the detection capability of the defect detection network is improved. BRIEF DESCRIPTION OF DRAWINGS

[0052] In order to make the purposes, technical solutions of the present application clearer, the present application provides the following figures and makes the following description:

[0053] Figure 1 The method flowchart provided for the embodiments of the present application;

[0054] Figure 2 The defect mask sample diagram provided for the embodiments of the present application;

[0055] Figure 3 The geometric mask diagram provided for the embodiments of the present application;

[0056] Figure 4 The first generative adversarial network data framework diagram of the defect detection method based on a two-stage generative adversarial network provided for the embodiments of the present application;

[0057] Figure 5 The process diagram of generating defect samples provided for the embodiments of the present application;

[0058] Figure 6 The second generative adversarial network data framework diagram of the defect detection method based on a two-stage generative adversarial network provided for the embodiments of the present application;

[0059] Figure 7 The defect detection framework diagram based on segmentation decision of the defect detection method based on a two-stage generative adversarial network provided for the embodiments of the present application; DETAILED DESCRIPTION

[0060] In order to make the purposes, technical solutions of the present application clearer, the present application provides the following figures and makes the following description:

[0061] In the following description, many specific details are set forth in order to provide a thorough understanding of the present application, however, the present application can also be implemented in other ways different from those described herein, therefore, the protection scope of the present application is not limited by the specific embodiments disclosed below.

[0062] The present application provides a defect detection method based on a two-stage generative adversarial network, as shown in Figure 1 The defect detection method flowchart based on a two-stage generative adversarial network is shown in the following specific steps:

[0063] S1: Obtain a data set, the data set includes defect samples and good samples, in a specific embodiment, 15 sample pictures with different defects and 15 good sample pictures are obtained on a production line, the defect samples are labeled and marked with a defect mask, so as to become defect mask samples, as shown in Figure 2 ​

[0064] S2: training a BayesGAN defect mask generation network with the defect mask samples, the network comprising a first generator and a first discriminator trained cooperatively;

[0065] The specific process is as follows:

[0066] The weight parameters of the first generator are fixed, and noise sampled from a Gaussian distribution z ~ p(z) is randomly sampled. In a specific embodiment, a 100-dimensional vector sampled from a normal distribution is converted into a 800x200 geometric mask image G(z) by the first generator G, as shown in Figure 3 The geometric mask image G(z) and the defect mask sample x are input into the first discriminator, as shown in Figure 4 The first discriminator is trained and its weight parameters are updated by using a loss function D loss defined as:

[0067] D loss = -log[D(x)] + log[D(G(z))]

[0068] wherein z is the randomly sampled Gaussian noise, and D and G represent the first discriminator and the first generator respectively;

[0069] The weight parameters of the first discriminator are fixed, and noise sampled from a Gaussian distribution z ~ p(z) is randomly sampled. The noise is converted into a geometric mask image by the first generator G, and the generated geometric mask image G(z) is input into the first discriminator. The first generator is trained and its weight parameters are updated by using a loss function G loss defined as:

[0070] G loss = -log[D(G(z))]

[0071] wherein z is the randomly sampled Gaussian noise, and D and G represent the first discriminator and the first generator respectively;

[0072] The above steps are repeated until the BayesGAN defect mask generation network model converges.

[0073] S3: obtaining a feature map of a good sample and calculating a style condition gram matrix by using a preset non-defect texture extraction network;

[0074] The specific process is as follows:

[0075] After the good sample picture is resized to 256x256, it is input into the preset non-defect texture extraction network, and a 16x16 size feature map is output;

[0076] The output feature map is used to calculate the style condition gram matrix, and the formula is as follows:

[0077]

[0078] wherein x is a feature map, i.e. matrix [C, WH]. W, H, C are width, height, channel number of the feature map respectively.

[0079] It should be noted that the preset non-defect texture extraction network includes the first five blocks of a VGG-16 network.

[0080] S4: training a mask-driven defect generation network using the style-conditioned Gram matrix of the good sample, the defect sample and the defect mask sample, the network including a second generator and a second discriminator trained cooperatively;

[0081] The specific process is as follows:

[0082] Fixing the weight parameters of the second generator, taking the defect mask sample picture x as the input of the second generator G t , taking the style-conditioned Gram matrix GM(y) of the good sample as the conditional input of the second generator G t , and generating a defect sample G t (x, GM(y)), as shown in Figure 5 . The generated defect sample G t (x, GM(y)), the defect sample picture l and the style-conditioned Gram matrix GM(y) of the good sample are input into the second discriminator, as shown in Figure 6 , a loss function is calculated and the weight parameters of the second discriminator are updated, and the loss function is defined as:

[0083]

[0084] wherein D t , G t represent the second discriminator and the second generator respectively;

[0085] Fixing the weight parameters of the second discriminator, taking the defect mask sample picture x as the input of the second generator G t , taking the style-conditioned Gram matrix GM(y) of the good sample as the conditional input of the second generator G t , and generating a defect sample G t (x, GM(y)) and the style-conditioned Gram matrix GM(y) of the good sample are input into the second discriminator, a loss function is calculated and the weight parameters of the second generator are updated, and the loss function is defined as:

[0086]

[0087] Repeat the above steps until the mask-driven defect generation network model converges.

[0088] S5: generating defect samples with annotation information using the trained first generator and second generator, supplementing the defects to the dataset, training a defect detection network based on segmentation decision using the supplemented dataset;

[0089] It should be noted that the trained first generator and second generator can generate any number of defect samples with annotation information.

[0090] The specific process is as follows:

[0091] Randomly sample a fixed batch of noise subject to Gaussian distribution, generate defect geometry mask through the trained first generator, and then generate defect picture through the trained second generator. The generated defect geometry mask, generated defect picture and original picture are input into a preset shared network. In one specific embodiment, the shared network is composed of 3 convolutional layers and pooling layers, each convolutional layer includes a BN layer and a RELU layer, and outputs a 128-channel high-level feature map;

[0092] The extracted feature map is input into a preset segmentation network. In one specific embodiment, the segmentation network is composed of 4 convolutional layers, the convolutional kernel size of the first two layers is 3x3, the convolutional kernel size of the third layer is 5x5, and the convolutional kernel size of the last layer is 1x1. Finally, the output feature map is pixel-level predicted by a sigmoid function, and a 1-channel defect segmentation heat map is output. The high-level feature map is input into a classification network embedded with a residual attention module RSAM. In one specific embodiment, after passing through an attention layer and 3 convolutional layers, a 256-channel feature map is output, as shown in Figure 7 ;

[0093] The defect segmentation heat map of the segmentation network is globally averaged and pooled, and then fused with the output feature map of the classification network. After passing through a fully connected layer and a sigmoid function, a probability representing the presence of defects in the entire picture is output. The loss function is calculated and the network parameters are updated. The loss function L is defined as:

[0094]

[0095] Where C is the number of classes, x is the detection sample, p(x i ) is the probability that the detection sample is of class i, and q(x i ) is the probability that the detection sample is predicted to be of class i.

[0096] S6: inputting the sample to be detected into the defect detection network to output a defect detection result.

[0097] The following experimental model code is implemented based on Pytorch, and is trained on a computer with two NVIDIA GeForce RTX2070Super graphics cards, with a data batch size of 32, using the stochastic gradient descent algorithm and a learning rate of 0.005. The experimental data is shown in Table 1.

[0098] Table 1

[0099]

[0100]

[0101] Among them, the U-Net and FCN segmentation decision methods are existing models. As can be seen from Table 1, the scheme of the present application has higher detection accuracy when detecting multiple defect types (such as scratches, dents, and stains) compared to existing methods. Taking scratch detection as an example, the average precision (AP) of the present application reaches 94.14%, which is significantly better than the 92.56% of U-Net and the 91.25% of FCN. Not only does it improve detection accuracy, but it also performs well in reducing false positives (FP) and false negatives (FN). In the detection of 45 samples of scratches, the FP+FN of the present application is 9, lower than the 13 of U-Net and the 14 of FCN. In the detection of stains, the FP+FN of the present application is only 5, which is also better than the performance of U-Net and FCN. In addition, the present application performs stably under different sample sizes, and even when the training sample is reduced to 30, the AP of scratch detection still reaches 93.19%, and the AP of dent detection is 95.83%, showing its superior performance when the sample size changes.

[0102] The second aspect of the present application provides a defect detection system based on a two-stage generative adversarial network, which comprises a memory and a processor, wherein the memory comprises a defect detection method based on a two-stage generative adversarial network, and the defect detection method based on the two-stage generative adversarial network is implemented when the processor is executed.

[0103] Obtain a data set, the data set comprising defect samples and good samples, label the defect samples and mark them with a defect mask to make them defect mask samples;

[0104] Train the BayesGAN defect mask generation network using the defect mask samples, wherein the network comprises a first generator and a first discriminator trained cooperatively;

[0105] Obtain a feature map of the good sample and calculate its style condition gram matrix using a pre-set non-defect texture extraction network;

[0106] A defect generation network driven by a mask is trained by using a style condition gram matrix of a good sample, a defect sample and a defect mask sample, and the network comprises a second generator and a second discriminator trained cooperatively;

[0107] Defect samples with label information are generated by using the trained first generator and the second generator, the defects are supplemented into a data set, and a defect detection network based on segmentation decision is trained by using the supplemented data set;

[0108] The sample to be detected is input into the defect detection network, and a defect detection result is output.

[0109] Obviously, the above embodiments of the present application are merely exemplary and are not intended to limit the embodiments of the present application. Based on the above description, other different forms of changes or variations can be made by those skilled in the art. Here, all the embodiments are not required to be exhausted. Any modification, equivalent replacement and improvement made within the spirit and principle of the present application shall be included in the protection scope of the claims of the present application.

Claims

1. A defect detection method based on a two-stage generative adversarial network, characterized in that, The method comprises the following steps: obtaining a data set comprising defect samples and good samples, labeling the defect samples and marking them with defect masks to make them defect mask samples; training a BayesGAN defect mask generation network using the defect mask samples, wherein the network comprises a first generator and a first discriminator trained in cooperation; obtaining a feature map of the good samples and calculating a style-conditioned Gram matrix thereof using a preset non-defect texture extraction network; training a mask-driven defect generation network using the style-conditioned Gram matrix of the good samples, the defect samples and the defect mask samples, wherein the network comprises a second generator and a second discriminator trained in cooperation; generating defect samples with labeled information using the trained first generator and the second generator, supplementing the defects to the data set, and training a defect detection network based on segmentation decision using the supplemented data set; inputting a sample to be detected into the defect detection network to output a defect detection result; the step of training the defect detection network based on segmentation decision comprises: randomly sampling a fixed batch of noise subject to a Gaussian distribution, generating a defect geometry mask through the trained first generator, generating a defect picture through the trained second generator, inputting the generated defect geometry mask, the generated defect picture and an original picture into a preset shared network to output a high-level feature map; inputting the extracted feature map into a preset segmentation network to output a defect segmentation heat map, and inputting the high-level feature map into a preset classification network to output a feature map; fusing the defect segmentation heat map of the segmentation network and the output feature map of the classification network after global average pooling, outputting a probability representing the presence of defects in the whole picture through a fully connected layer and a preset function, calculating a loss function and updating network parameters, and defining the loss function L as: where C is the number of classes, x is the test sample, p(x i ) is the probability that the test sample is of class i, and q(x i ) is the probability that the test sample is predicted to be of class i. 2.The defect detection method based on the two-stage generative adversarial network according to claim 1, characterized in that, the step of training the BayesGAN defect mask generation network comprises: The weight parameters of the first generator are fixed, noise sampled randomly in accordance with a Gaussian distribution is converted into a geometric mask by the first generator G, the geometric mask G(z) and the defect mask sample x are input into the first discriminator, the first discriminator is trained by using a loss function and its weight parameters are updated, and the loss function D loss is defined as: D loss = -log[D(x)] + log[D(G(z))] wherein x is a defect mask sample picture, z is randomly sampled noise subject to a Gaussian distribution, and D and G represent the first discriminator and the first generator, respectively; The weight parameters of the first discriminator are fixed, noise obeying a Gaussian distribution is randomly sampled, the noise is converted into a geometric mask image by the first generator G, the generated geometric mask image G(z) is input into the first discriminator, the first generator is trained by using a loss function and the weight parameters thereof are updated, and the loss function G loss is defined as: G loss = -log[D(G(z))] wherein z is randomly sampled noise subject to a Gaussian distribution, and D and G represent the first discriminator and the first generator, respectively; repeating the above steps until the BayesGAN defect mask generation network model converges. 3.The defect detection method based on two-stage generative adversarial network according to claim 1, characterized in that, the step of obtaining a feature map of the good samples and calculating a style-conditioned Gram matrix thereof comprises: inputting a good sample picture into a preset non-defect texture extraction network to output a feature map; calculating a style-conditioned Gram matrix using the output feature map, and the formula is as follows: wherein x is a feature map, i.e. matrix [C, WH], W, H, C are width, height, channel number of the feature map respectively.

4. The defect detection method based on the two-stage generative adversarial network according to claim 1 or 3, characterized in that, the preset non-defect texture extraction network comprises the first five blocks of a VGG-16 network. 5.The defect detection method based on two-stage generative adversarial network according to claim 1, characterized in that, the step of training the mask-driven defect generation network comprises: Fix the weight parameters of the second generator, input the defect mask sample picture x as the input of the second generator G t , input the style condition gram matrix GM(y) of the good sample as the condition input of the second generator G t , generate the defect sample G t (x, GM(y)), input the defect sample picture l and the style condition gram matrix GM(y) of the good sample into the second discriminator, calculate the loss function and update the weight parameters of the second discriminator, and the loss function is defined as: wherein D t , G t represent the second discriminator and the second generator, respectively; The weight parameters of the second discriminator are fixed, the defect mask sample picture x is taken as the input of the second generator G t , the style condition gram matrix GM(y) of the good sample is taken as the condition input of the second generator G t , and the defect sample G t (x, GM(y)) is generated. The style condition gram matrix GM(y) of the good sample is input into the second discriminator, a loss function is calculated, and the weight parameters of the second generator are updated. The loss function is defined as: repeating the above steps until the mask-driven defect generation network model converges. 6.The defect detection method based on two-stage generative adversarial network of claim 1, wherein, The trained first generator and the second generator can generate any number of defect samples with labeled information.

7. The defect detection method based on two-stage generative adversarial network according to claim 1, characterized in that, The preset classification network is a classification network embedded with a residual attention module; the preset shared network comprises three convolutional layers and three pooling layers, and each convolutional layer comprises a BN layer and a RE LU layer; and the preset segmentation network comprises four convolutional layers. 8.The defect detection method based on two-stage generative adversarial network of claim 1, wherein, The preset function representing the probability of the presence of defects in the whole picture is a sigmoid function. 9.A defect detection system based on a two-stage generative adversarial network, characterized in that, The system comprises a memory and a processor, the memory comprising a defect detection method based on a two-stage generative adversarial network program, and the defect detection method based on the two-stage generative adversarial network program is executed by the processor to realize the steps of the defect detection method based on the two-stage generative adversarial network in any one of claims 1 to 8.

Images