主要内容如下：
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
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%，甚至更小，因此基本都是很小的目标。
'''
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
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))
到此，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)
生成结果如下：
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))
运行即可，注意创建多级文件夹报错时，手动创建！
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
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：若爆显存，降低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')
预测结果：
1646
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
