利用YOLOv5算法实现对PCB板上的缺陷进行检测识别。
一、数据集介绍
使用的DeepPCB缺陷数据集中的所有图像都是从线性扫描CCD获得的,分辨率约为每1毫米48个像素,以上述方式从采样图像中手动检查
测试图像的原始大小约为16k x 16k像素, 然后将它们裁剪成许多大小为640 x 640的子图像,共1500张图片,DeepPCB数据集中的部分图片如下图所示。
PCB面板数据集自己去找
对于测试图像中的每个缺陷,我们使用轴对齐的边界框和一个类ID。如上所示,我们标注了六种常见的PCB缺陷类型:open、short、mousebite、spur、pin-hole、spur。
由于实际测试图像中只有少数缺陷,我们根据 PCB 缺陷模式在每个测试图像上手动论证一些人工缺陷,这导致每个640 x 640图像中大约有3到12个缺陷。
PCB缺陷数如下图所示。我们将1000 张图像作为训练集,剩下的作为测试集。
二、构建训练数据集
1、先构建数据集文件夹
本人按照VOC格式创建数据集,具体格式如下:
2、数据集格式转换
原始的标注信息是保存成txt文件,txt文件里面的每一行都包含一个标注信息,格式为x1,y1,x2,y2,type,这里 (x1,y1) 和 (x2,y2) 是缺陷边界框的左上角和右下角
type是匹配后的整数 ID:0-background、1-open、2-short、3-mousebite、4-spur、5-copper、6-pin-hole。通过一下代码进行转换:
3、训练集划分代码
主要是将数据集分类成训练数据集和测试数据集,默认train,val,test按照比例进行随机分类,运行后dataSet文件夹中会出现四个文件
主要是生成的训练数据集和测试数据集的图片名称,如下图。同时data目录下也会出现这四个文件,内容是训练数据集和测试数据集的图片路径。
4、生成yolo格式的标签
主要是将图片数据集标注后的xml文件中的标注信息读取出来并写入txt文件,运行后在label文件夹中出现所有图片数据集的标注信息
label文件夹中某文件内容如下:
三、修改配置文件
1、数据配置文件
首先需要在/yolov5-master/data文件夹中,新建一个PCBDetect.yaml文件,内容设置如下:
2、网络参数修改
对yolov5-master/model文件夹中,对yolov5x.yaml(根据自己选择的模型而定)文件内容修改。
3、trian.py修改
主要用到的几个参数:–weights,–cfg,–data,–epochs,–batch-size,–img-size,–project,-workers
重点注意:–weights,–cfg,–data,其他的默认即可(batch_size,workers根据自己电脑属性进行设置)。
四、训练及测试
1、训练
在完成上述所有的操作之后,就可以进行训练,在命令窗口输入python train.py即可以进行训练。
2、测试
在训练完成后可以利用测试集对训练后的模型进行测试,利用val.py文件进行测试,主要修改一下地方:
测试完成后会输出map、precision、recall等指标,具体如下图所示:
P-R曲线如下图所示:
同时也可以利用detect.py文件对测试集进行测试,将检测后的框绘制在图像上,部分测试结果如下图所示:
大大的demo 完事了 不过靠不靠谱我也不敢保证~~