【超详细】基于YOLOv8的PCB缺陷检测_yolo pcb

主要内容如下：

1、数据集介绍
2、下载PCB数据集
3、不同格式数据集预处理（Json/xml），制作YOLO格式训练集
4、模型训练及可视化

运行环境：Python=3.8（要求>=3.8），torch1.12.0+cu113（要求>=1.8）
百度AI stduio下载链接：https://aistudio.baidu.com/datasetdetail/272346

往期内容：

【超详细】跑通YOLOv8之深度学习环境配置1-Anaconda安装
【超详细】跑通YOLOv8之深度学习环境配置2-CUDA安装
【超详细】跑通YOLOv8之深度学习环境配置3-YOLOv8安装

1 数据集介绍

1.1 简介

印刷电路板（PCB）瑕疵数据集：是一个公共的合成PCB数据集，由北京大学发布，其中包含1386张图像以及6种缺陷（缺失孔，鼠标咬伤，开路，短路，杂散，伪铜），用于检测，分类和配准任务。本文我们选取了其中适用与检测任务的693张图像，随机选择593张图像作为训练集，100张图像作为验证集。

1.2 示例

2 下载数据集

官方链接：https://robotics.pkusz.edu.cn/resources/dataset/
注意：百度网盘下载，速度很慢，不推荐！推荐去百度AI stduio数据集下载，速度快！

百度AI stduio下载链接：https://aistudio.baidu.com/datasetdetail/272346

3 制作YOLO格式训练集

# 1.YOLO训练标签示例如下（转换目标）：
# 类别id 归一化后框中心点x坐标, 归一化后框中心点y坐标, 归一化后框的宽度, 归一化后框的高度
0 0.292792 0.729031 0.367417 0.246281  

# 2.训练YOLOv8路径格式如下（图像与标签路径必须对应，注意新版本不需要train.txt和val.txt文件）：
--datasets
	--images
		--train2017   # jpg/png
		--val2017
	--labels
		--train2017  # txt
		--val2017
1
2
3
4
5
6
7
8
9
10
11
12

3.1 Json格式数据集转换

3.1.1 下载后数据集布局

数据标签框分析：

import json
from collections import defaultdict
import matplotlib.pyplot as plt

# 修改1
with open("E:\\datasets\\PCB\\PCB_DATASET\\Annotations\\train.json") as f:
    data = json.load(f)

imgs = {}
for img in data['images']:
    imgs[img['id']] = {
        'h': img['height'],
        'w': img['width'],
        'area': img['height'] * img['width'],
    }

hw_ratios = []
area_ratios = []
label_count = defaultdict(int)
for anno in data['annotations']:
    hw_ratios.append(anno['bbox'][3]/anno['bbox'][2])
    area_ratios.append(anno['area']/imgs[anno['image_id']]['area'])
    label_count[anno['category_id']] += 1

print(label_count, len(data['annotations']) / len(data['images']))

plt.hist(hw_ratios, bins=100, range=[0, 2])
plt.show()

plt.hist(area_ratios, bins=100, range=[0, 0.005])
plt.show()
'''
结果如下：
(defaultdict(int, {3: 399, 5: 416, 2: 435, 6: 447, 4: 412, 1: 418}),
 4.261382799325464)
1、从标签来看，总共6个类别，如果加上背景类，总共7个类别；
2、各类别之间的框数量相对较平均，不需要调整默认的损失函数。（如果类别之间相差较大，建议调整损失函数，如BalancedL1Loss）；
3、平均每张图的框数量在4张左右，属于比较稀疏的检测；
4、真实框的宽高比，可以看到大部分集中在1.0左右，但也有部分在0.5-1之间，少部分在1.25-2.0之间；
5、真实框在原图的大小比例，可以看到大部分框只占到了原图的0.1%，甚至更小，因此基本都是很小的目标。
'''
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

3.1.2 转换代码

import json
import os

# 注意运行2次，train/val
# 修改1 json存储地址
json_path = "E:\\datasets\\PCB\\PCB_DATASET\\Annotations\\val.json"
# # 修改2，保存位置yolo存储地址
yolo_paths= "E:\\datasets\\PCB\\PCB_DATASET\\labels\\val2017/"   


with open(json_path) as f:
    data = json.load(f)

imgs = {}
for img in data['images']:
    imgs[img['id']] = {
        'h': img['height'],
        'w': img['width'],
        'file_name': img['file_name'],
    }


tmp = ''
for anno in data['annotations']:
    print(imgs[anno['image_id']]['file_name'])
    
    if imgs[anno['image_id']] != tmp:
        txt_path = os.path.join(yolo_paths, imgs[anno['image_id']]['file_name'].split('.')[0] + '.txt')
        txt_file = open(txt_path, 'w')
        # xywh --> xywh(归一化)
        bbox = [anno['bbox'][0] / imgs[anno['image_id']]['w'],
                anno['bbox'][1] / imgs[anno['image_id']]['h'],
                anno['bbox'][2] / imgs[anno['image_id']]['w'],
                    anno['bbox'][3] / imgs[anno['image_id']]['h']]
        
        cls_id = anno['category_id']
        # 保存
        txt_file.write(str(cls_id) + ' ' +" ".join([str(a) for a in bbox])+"\n")  # 生成格式0 cx,cy,w,h
        tmp = imgs[anno['image_id']]
    
    else:
        # xywh --> xywh(归一化)
        bbox = [anno['bbox'][0] / imgs[anno['image_id']]['w'],
                anno['bbox'][1] / imgs[anno['image_id']]['h'],
                anno['bbox'][2] / imgs[anno['image_id']]['w'],
                    anno['bbox'][3] / imgs[anno['image_id']]['h']]
        
        cls_id = anno['category_id']
        # 保存
        txt_file.write(str(cls_id) + ' ' +" ".join([str(a) for a in bbox])+"\n")  # 生成格式0 cx,cy,w,h    
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

3.1.3 生成训练集和验证集

注意：由于已有train2017和val2017的labels文件夹，直接按名字移动图像到对应文件夹即可。

import os
import shutil

Images_path = 'E:\\datasets\\PCB\\PCB_DATASET\\images'  # 源图路径

train_labels = 'E:\\datasets\\PCB\\PCB_DATASET\\labels\\train2017'  # train标签路径
val_labels = 'E:\\datasets\\PCB\\PCB_DATASET\\labels\\val2017'  # val标签路径

train_images = 'E:\\datasets\\PCB\\PCB_DATASET\\images\\train2017'  # 保存train图像路径
val_images = 'E:\\datasets\\PCB\\PCB_DATASET\\images\\val2017'  # 保存val图像路径


# 判断文件夹是否存在，不存在即创建
if not os.path.exists(train_images):
    os.mkdir(train_images)
if not os.path.exists(val_images):
    os.mkdir(val_images)


# 按照标签名移动对应图像
for label_name in os.listdir(train_labels):
    img_name = label_name[:-3] + 'jpg'  # txt2jpg
    shutil.move(os.path.join(Images_path, img_name), os.path.join(train_images, img_name))


for label_name in os.listdir(val_labels):
    img_name = label_name[:-3] + 'jpg'  # txt2jpg
    shutil.move(os.path.join(Images_path, img_name), os.path.join(val_images, img_name))

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

到此，YOLOv8训练数据集已经制作好，包括images(train2017与val2017)和labels(train2017与val2017)文件夹！

3.2 VOC格式数据集转换

3.2.1 下载后数据集布局

xml文件有用信息：对应图像文件名、对应图像大小、每个框信息（包括：类别、xyxy具体框位置信息）

3.2.2 转换代码

# 实现xml格式转yolov5格式
import xml.etree.ElementTree as ET
import os

# box [xmin,ymin,xmax,ymax]
def convert(size, box):
    x_center = (box[2] + box[0]) / 2.0
    y_center = (box[3] + box[1]) / 2.0
    # 归一化
    x = x_center / size[0]
    y = y_center / size[1]
    # 求宽高并归一化
    w = (box[2] - box[0]) / size[0]
    h = (box[3] - box[1]) / size[1]
    return (x, y, w, h)

# xml2yolo
def convert_annotation(xml_paths, yolo_paths, classes):
    xml_files = os.listdir(xml_paths)
    # 生成无序文件列表
    for file in xml_files:
        
        xml_file_path = os.path.join(xml_paths, file)
        yolo_txt_path = os.path.join(yolo_paths, file.split(".")[0]
                                     + ".txt")
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        size = root.find("size")
        # 获取xml的width和height的值
        w = int(size.find("width").text)
        h = int(size.find("height").text)
        # object标签可能会存在多个，所以要迭代
        with open(yolo_txt_path, 'w') as f:
            for obj in root.iter("object"):
                difficult = obj.find("difficult").text
                # 类别
                cls = obj.find("name").text
                if cls not in classes or difficult == 1:
                    continue
                # 转换成训练模式读取的标签
                cls_id = classes.index(cls)
                xml_box = obj.find("bndbox")
                box = (float(xml_box.find("xmin").text), float(xml_box.find("ymin").text),
                       float(xml_box.find("xmax").text), float(xml_box.find("ymax").text))
                boxex = convert((w, h), box)
                # yolo标准格式类别 x_center,y_center,width,height
                f.write(str(cls_id) + " " + " ".join([str(s) for s in boxex]) + '\n')
        print(f'xml_file :{file} --> txt Saved!')


if __name__ == "__main__":
    # PCB数据的类别
    classes_train = ['missing_hole', 'mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper']   # 修改1，类别
    
    # xml存储地址
    xml_dir = "E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\Annotations/"  # 修改2，读取位置
    
    # yolo存储地址
    yolo_txt_dir = "E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\labels/"   # 修改3，保存位置
    
    # voc转yolo
    convert_annotation(xml_paths=xml_dir, yolo_paths=yolo_txt_dir,
                       classes=classes_train)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

生成结果如下：

3.2.3 生成训练集和验证集

注意：这边下载数据集没有分train2017和val2017文件夹，所以需自己按比例或按数据量进行划分！

import os
import shutil
import random

Images_path = 'E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\JPEGImages'  # 源图路径
Labels_path = 'E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\labels'  # 源标签路径

train_labels = 'E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\labels\\train2017'  # train标签路径
val_labels = 'E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\labels\\val2017'  # val标签路径

train_images = 'E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\images\\train2017'  # 保存train图像路径
val_images = 'E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\images\\val2017'  # 保存val图像路径

radio = 0.2  # 按照比例划分的验证集比例
nums = 100  # 按照数量划分的验证集数量
is_radio = False  # 如果为True，则按照比例进行划分，否则按照数量划分

# 判断文件夹是否存在，不存在即创建
if not os.path.exists(train_images):
    os.mkdir(train_images)
if not os.path.exists(val_images):
    os.mkdir(val_images)

if not os.path.exists(train_labels):
    os.mkdir(train_labels)
if not os.path.exists(val_labels):
    os.mkdir(val_labels)


Imgs = os.listdir(Images_path)
if is_radio:
    val_nums = int(len(Imgs) * radio)
else:
    val_nums = nums
val_Imgs = random.sample(Imgs, val_nums)

for val_name in val_Imgs:
    shutil.move(os.path.join(Images_path, val_name), os.path.join(val_images, val_name))

    val_name = val_name[:-3] + 'txt'  # jpg2txt
    shutil.move(os.path.join(Labels_path, val_name), os.path.join(val_labels, val_name))

if (len(Imgs) - len(val_Imgs)) > 0:
    for i in val_Imgs:
        if i in Imgs:
            Imgs.remove(i)
    train_Imgs = Imgs
    for train_name in train_Imgs:
        shutil.move(os.path.join(Images_path, train_name), os.path.join(train_images, train_name))

        train_name = train_name[:-3] + 'txt'  # jpg2txt
        shutil.move(os.path.join(Labels_path, train_name), os.path.join(train_labels, train_name))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

运行即可，注意创建多级文件夹报错时，手动创建！

4 模型训练及可视化

4.1 创建数据集yaml文件

注意：路径一定填对，类别与id一定要对应（下面用的是上面xml格式生成的数据集训练）！！！
创建ultralytics\cfg\datasets\PCB.yaml文件，内容如下：

path: E:\\datasets\\PCB\\VOCdevkit\\VOC2007 # dataset root dir
train: images/train2017 # train images (relative to 'path') 4 images
val: images/val2017 # val images (relative to 'path') 4 images

# Classes for DOTA 1.0
names:
  0: missing_hole
  1: mouse_bite
  2: open_circuit
  3: short
  4: spur
  5: spurious_copper
1
2
3
4
5
6
7
8
9
10
11
12

4.2 创建一个训练脚本

在主目录下创建一个train.py，内容如下：

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    # model = YOLO("yolov8n.yaml")  # build a new model from scratch
    model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)

    # Use the model
    model.train(data="PCB.yaml", imgsz=640, batch=16, workers=8, cache=True, epochs=100)  # train the model
    metrics = model.val()  # evaluate model performance on the validation set
    # results = model("ultralytics\\assets\\bus.jpg")  # predict on an image
    path = model.export(format="onnx", opset=13)  # export the model to ONNX format
1
2
3
4
5
6
7
8
9
10
11
12

问题1：若爆显存，降低batch和workers大小！

训练结果如下：

4.2 创建一个预测脚本

在主目录下创建一个detect.py，内容如下：

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO("runs\\detect\\train\\weights\\best.pt")  # load model
    model.predict(source="E:\\datasets\\PCB\\VOCdevkit\\VOC2007\\images\\val2017", save=True, save_conf=True, save_txt=True, name='output')
1
2
3
4
5
6

预测结果：