基于YOLOv5-7.0的PCB板缺陷检测

原文发布于 CSDN：https://blog.csdn.net/m0_49683806/article/details/133899280

参考引用

你的陈某某-基于YOLOv5的PCB板缺陷检测

一、数据集介绍

印刷电路板（PCB）瑕疵数据集。它是一个公共合成PCB数据集，包含1386张图像，具有6种缺陷（漏孔、鼠咬、开路、短路、杂散、杂铜），用于图像检测、分类和配准任务。

下载地址：数据集

数据样本示例：

二、环境配置

1、Gitub官网下载yolov5源码：官方地址

2、Anaconda 安装配置（省略）

3、创建新的环境（python=3.8就行，因为使用的是7.0的版本，3.6有点低。）

conda create -n pytorch python=3.8.5

4、安装pytorch：pytorch官网

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

5、yolov5的依赖下载

根据下载的 yolov5中的requirements.txt进行安装（缺啥补啥）【注意：具体安装以yolov5的readme.md为主！】

pip install -r requirements.txt

三、构建训练数据集

1、先构建数据集文件夹

下载好的PCB_DATASET解压后。是下面这种格式。

官网下载的PCB数据按照缺陷类划分文件夹的

将每个文件夹内的东西都复制出来到这个文件中。

然后删掉分类的文件夹。

images文件夹内的内容也是一样的。

├── PCB_DATASET
│   ├── Annotations  进行 detection 任务时的标签文件，xml 形式，文件名与图片名一一对应
│   ├── images  存放 .jpg 格式的图片文件
│   ├── ImageSets  存放的是分类和检测的数据集分割文件，包含 train.txt,val.txt ,trainval.txt,test.txt
│   ├── labels  存放label标注信息的txt文件，与图片一一对应

├── ImageSets(train，val，test建议按照8：1：1比例划分)
│   ├── train.txt  写着用于训练的图片名称
│   ├── val.txt  写着用于验证的图片名称
│   ├── trainval.txt  train与val的合集
│   ├── test.txt  写着用于测试的图片名称

2、训练数据生成，分为两个代码（训练集划分代码与用于yolo训练的txt格式代码）

（1）训练集划分代码

在PCB_DATASET文件夹下新建ImageSets文件夹。

用途：主要是将数据集分类成训练数据集和测试数据集，默认train，val，test按照比例进行随机分类，运行后ImagesSets文件夹中会出现四个文件，主要是生成的训练数据集和测试数据集的图片名称。

import os
import random

trainval_percent = 0.9
train_percent = 0.9
xmlfilepath = 'D:\\pcb\\PCB_DATASET\\Annotations'
txtsavepath = 'D:\\pcb\\PCB_DATASET\\ImageSets'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open(txtsavepath+'\\trainval.txt', 'w')
ftest = open(txtsavepath+'\\test.txt', 'w')
ftrain = open(txtsavepath+'\\train.txt', 'w')