【表面缺陷检测】基于yolov5的PCB表面缺陷检测(附代码)

今天手把手带大家做一个深度学习实战项目：基于yolov5的PCB(印刷电路板)表面缺陷检测。yolov5模型的训练在Windows系统下进行，CUDA版本:11.6，驱动版本:512.77，显卡：NVIDIA GeForce RTX 2060。

阅读完本文后，你将知道：

• PCB的全称是printed circuit board，中文名称为印制电路板，是电子工业的重要部件之一

• PCB的创造者是奥地利人保罗·爱斯勒(Paul eisler)

• PCB表面缺陷数据的获取方式

• 如何使用yolov5训练自己的数据集

多的不说，少的不唠，下面开始今天的教程。

项目背景

  PCB是printed circuit board的缩写，中文名称为印制电路板，由于它是采用电子印刷术制作的，又被称为“印刷电路板”，它是电子工业的重要部件之一，计算器、计算机、电子手表、收音机、手机等都需要PCB。  PCB产品的优良，关系到企业的发展。如果良品率太低，出厂的PCB被用于电子产品，将会损害客户的利益，并影响企业的口碑和发展。因此，遏制不合格PCB产品的出厂，将会给企业降低风险，树立优质的企业形象，增加企业在行业的竞争力。  人工检测生产出来的PCB产品表面缺陷，不仅效率低，而且容易出现漏检和错检。采用深度学习的方法自动检测生产出来的PCB产品表面缺陷，一方面降低了人力成本，另一方面提高了检出率，降低了漏检率和错检率，还可以对检测出来的表面缺陷进行分析，及时的优化产品的产业链，从而提高良品率。

获取数据集

  目前，可以获取由北京大学发布的一个公共的PCB表面缺陷数据集，其中包含1386张图像以及6种缺陷（缺失孔，老鼠咬坏，开路，短路，毛刺，杂散铜），可直接用于目标检测的数据信息如下：

缺陷类型(英文名称)缺陷类型(中文名称)图片个数标签个数

693张缺陷图片可直接用于目标检测模型训练，使用labelimg工具查看数据：

数据存放结构如下：VOCdevkit  --VOC2007    --Annotations     --01_missing_hole_01.xml     --01_missing_hole_02.xml     --...    --JPEGImages     --01_missing_hole_01.jpg     --01_missing_hole_02.jpg     --...

下载yolov5代码

yolov5下载地址：https://github.com/ultralytics/yolov5

如上图所示，本次教程下载的是yolov5-6.2版本，将yolov5-6.2.zip压缩包放入新建的yolov5_pcb文件夹下并解压：

使用PyCharm编辑器打开yolov5_pcb/yolov5-6.2项目，并将存放PCB表面缺陷数据集的VOCdevkit 文件夹考本过去，如下图所示：

环境配置

4.1创建虚拟环境

打开本地的Anaconda Prompt：

在Anaconda Prompt终端输入命令，创建名为yolov5的虚拟环境：

# 环境名称：yolov5
# python版本：3.9.12
conda create --name yolov5 python=3.9.12

这样，Anaconda中就拥有了名为yolov5的虚拟环境。

有3点需要说明：

• 1.创建新环境需要连接网络

• 2.环境名称可以任意起，请自行安排

• 3.可以不创建新的虚拟环境，使用现有的环境（创建虚拟环境的优点是避免破坏之前项目的环境）

4.2安装依赖库

4.2.1项目解释器

返回到PyCharm编辑器打开yolov5_pcb/yolov5-6.2项目，点击File -> Settings -> Project:yolov5-6.2 -> Python Interpreter -> Add Interpreter -> Add Local Interpreter... -> System Interpreter -> ... > 找到Anaconda路径下的yolov5虚拟环境

4.2.2安装项目依赖库

打开PyCharm底部的Terminal，激活yolov5环境：

根据requirements.txt中的依赖库，使用pip命令进行依次安装：

也可以使用pip install -r requirements.txt命令进行全部安装。

4.3下载预训练权重

yolov5有不同的模型结构，因此提供了不同的预训练权重文件：

下载yolov5s.pt权重（当然，也可以下载其他的模型权重），将其放在新建的weights文件夹，并拷贝到yolov5_pcb/yolov5-6.2项目下：

4.4测试环境

在终端输入命令，测试环境是否可用：

python detect.py --weights weights/yolov5s.pt --source data/images/bus.jpg

如果看到以下运行结果，说明环境配置成功：

如果出现如下提示信息，说明缺少依赖库：

ModuleNotFoundError: No module named xxx

请根据提示信息，使用pip命令安装依赖库：

pip install xxx

测试图片的检测结果保存在runs/detect/exp/文件夹下：

配置文件

5.1数据配置

在data文件夹下新建一个VOC_PCB.yaml文件（可以自定义命名）：

写入内容如下：

# 注意路径，否则会报错，和6.1节数据划分相关联
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../yolov5-6.2/VOCdevkit/
train: # train images (relative to 'path')  16551 images
  - images/train
val: # val images (relative to 'path')  4952 images
  - images/val
test: # test images (optional)
  - images/test

# Classes
nc: 6  # number of classes
names: ["missing_hole", "mouse_bite", "open_circuit", "short", "spur", "spurious_copper"]  # class names

5.2模型配置

由于yolov5提供了好几个模型，本次教程我们训练yolov5s，将models/yolov5s.yaml的内容复制到新建的yolov5s_pcb.yaml:

修改yolov5s_pcb.yaml中类别数量：

训练

6.1数据集划分

编写prepare_train_val.py脚本，将数据集划分为训练集和测试集：

训练集图片数：554测试集图片数：139

获取prepare_train_val.py脚本：https://blog.csdn.net/AugustMe/article/details/127320880

6.2模型训练

在终端执行命令：

# 指定预训练权重、模型结构、图片路径、训练轮数、批次大小、图片大小
python train.py --weights weights/yolov5s.pt --cfg models/yolov5s_pcb.yaml --data data/VOC_PCB.yaml --epochs 50 --batch-size 4 --imgsz 1280

开始训练：

6.3过程及结果

查看GPU使用情况：

训练过程可视化：

打开一个Anaconda Prompt终端，切换到yolov5_pcb/yolov5-6.2目录下，激活yolov5虚拟环境，在终端输入：

tensorboard --logdir=runs

也可以直接在PyCharm里面新开一个终端，在终端输入：

tensorboard --logdir=runs

浏览器输入：http://localhost:6006/

训练结束后，最新模型在测试集上各个缺陷指标数据：

测试集上PR曲线：

训练好的模型保存在runs/train/exp/weights/目录下：

测试

新建一个testpics文件夹，拷贝2张新收集的PCB缺陷图片到testpics，在终端执行检测命令：

python detect.py --weights runs/train/exp/weights/best.pt --source  testpics

在runs/detect/文件夹下查看检测结果：

总结

在Windows系统下，我们使用一张NVIDIA GeForce RTX 2060显卡，基于PCB表面缺陷检测数据集，训练yolov5中的5s模型，50个epoch过后，测试集上六类缺陷的mAP@0.5为0.982，说明我们训练得到的模型可以准确地实现PCB表面缺陷检测。

Bug修复

训练yolov5可能出现的bug：

cv2.error: OpenCV(4.5.5) D:\a\opencv-python\opencv-python\opencv\modules\core\src\alloc.cpp:73: error: (-4:Insufficient memory) Failed to allocate 19926864 bytes in function 'cv::OutOfMemoryError'

OSError: [WinError 1455] 页面文件太小，无法完成操作。Error loading "D:\anaconda3\envs\yolov5\lib\site-packages\torch\lib\shm.dll" or one of its dependencies.

以上报错信息我已在文章中详细地撰写了解决方案，请点击链接进行查看：https://blog.csdn.net/AugustMe/article/details/127497030

参考资料

[1]https://blog.csdn.net/AugustMe/article/details/127368318[2]https://blog.csdn.net/AugustMe/article/details/127322898[3]https://blog.csdn.net/AugustMe/article/details/127497030[4]https://blog.csdn.net/weixin_45679938/article/details/118803745