Abstract:
Generative Adversarial Network (GAN) often demands extensive training data to produce high-quality images. However, the manufacturing industry grapples with a significant shortage of training data for industrial workpieces, presenting challenges in applying conventional GAN models to data augmentation. Consequently, we propose a GAN model capable of training on a small-scale artifact dataset to generate high-quality artifact images, serving as a method to augment the few-sample artifact dataset. Self-attention mechanisms are integrated into both the generator and discriminator, incorporating attention mask and attention mapping based on the positions of workpiece voids. This weighted approach enhances the correlation between the regions of workpiece voids and surrounding pixels, thereby mitigating the reliance on extensive training data to a certain extent. A redesigned residual structure is implemented in the generator, employing a combination of upsampling and convolution to ameliorate the "chessboard artifact" phenomenon in generated images, thereby enhancing their realism. The loss function is formulated as a combination of Wasserstein distance and weighted feature matching loss. In comparison to traditional GANs, the proposed model demonstrates a reduction in FID score to 100. 91 and an elevation in SSIM score to 0. 906 for generated workpiece images. After employing the proposed GAN model for data augmentation, the mean Average Precision (mAP) value based on the YOLOv8 defect detection algorithm is elevated to 92. 7%. This method can offer a solution to the inadequacy of training samples in industrial inspection.