基于CycleGAN和Pix2Pix的轨面缺陷图像智能生成技术

Intelligent Generation of Rail Surface Defect Images Based on CycleGAN and Pix2Pix

为解决基于人工智能的高动态高精度轨道巡检技术的小样本学习难题,提出基于CycleGAN和Pix2Pix生成的对抗网络模型,实现小样本数据集语义特征学习和钢轨轨面缺陷数据的智能生成。其中,Pix2Pix模型生成特定类别的轨面图像;而CycleGAN模型将无缺陷轨面图像转换成有缺陷轨面图像,且缺陷样式不受定式约束。因此,在维持轨面缺陷类别不变而形态各异的基础上,实现了轨面缺陷数据集的大规模增强,解决了小样本数据集存在的数据分布不平衡、数据缺乏多样性以及数据标注难度高等难题。利用VGG19、YOLOv5和UNet进行性能测试,试验表明:生成的轨面缺陷图像增强数据集在图像分类任务中的准确率为81.177%,较原数据集增加了23.138%;在目标检测任务中,准确率为91.90%,增加了26.60%,召回率为87.20%,增加了16.00%,均值平均精度为93.50%,增加了18.30%;在语义分割任务中Dice得分为71.015,较原数据集提高6%。研究结果对解决轨道巡检技术小样本学习难题具有重要应用价值。

In order to solve the small-sample learning problem of highly dynamic and high-precision track inspection technology based on artificial intelligence,two generative adversarial network models based on CycleGAN and Pix2Pix were proposed to achieve semantic feature learning of small-sample datasets and intelligent generation of rail surface defect data.The Pix2Pix model generated class-specific rail surface images while the CycleGAN model converted defect-free rail surface images into defective rail surface images,with the defect style not constrained by the stereotype.In this way,a large-scale enhancement of the rail surface defect dataset was achieved on the basis of maintaining the same category of rail surface defects with different forms,solving the problems of imbalanced data distribution in small sample datasets,lack of diversity in data,and high difficulty in data annotation.Performance experiments using VGG19,YOLOv5 and UNet show that the accuracy of the rail surface defect image enhancement dataset generated in this paper is 81.177%in the image classification task,23.138%better than the original dataset.In the target detection task,the accuracy is 91.90%,up 26.60%,with respect to the recall rate of 87.20%,an increase of 16.00%,and mean average precision of 93.50%,an improvement of 18.30%.In the semantic segmentation task,the Dice score is 71.015,an improvement of 6%over the original dataset.The research results demonstrate important application value in solving the small sample learning problem of track inspection technology.