一种基于生成对抗网络模型的工件图像数据增广方法

A Method of Augmenting Workpiece Image Data Based on Generative Adversarial Network Model

  • 摘要: 生成对抗网络(generative adversarial network,GAN)往往需要大量的训练数据才能生成高质量图像,而制造业工件训练数据严重匮乏,难以利用传统GAN模型进行数据增广。为此,提出一种能够以小规模工件数据集训练生成高质量工件图像的GAN模型,作为少样本工件数据集的增广方法。对生成器和鉴别器融入自注意力机制,依据工件孔洞分布特点创建注意力掩码与注意力映射进行加权,以提高工件孔洞区域与周围像素点的相关性,在一定程度上减少对大规模训练数据的依赖。重新设计残差结构并应用于生成器,利用上采样和卷积组合的方式改善生成图像的“棋盘格伪影”现象,以提高生成图像的逼真度。损失函数采用Wasserstein距离和特征匹配损失加权相结合的形式。与传统GAN对比,所提模型生成工件图像的FID分数降低至100. 91,SSIM分数提升至0. 906。经所提GAN模型数据增广后,基于YOLOv8算法的工件缺陷检测模型的mAP值提升至92. 7%,可为工业检测训练样本不足提供解决方案。

     

    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.

     

/

返回文章
返回