A composite coating defect dataset augmentation method based on improved Pix2Pix

Abstract

The application relates to the technical field of deep learning, and discloses a composite coating defect data set amplification method based on improved Pix2Pix, which comprises data preprocessing, conditional adversarial generation network, generator design of Pix2Pix, discriminator design, and adversarial ability control design of the generator and the discriminator; the data preprocessing comprises screening of original defect pictures, annotation of semantic segmentation pictures, conversion of the prepared data set pictures and labels into a tensor type that can be processed by a deep learning model, and normalization of the pictures to-1~1; the generator design is used for converting a given semantic segmentation picture into a defect picture, using a U-Net structure and improving the performance of the model by introducing an attention mechanism; the discriminator design is used for evaluating the generated defect sample pictures; and the adversarial ability control design of the generator and the discriminator is used for balancing the abilities of the generator and the discriminator, so that the generator can better generate pictures containing defect features. Through the application, the problems of small sample quantity and imbalance in a defect recognition task are solved, and the defect classification model has stronger generalization ability.

Description

Technical Field

[0001] This invention relates to the field of deep learning, specifically to a method for augmenting composite material coating defect datasets based on an improved Pix2Pix. Background Technology

[0002] Composite material coatings are crucial functional materials for bonding engine propellant to the thermal insulation layer. When using image processing methods to detect defects in engine composite material coatings, insufficient sample sizes can lead to poor detection and identification results. Generative Adversarial Networks (GANs) are powerful deep learning models. By utilizing the generator network of a GAN, combined with advanced image processing techniques and the characteristics of composite material coating defects, highly realistic image samples with specific defect features can be generated. This method can be used to generate a large number of composite material coating defect sample images of different types and severity, providing reliable data support for the detection and evaluation of composite material coating defects. Furthermore, it can be used to optimize and improve other defect detection algorithms and techniques. Summary of the Invention

[0003] (a) Technical problems to be solved

[0004] To address the shortcomings of existing technologies, this invention provides a method for augmenting composite material coating defect datasets based on an improved Pix2Pix, thus solving the problems mentioned in the background technology.

[0005] (II) Technical Solution

[0006] To achieve the above objectives, the present invention provides the following technical solution: a method for augmenting composite material coating defect datasets based on an improved Pix2Pix, comprising the following steps:

[0007] S1. Obtaining semantic segmentation samples based on conditional generative adversarial networks:

[0008] S11, Data Preprocessing:

[0009] S111. Semantic Segmentation Image Creation: Create three folders, "0", "1", and "2", to represent defect-free images, images with missing coating defects, and images with excessive coating, respectively. Manually screen the obtained original images, determine their defect types, and place them in the corresponding folders. A total of 450 defect-free images, 486 images with missing coating defects, and 443 images with excessive coating were used. The labelme tool was used to annotate the defects in the original images. The labeled files were in JSON format. The json library in Python was used to parse the JSON files, and the OpenCV library was used to draw the semantic segmentation images.

[0010] S112. Labels converted to one-hot encoding: Converting labels to one-hot encoding can provide a unified representation, enabling the generator and discriminator networks to work together better and improving the performance of generative adversarial networks; after the labels are converted to one-hot encoding, they can be more easily used in convolution or deconvolution.

[0011] S113. Image tensorization and normalization: Use the ToTensor and Normalize methods in PyTorch to convert each input grayscale image to a value between -1 and 1.

[0012] S114. Custom Dataset Class: Define a custom dataset class that inherits from the Dataset class to process your own dataset. In the class, define members such as getting the number of files, getting the index value, getting the label value, and getting the image path value for flexible image processing.

[0013] S115. Dataset Creation: Load semantic segmentation images onto the custom dataset class described in S114, set the batch size BATCH_SIZE to 16, and enable random and repeated reading.

[0014] S12, Generator Design:

[0015] S121. Create random noise: Define a noise of length 500 that follows a normal distribution; this noise is the input to the generator. After the noise data passes through a series of fully connected layers and deconvolution layers, it is transformed into a 1x256x256 grayscale image, thus forming the target image.

[0016] S122. Change the shape of the label data: After one-hot encoding, the label will become a tensor of length 3. After passing through several fully connected layers of neural network, the shape of the vector will be transformed into (128, 16, 16).

[0017] S123. Change the shape of the input noise: The original shape of the input noise is a random tensor of length 500 that follows a normal distribution. After passing through several fully connected layers of neural network, the shape of the vector is transformed into (128, 16, 16).

[0018] S124. Merge labels and noise: Merge the input noise and labels in the channel dimension after the shape is changed, that is, into data with the shape of (256, 16, 16);

[0019] S125, Deconvolution and Batch Normalization: Deconvolution is performed on the merged data. In order to form the final target image size and obtain truly useful defect feature information, each deconvolution is activated using the leaky_relu function, ultimately forming data of shape (1, 256, 256), i.e., a grayscale image.

[0020] S13. Discriminator Design:

[0021] S131. Change the shape of the label data: The discriminator receives two types of data: one is the generated image, and the other is the label after one-hot encoding. In order for the discriminator to judge the two, the label needs to be converted into a 1x256x256 shape through a fully connected neural network, and then flattened into a (1, 256, 256) shape.

[0022] S132. Merge labels and generated images: Merge the labels with shape (1, 256, 256) and the generated images with shape (1, 256, 256) into data with shape (2, 256, 256).

[0023] S133. Design Convolutional Layers: Design convolutional layers by sequentially increasing the number of channels in the merged shape while decreasing the width and height. Each convolution is activated using the leaky_relu function; ultimately forming a tensor with the shape (128, 63, 63).

[0024] S134, Batch Standardization: Perform batch standardization on tensors of shape (128, 63, 63);

[0025] S135, Data Dimensionality Reduction: Flatten the tensor of (128, 63, 63) into a tensor of length 128*63*63;

[0026] S136, Fully Connected Layer: Use a fully connected layer to connect a 128*63*63 tensor to a unit 1 tensor, and then use sigmoid to activate it;

[0027] S14. Setting the learning rate: Set the learning rate of the generator to 0.001 and the learning rate of the discriminator to 0.0001. The advantage of this is that it weakens the ability of the discriminator, so that the generator has the ability to compete with the discriminator in the early stage.

[0028] S15. Configure model training parameters:

[0029] S151. Device Selection: To ensure better compatibility of deep learning code, the availability of GPUs in the hardware devices is first determined. If the GPU is unavailable, the CPU is selected as the device.

[0030] S152. Loss Function: Since the final output of the discriminator is a probability value of 0 to 1, which belongs to a binary classification problem, BCELoss is chosen as the loss function for training.

[0031] S153, Optimizer Selection: Adam is selected as the optimizer for both the generator and the discriminator. Except for the learning rate as described in S14, all other parameters of Adam are set to default values.

[0032] S16. Saving the Model: A trained neural network consists of two parts: structure and weights. After training is complete, all the information is saved in the model, but at this point it is only in memory. The save method saves the network's framework and weights to the "CGAN_model.pth" file. When needed, only the pth file needs to be loaded.

[0033] S2. Generating defect images based on the improved Pix2Pix:

[0034] S21. Data Preprocessing:

[0035] S211, Threshold Segmentation: Using the method in S1, images obtained during training can be acquired. Images with good quality and defect features are manually collected and placed in the "Pix2PixData" dataset. Initial processing of the generated images is performed using OpenCV-based threshold segmentation, median filtering, and Gaussian filtering techniques to remove low-brightness pixel values, salt-and-pepper noise, and Gaussian noise. A total of 1600 images of missing coating defects and 1500 images of over-coating defects are obtained.

[0036] S212, Semantic Segmentation: The Labelme tool is used to perform semantic segmentation on the noise-removed defect image. The semantically segmented image drawn using the OpenCV library will be used as the label input for Pix2Pix to generate the defect image.

[0037] S213. Create two variables, imgs_path and annos_path, and store all paths of the original defect image and the semantic segmentation image into the variables using the glob method;

[0038] S214. Image tensorization and normalization: Use the ToTensor and Normalize methods in PyTorch to convert each input grayscale image to a value between -1 and 1.

[0039] S215. Custom Dataset Class: Define a custom dataset class that inherits from the Dataset class to process your own dataset. In the class, define members such as getting the number of files, getting the index value, getting the label value, and getting the image path value for flexible image processing.

[0040] S216. Dataset Creation: Based on the custom dataset class described in S111, load the original images and corresponding semantic segmentation images, set the batch value BATCH_SIZE to 16, and enable random reading and repeated reading;

[0041] S22, Pix2Pix generator design:

[0042] S221. Define the downsampling class: This class is named Downsample and consists of a convolutional layer, a LeakyReLU activation function, and a batch normalization. The batch normalization has a trigger condition that requires the parameter is_bn to be set to True for activation. This design allows for flexible use of batch normalization.

[0043] S222. Define the upsampling class: The class is named Upsample. This class consists of a deconvolutional layer, a LeakyReLU activation function, and a batch normalization. The dropout trigger condition is set. It will only be activated when the parameter is_drop is True, which is used to increase the randomness of the generator.

[0044] S223. Define a channel attention mechanism class: Name the class ChannelAttention. This class defaults to a channel reduction ratio of 16 and creates an adaptive average pooling layer to pool the input feature map into a 1x1 feature map. Define a sequence model consisting of two linear layers and an activation function to learn channel weights. First, a linear layer reduces the number of channels in the input feature map to 1 / 16 of its original value. Then, ReLU activation is applied. Next, another linear layer restores the number of channels to its original value. Finally, a Sigmoid activation function is applied to obtain the channel attention weights. A forward propagation function is used to calculate the output of the channel attention. Global average pooling reduces the dimensionality of the input feature map to the average value of each channel. The learned channel weights are then applied to obtain weighted channel features. This weighted feature is multiplied by the original input feature map to produce enhanced channel attention features.

[0045] S224. Define a spatial attention mechanism class: Name the class SpatialAttention; create a convolutional layer with 2 input channels (a stack of the average and maximum values) and 1 output channel, with a kernel size of 3 or 7, and no bias. By calculating the average and maximum values ​​of the input feature maps, two feature maps are obtained. These two feature maps are then stacked and input into the convolutional layer to obtain enhanced spatial attention features.

[0046] S225. Creating a generator: By combining upsampling and downsampling, a U-Net structure is constructed, outputting an image of (3, 256, 256). This data is then converted to grayscale to obtain the required grayscale image. In order to enhance the model's ability to pay attention to the input data and improve model performance, an attention mechanism is introduced into the U-Net structure, and an attention module is added to the layer between upsampling and downsampling.

[0047] S23, Pix2Pix Discriminator Design:

[0048] S231. Discriminator input image merging: The discriminator receives two data: one is the generated image, and the other is the semantic segmentation image. Both are (1, 256, 256) images. The two are merged to become (2, 256, 256) data.

[0049] S232, Output Discriminant Value: Through convolution, batch normalization, and fully connected layers, the output is a data of size 1. Finally, the sigmoid function is used for activation to obtain a value between 0 and 1, which is used as the discriminant value.

[0050] S24. Setting the learning rate: Set the learning rate of the generator to 0.0001 and the learning rate of the discriminator to 0.00002; and set a decay coefficient for the learning rate. The advantage of this is that it weakens the ability of the discriminator, so that the generator has the ability to compete with the discriminator in the early stage.

[0051] S25. Configure model training parameters:

[0052] S251. Device Selection: To ensure better compatibility of deep learning code, the availability of GPUs in the hardware devices is first determined. If the GPU is unavailable, the CPU is selected as the device.

[0053] S252. Generator Loss Function: The cross-entropy loss function of binary classification is used to compare the output of the discriminator to the generated image with the all-1 vector and calculate the loss. In addition, L1 loss is introduced, which is used to constrain the similarity between the generator output and the real image. Taking all factors into account, the total loss of the generator is a linear combination of the cross-entropy loss and the L1 loss.

[0054] S253. Discriminator Loss Function: The discriminator loss function is used to encourage the discriminator to classify the images generated by the generator as real images. It usually consists of two parts: real image loss and generated image loss. The real image loss is used to train the discriminator to distinguish real images. It uses a binary classification loss function to convert the discriminator's output for real images into the generated image loss. The total discriminator loss is the sum of these two losses.

[0055] S254, Optimizer Selection: Adam is selected as the optimizer for both the generator and the discriminator. Except for the learning rate as described in S24, all other parameters of Adam are set to default values.

[0056] S26. Model Saving: A trained generator consists of two parts: structure and weights. After training is complete, all information is saved in the model, but at this point it is only in memory. The save method saves the network framework and weights to the "Pix2Pix_model.pth" file. When needed, only the pth file needs to be loaded.

[0057] S27. Generation of defective images: Based on the generator model described in S26, a loop is constructed, taking the semantic segmentation image as described in S212 as input, and using the model to predict and obtain the output defective image.

[0058] Preferably, in S111, Labelme is one of the most commonly used tools for marking defects on images. In addition, tools such as Labelmeg or CVAT are also suitable as marking tools of the present invention.

[0059] Preferably, in step S115, BATCH_SIZE is set to 16, meaning 16 images and their corresponding labels are retrieved each time for training. If the computer has a high configuration and a large amount of graphics card memory, BATCH_SIZE can be set to a larger value. Since repeated training is required, the repeat function is used to generate data repeatedly. In order to enable the network to achieve better generalization ability, shuffle is used to shuffle the data.

[0060] Preferably, in step S121, the magnitude of the random noise, 500, is a hyperparameter that can be adjusted according to the actual usage. The random pattern is generally selected as a normal distribution, but other distribution types such as Poisson distribution or exponential distribution can also be tried.

[0061] Preferably, in step S14, the learning rate of the generator should be higher than that of the discriminator, as this can prevent the loss value of the discriminator from decreasing too quickly, which would prevent the discriminator from guiding the generator.

[0062] Preferably, in step S225, the channel attention mechanism class and the spatial attention mechanism class can be placed at any position between upsampling and downsampling.

[0063] Preferably, in S252, the generator loss function is a linear combination of the cross-entropy loss function and the L1 loss.

[0064] (III) Beneficial Effects

[0065] This invention provides a method for augmenting composite material coating defect datasets based on an improved Pix2Pix, which has the following advantages:

[0066] This invention is applied to the data augmentation stage of composite material coating defect detection tasks. By using an improved Pix2Pix to generate samples with different defect types, covering different defect shapes and details, it can increase the quantity and quality of training data, thereby improving the robustness and generalization ability of the defect classification model. In addition, it can balance the sample distribution of each category, so that the subsequent defect identification and classification model can be fully identified in all categories. Attached Figure Description

[0067] Figure 1 A generator structure diagram for a conditional generative adversarial network;

[0068] Figure 2 A diagram of the discriminator structure for a conditional generative adversarial network;

[0069] Figure 3 Example of an image generated by a conditional adversarial generative network;

[0070] Figure 4 Here is a diagram of the generator structure for Pix2Pix;

[0071] Figure 5 This is a diagram of the discriminator structure for Pix2Pix;

[0072] Figure 6 Example of an image generated for Pix2Pix;

[0073] Figure 7 This is a technical roadmap for the overall solution; Detailed Implementation

[0074] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0075] This invention provides a technical solution: a method for augmenting composite material coating defect datasets based on an improved Pix2Pix, comprising the following steps:

[0076] S1. Obtaining semantic segmentation samples based on conditional generative adversarial networks:

[0077] S11, Data Preprocessing:

[0078] S111. Semantic Segmentation Image Creation: Create three folders, "0", "1", and "2", to represent defect-free images, images with missing coating defects, and images with excessive coating, respectively. Manually screen the obtained original images, determine their defect types, and place them in the corresponding folders. A total of 450 defect-free images, 486 images with missing coating defects, and 443 images with excessive coating were used. Labelme was used to annotate the defects in the original images. The files annotated with Labelme were in JSON format. The JSON file was parsed using the json library in Python, and the semantic segmentation images were drawn using the OpenCV library.

[0079] S112. Labels converted to one-hot encoding: Converting labels to one-hot encoding can provide a unified representation, enabling the generator and discriminator networks to work together better and improving the performance of generative adversarial networks; after the labels are converted to one-hot encoding, they can be more easily used in convolution or deconvolution.

[0080] S113. Image tensorization and normalization: Use the ToTensor and Normalize methods in PyTorch to convert each input grayscale image to a value between -1 and 1.

[0081] S114. Custom Dataset Class: Define a custom dataset class that inherits from the Dataset class to process your own dataset. In the class, define members such as getting the number of files, getting the index value, getting the label value, and getting the image path value for flexible image processing.

[0082] S115. Dataset Creation: Load semantic segmentation images onto the custom dataset class described in S114, set the batch size BATCH_SIZE to 16, and enable random and repeated reading.

[0083] S12, Generator Design:

[0084] S121. Create random noise: Define a noise of length 500 that follows a normal distribution; this noise is the input to the generator. After a series of fully connected layers and deconvolutional layers, the noise data is transformed into a 1x256x256 grayscale image, which forms the target image.

[0085] S122. Change the shape of the label data: After one-hot encoding, the label will become a tensor of length 3. After passing through several fully connected layers of neural network, the shape of the vector will be transformed into (128, 16, 16).

[0086] S123. Change the shape of the input noise: The original shape of the input noise is a random tensor of length 500 that follows a normal distribution. After passing through several fully connected layers of neural network, the shape of the vector is transformed into (128, 16, 16).

[0087] S124. Merge labels and noise: Merge the input noise and labels in the channel dimension after the shape is changed, that is, into data with the shape of (256, 16, 16);

[0088] S125, Deconvolution and Batch Normalization: Deconvolution is performed on the merged data to form the final target image size and obtain truly useful defect feature information. Each deconvolution operation uses the leaky_relu function for activation, ultimately forming data of shape (1, 256, 256), i.e., a grayscale image. The deconvolution and batch normalization processes are as follows: Figure 1 As shown;

[0089] S13. Discriminator Design:

[0090] S131. Change the shape of the label data: The discriminator receives two types of data: one is the generated image, and the other is the label after one-hot encoding. In order for the discriminator to judge the two, the label needs to be converted into a 1x256x256 shape through a fully connected neural network, and then flattened into a (1, 256, 256) shape.

[0091] S132. Merge labels and generated images: Merge the labels with shape (1, 256, 256) and the generated images with shape (1, 256, 256) into data with shape (2, 256, 256).

[0092] S133. Design Convolutional Layers: Design convolutional layers by sequentially increasing the number of channels in the merged shape while decreasing the width and height. Each convolution is activated using the leaky_relu function; ultimately forming a tensor with the shape (128, 63, 63).

[0093] S134, Batch Standardization: Perform batch standardization on tensors of shape (128, 63, 63);

[0094] S135, Data Dimensionality Reduction: Flatten the tensor of (128, 63, 63) into a tensor of length 128*63*63;

[0095] S136, Fully Connected Layer: Use a fully connected layer to connect a 128*63*63 tensor to a unit 1 tensor, and then use sigmoid to activate it;

[0096] S14. Setting the learning rate: Set the learning rate of the generator to 0.001 and the learning rate of the discriminator to 0.0001. The advantage of this is that it weakens the ability of the discriminator, so that the generator has the ability to compete with the discriminator in the early stage.

[0097] S15. Configure model training parameters:

[0098] S151. Device Selection: To ensure better compatibility of deep learning code, the availability of GPUs in the hardware devices is first determined. If the GPU is unavailable, the CPU is selected as the device.

[0099] S152. Loss Function: Since the final output of the discriminator is a probability value of 0 to 1, which belongs to a binary classification problem, BCELoss is chosen as the loss function for training.

[0100] S153, Optimizer Selection: Adam is selected as the optimizer for both the generator and the discriminator. Except for the learning rate as described in S14, all other parameters of Adam are set to default values.

[0101] S16. Model Saving: A trained neural network consists of two parts: structure and weights. After training is complete, all information is saved in the `model` object, but it's currently only in memory. The `save` method saves the network's framework and weights to the "CGAN_model.pth" file. When needed, simply load the `.pth` file. Some generated images from the training process are also included. Figure 3 As shown;

[0102] S2. Generating defect images based on the improved Pix2Pix:

[0103] S21, Data Preprocessing

[0104] S211, Threshold Segmentation: Using the method in S1, images obtained during training can be acquired. Images with good quality and defect features are manually collected and placed in the "Pix2PixData" dataset. Initial processing of the generated images is performed using OpenCV-based threshold segmentation, median filtering, and Gaussian filtering techniques to remove low-brightness pixel values, salt-and-pepper noise, and Gaussian noise. A total of 1600 images of missing coating defects and 1500 images of over-coating defects are obtained.

[0105] S212, Semantic Segmentation: The Labelme tool is used to perform semantic segmentation on the noise-removed defect image. The semantically segmented image drawn using the OpenCV library will be used as the label input for Pix2Pix to generate the defect image.

[0106] S213. Create two variables, imgs_path and annos_path, and store all paths of the original defect image and the semantic segmentation image into the variables using the glob method;

[0107] S214. Image tensorization and normalization: Use the ToTensor and Normalize methods in PyTorch to convert each input grayscale image to a value between -1 and 1.

[0108] S215. Custom Dataset Class: Define a custom dataset class that inherits from the Dataset class to process your own dataset. In the class, define members such as getting the number of files, getting the index value, getting the label value, and getting the image path value for flexible image processing.

[0109] S216. Dataset Creation: Based on the custom dataset class described in S111, load the original images and corresponding semantic segmentation images, set the batch value BATCH_SIZE to 16, and enable random reading and repeated reading;

[0110] S22, Pix2Pix generator design:

[0111] S221. Define the downsampling class: This class is named Downsample and consists of a convolutional layer, a LeakyReLU activation function, and a batch normalization. The batch normalization has a trigger condition that requires the parameter is_bn to be set to True for activation. This design allows for flexible use of batch normalization.

[0112] S222. Define the upsampling class: The class is named Upsample. This class consists of a deconvolutional layer, a LeakyReLU activation function, and a batch normalization. The dropout trigger condition is set. It will only be activated when the parameter is_drop is True, which is used to increase the randomness of the generator.

[0113] S223. Define a channel attention mechanism class: Name the class ChannelAttention. This class defaults to a channel reduction ratio of 16 and creates an adaptive average pooling layer to pool the input feature map into a 1x1 feature map. Define a sequence model consisting of two linear layers and an activation function to learn channel weights. First, a linear layer reduces the number of channels in the input feature map to 1 / 16 of its original value. Then, ReLU activation is applied. Next, another linear layer restores the number of channels to its original value. Finally, a Sigmoid activation function is applied to obtain the channel attention weights. A forward propagation function is used to calculate the output of the channel attention. Global average pooling reduces the dimensionality of the input feature map to the average value of each channel. The learned channel weights are then applied to obtain weighted channel features. This weighted feature is multiplied by the original input feature map to produce enhanced channel attention features.

[0114] S224. Define a spatial attention mechanism class: Name the class SpatialAttention; create a convolutional layer with 2 input channels (a stack of the average and maximum values) and 1 output channel, with a kernel size of 3 or 7, and no bias. By calculating the average and maximum values ​​of the input feature maps, two feature maps are obtained. These two feature maps are then stacked and input into the convolutional layer to obtain enhanced spatial attention features.

[0115] S225. Generator Creation: By combining upsampling and downsampling, a U-Net structure is constructed, outputting an image of (3, 256, 256). This data is then converted to grayscale to obtain the desired grayscale image. To enhance the model's ability to pay attention to input data and improve model performance, an attention mechanism is introduced into the U-Net structure. An attention module is added to the layer between upsampling and downsampling. The generator structure is as follows: Figure 4 As shown;

[0116] S23, Pix2Pix Discriminator Design:

[0117] S231. Discriminator input image merging: The discriminator receives two data: one is the generated image, and the other is the semantic segmentation image. Both are (1, 256, 256) images. The two are merged to become (2, 256, 256) data.

[0118] S232. Output Discriminant Value: Through convolution, batch normalization, and fully connected layers, the output is a data value of size 1. Finally, the sigmoid function is used for activation to obtain a value between 0 and 1, which is used as the discriminant value. The structure of the discriminator is as follows: Figure 5 As shown;

[0119] S24. Setting the learning rate: Set the learning rate of the generator to 0.0001 and the learning rate of the discriminator to 0.00002; and set a decay coefficient for the learning rate. The advantage of this is that it weakens the ability of the discriminator, so that the generator has the ability to compete with the discriminator in the early stage.

[0120] S25. Configure model training parameters

[0121] S251. Device Selection: To ensure better compatibility of deep learning code, the availability of GPUs in the hardware devices is first determined. If the GPU is unavailable, the CPU is selected as the device.

[0122] S252. Generator Loss Function: The cross-entropy loss function of binary classification is used to compare the output of the discriminator to the generated image with the all-1 vector and calculate the loss. In addition, L1 loss is introduced, which is used to constrain the similarity between the generator output and the real image. Taking all factors into account, the total loss of the generator is a linear combination of the cross-entropy loss and the L1 loss.

[0123] S253. Discriminator Loss Function: The discriminator loss function is used to encourage the discriminator to classify the images generated by the generator as real images. It usually consists of two parts: real image loss and generated image loss. The real image loss is used to train the discriminator to distinguish real images. It uses a binary classification loss function to convert the discriminator's output for real images into the generated image loss. The total discriminator loss is the sum of these two losses.

[0124] S254, Optimizer Selection: Adam is selected as the optimizer for both the generator and the discriminator. Except for the learning rate as described in S24, all other parameters of Adam are set to default values.

[0125] S26. Model Saving: A trained generator consists of two parts: structure and weights. After training is complete, all information is saved in the model, but at this point it is only in memory. The save method saves the network framework and weights to the "Pix2Pix_model.pth" file. When needed, only the pth file needs to be loaded.

[0126] S27. Generation of Defect Images: Based on the generator model described in S26, a loop is constructed, taking the semantically segmented image as described in S212 as input, and using the model to predict and obtain the output defect image. Some example images generated are shown below. Figure 6 As shown.

[0127] Although embodiments of the invention have been shown and described, it will be understood by those skilled in the art that various changes, modifications, substitutions and alterations can be made to these embodiments without departing from the principles and spirit of the invention, the scope of which is defined by the appended claims and their equivalents.

Claims

1. A method for augmenting composite material coating defect datasets based on an improved Pix2Pix, characterized in that, Includes the following steps: S1. Obtaining semantic segmentation samples based on conditional generative adversarial networks: S11, Data Preprocessing: S111. Semantic Segmentation Image Creation: Create three folders, "0", "1", and "2", to represent defect-free images, images with missing coating defects, and images with excessive coating, respectively. Manually screen the obtained original images, determine their defect types, and place them in the corresponding folders. A total of 450 defect-free images, 486 images with missing coating defects, and 443 images with excessive coating were used. Labelme was used to annotate the defects in the original images. The files annotated with Labelme were in JSON format. The JSON file was parsed using the json library in Python, and the semantic segmentation images were drawn using the OpenCV library. S112. Labels converted to one-hot encoding: Converting labels to one-hot encoding can provide a unified representation, enabling the generator and discriminator networks to work together better and improving the performance of generative adversarial networks; after the labels are converted to one-hot encoding, they can be more easily used in convolution or deconvolution. S113. Image tensorization and normalization: Use the ToTensor and Normalize methods in PyTorch to convert each input grayscale image to a value between -1 and 1. S114. Custom Dataset Class: Define a custom dataset class that inherits from the Dataset class to process your own dataset. In the class, define members such as getting the number of files, getting the index value, getting the label value, and getting the image path value for flexible image processing. S115. Dataset Creation: Load semantic segmentation images onto the custom dataset class described in S114, set the batch size BATCH_SIZE to 16, and enable random and repeated reading. S12, Generator Design: S121. Create random noise: Define a noise of length 500 that follows a normal distribution; this noise is the input to the generator. After a series of fully connected layers and deconvolutional layers, the noise data is transformed into a 1x256x256 grayscale image, which forms the target image. S122. Change the shape of the label data: After one-hot encoding, the label will become a tensor of length 3. After passing through several fully connected layers of neural network, the shape of the vector will be transformed into (128, 16, 16). S123. Change the shape of the input noise: The original shape of the input noise is a random tensor of length 500 that follows a normal distribution. After passing the tensor through several fully connected layers of a neural network, the shape of the vector is transformed into (128, 16, 16). S124. Merge labels and noise: Merge the input noise and labels in the channel dimension after the shape is changed, that is, into data with the shape of (256, 16, 16); S125, Deconvolution and Batch Normalization: Deconvolution is performed on the merged data. In order to form the final target image size and obtain truly useful defect feature information, each deconvolution is activated using the leaky_relu function, which ultimately forms data of shape (1, 256, 256), i.e., a grayscale image. S13, Discriminator Design S131. Change the shape of the label data: The discriminator receives two types of data: one is the generated image, and the other is the label after one-hot encoding. In order for the discriminator to judge the two, the label needs to be converted into a 1x256x256 shape through a fully connected neural network, and then flattened into a (1, 256, 256) shape. S132. Merge labels and generated images: Merge the labels with shape (1, 256, 256) and the generated images with shape (1, 256, 256) into data with shape (2, 256, 256). S133. Design Convolutional Layers: Design convolutional layers by sequentially increasing the number of channels and decreasing the width and height of the merged shape. Each deconvolution is activated using the leaky_relu function; the final result is a tensor with the shape (128, 63, 63). S134, Batch Standardization: Perform batch standardization on tensors of shape (128, 63, 63); S135, Data Dimensionality Reduction: Flatten the tensor of (128, 63, 63) into a tensor of length 128 x 63 x 63; S136, Fully Connected Layer: Use a fully connected layer to connect (128, 63, 63) to a tensor of unit 1, and then use sigmoid for activation; S14. Setting the learning rate: Set the learning rate of the generator to 0.001 and the learning rate of the discriminator to 0.0001. The advantage of this is that it weakens the ability of the discriminator, so that the generator has the ability to compete with the discriminator in the early stage. S15, Configure model training parameters S151. Device Selection: To ensure better compatibility of deep learning code, the availability of GPUs in the hardware devices is first determined. If the GPU is unavailable, the CPU is selected as the device. S152. Loss Function: Since the final output of the discriminator is a probability value of 0 to 1, which belongs to a binary classification problem, BCELoss is chosen as the loss function for training. S153, Optimizer Selection: Adam is selected as the optimizer for both the generator and the discriminator. Except for the learning rate as described in S14, all other parameters of Adam are set to default values. S16. Saving the Model: A trained neural network consists of two parts: structure and weights. After training is complete, all the information is saved in the model, but at this point it is only in memory. The save method saves the network's framework and weights to the "CGAN_model.pth" file. When needed, simply load the pth file. S2. Generating defect images based on the improved Pix2Pix: S21, Data Preprocessing S211, Threshold Segmentation: Using the method in S1, images obtained during training can be acquired. Images with good quality and defect features are manually collected and placed in the "Pix2PixData" dataset. Initial processing of the generated images is performed using OpenCV-based threshold segmentation, median filtering, and Gaussian filtering techniques to remove low-brightness pixel values, salt-and-pepper noise, and Gaussian noise. A total of 1600 images of missing coating defects and 1500 images of over-coating defects are obtained. S212, Semantic Segmentation: The Labelme tool is used to perform semantic segmentation on the noise-removed defect image. The semantically segmented image drawn using the OpenCV library will be used as the label input for Pix2Pix to generate the defect image. S213. Create two variables, imgs_path and annos_path, and store all paths of the original defect image and the semantic segmentation image into the variables using the glob method; S214. Image tensorization and normalization: Use the ToTensor and Normalize methods in PyTorch to convert each input grayscale image to a value between -1 and 1. S215. Custom Dataset Class: Define a custom dataset class that inherits from the Dataset class to process your own dataset. In the class, define members such as getting the number of files, getting the index value, getting the label value, and getting the image path value for flexible image processing. S216. Dataset Creation: Based on the custom dataset class described in S111, load the original images and corresponding semantic segmentation images, set the batch value BATCH_SIZE to 16, and enable random reading and repeated reading; S22, Pix2Pix generator design: S221. Define the downsampling class: This class is named Downsample and consists of a convolutional layer, a LeakyReLU activation function, and a batch normalization. The batch normalization has a trigger condition that requires the parameter is_bn to be set to True for activation. This design allows for flexible use of batch normalization. S222. Define the upsampling class: The class is named Upsample. This class consists of a deconvolutional layer, a LeakyReLU activation function, and a batch normalization. The dropout trigger condition is set. It will only be activated when the parameter is_drop is True, which is used to increase the randomness of the generator. S223. Define a channel attention mechanism class: The class is named ChannelAttention. The default channel reduction ratio in this class is 16, and an adaptive average pooling layer is created to pool the input feature map into a 1x1 feature map. Define a sequence model, including two linear layers and an activation function, to learn the channel weights. First, a linear layer is used to reduce the number of channels in the input feature map to 1 / 16 of the original number, then ReLU activation is applied, followed by another linear layer to restore the number of channels to the original value, and finally the Sigmoid activation function is applied to obtain the channel attention weights. The output of the channel attention is calculated using the forward propagation function; The input feature map is reduced to the average value of each channel using global average pooling. Then, it is processed with learned channel weights to obtain weighted channel features. This weighted feature is multiplied by the original input feature map to produce enhanced channel attention features. S224. Define a spatial attention mechanism class: Name the class SpatialAttention; create a convolutional layer with 2 input channels (a stack of the average and maximum values) and 1 output channel, with a kernel size of 3 or 7, and no bias. By calculating the average and maximum values ​​of the input feature maps, two feature maps are obtained. These two feature maps are then stacked and input into the convolutional layer to obtain enhanced spatial attention features. S225. Creating a generator: By combining upsampling and downsampling, a U-Net structure is constructed, outputting an image of (3, 256, 256). This data is then converted to grayscale to obtain the required grayscale image. In order to enhance the model's ability to pay attention to the input data and improve model performance, an attention mechanism is introduced into the U-Net structure, and an attention module is added to the layer between upsampling and downsampling. S23, Pix2Pix Discriminator Design: S231. Discriminator input image merging: The discriminator receives two data: one is the generated image, and the other is the semantic segmentation image. Both are (1, 256, 256) images. The two are merged to become (2, 256, 256) data. S232, Output Discriminant Value: Through convolution, batch normalization, and fully connected layers, the output is a data of size 1. Finally, the sigmoid function is used for activation to obtain a value between 0 and 1, which is used as the discriminant value. S24. Setting the learning rate: Set the learning rate of the generator to 0.0001 and the learning rate of the discriminator to 0.00002; and set a decay coefficient for the learning rate. The advantage of this is that it weakens the ability of the discriminator, so that the generator has the ability to compete with the discriminator in the early stage. S25. Configure model training parameters S251. Device Selection: To ensure better compatibility of deep learning code, the availability of GPUs in the hardware devices is first determined. If the GPU is unavailable, the CPU is selected as the device. S252. Generator Loss Function: The cross-entropy loss function of binary classification is used to compare the output of the discriminator to the generated image with the all-1 vector and calculate the loss. In addition, L1 loss is introduced, which is used to constrain the similarity between the generator output and the real image. Taking all factors into account, the total loss of the generator is a linear combination of the cross-entropy loss and the L1 loss. S253. Discriminator Loss Function: The discriminator loss function is used to encourage the image generated by the generator to be judged as a real image by the discriminator. It usually includes two parts: real image loss and generated image loss. The real image loss is used to train the discriminator to distinguish real images. It uses a binary classification loss function to generate the image loss from the discriminator's output on real images. The total discriminator loss is the sum of these two losses. S254, Optimizer Selection: Adam is selected as the optimizer for both the generator and the discriminator. Except for the learning rate as described in S24, all other parameters of Adam are set to default values. S26. Model Saving: A trained generator consists of two parts: structure and weights. After training is complete, all information is saved in the model, but at this point it is only in memory. The save method saves the network framework and weights to the "Pix2Pix_model.pth" file. When needed, only the pth file needs to be loaded. S27. Generation of defective images: Based on the generator model described in S26, a loop is constructed, taking the semantic segmentation image as described in S212 as input, and using the model to predict and obtain the output defective image.

2. The method for augmenting composite material coating defect datasets based on improved Pix2Pix according to claim 1, characterized in that: In step S112, the label data of length 1 is converted into one-hot encoding, which allows it to participate better in convolution or deconvolution operations.

3. The method for augmenting composite material coating defect datasets based on improved Pix2Pix according to claim 1, characterized in that: In S124, the input noise and the label data have the same shape after passing through the fully connected layer, and the two are merged in the channel dimension.

4. The method for augmenting composite material coating defect datasets based on improved Pix2Pix according to claim 1, characterized in that: In S22, the generator design incorporates channel attention mechanism and spatial attention mechanism.

5. The method for augmenting composite material coating defect datasets based on improved Pix2Pix according to claim 1, characterized in that: In S252, the generator loss function is a linear combination of cross-entropy loss and L1 loss.

Images