一、数据集获取与介绍
"刷电路板（PCB）瑕疵数据集" 是一个公开可用的合成PCB图像集合，由北京大学的研究团队精心制作并发布。该数据集涵盖了1386幅高分辨率图像，详细记录了六类常见的制造缺陷：孔洞缺失、边缘损伤（俗称“鼠标咬”）、线路断开（开路）、短路连接、非结构化残留物（杂散）及虚假铜层现象。这些图像为研究人员提供了宝贵资源，适用于开发和测试针对PCB缺陷的检测、分类及对齐算法。
在本研究中，我们精选了数据集中适合进行缺陷检测任务的693幅图像。为了确保模型训练的有效性和泛化能力，我们采取了科学合理的数据分割策略，即从这693幅图像中随机挑选出593幅作为训练样本，剩余的100幅则用作验证集，以评估模型性能。通过这种方式，我们不仅能够充分利用现有资源，还能确保研究结果的可靠性和可重复性。
数据集目录如下所示：
PCB
├── Annotations
│ ├── 01_missing_hole_01.xml
│ ├── 01_missing_hole_02.xml
│ ├── ............................................
├──images
│ ├──01_missing_hole_01.jpg
│ ├──01_missing_hole_02.jpg
│ ├── .....................................
有需要的小伙伴可以点击我上传的资源进行数据集下载，自行进行训练。
https://download.youkuaiyun.com/download/weixin_44765053/89936044二、模型训练
1.划分数据集
通过以下脚本进行训练集、验证集以及测试集的划分，例如我以下的比例0.7：0.15：0.15。
import os
import random
import shutil
import time
import yaml
class YOLOTrainDataSetGenerator:
def __init__(self, origin_dataset_dir, train_dataset_dir, train_ratio=0.7, val_ratio=0.15, test_ratio=0.15,
clear_train_dir=False):
# 设置随机数种子
random.seed(1233)
self.origin_dataset_dir = origin_dataset_dir
self.train_dataset_dir = train_dataset_dir
self.train_ratio = train_ratio
self.val_ratio = val_ratio
self.test_ratio = test_ratio
self.clear_train_dir = clear_train_dir
assert self.train_ratio > 0.5, 'train_ratio must larger than 0.5'
assert self.val_ratio > 0.01, 'train_ratio must larger than 0.01'
assert self.test_ratio > 0.01, 'test_ratio must larger than 0.01'
total_ratio = round(self.train_ratio + self.val_ratio + self.test_ratio)
assert total_ratio == 1.0, 'train_ratio + val_ratio + test_ratio must equal 1.0'
def generate(self):
time_start = time.time()
# 原始数据集的图像目录，标签目录，和类别文件路径
origin_image_dir = os.path.join(self.origin_dataset_dir, 'images')
origin_label_dir = os.path.join(self.origin_dataset_dir, 'labels')
origin_classes_file = os.path.join(self.origin_dataset_dir, 'classes.txt')
if not os.path.exists(origin_classes_file):
return
else:
origin_classes = {}
with open(origin_classes_file, mode='r') as f:
for cls_id, cls_name in enumerate(f.readlines()):
cls_name = cls_name.strip()
if cls_name != '':
origin_classes[cls_id] = cls_name
# 获取所有原始图像文件名（包括后缀名）
origin_image_filenames = os.listdir(origin_image_dir)
# 随机打乱文件名列表
random.shuffle(origin_image_filenames)
# 计算训练集、验证集和测试集的数量
total_count = len(origin_image_filenames)
train_count = int(total_count * self.train_ratio)
val_count = int(total_count * self.val_ratio)
test_count = total_count - train_count - val_count
# 定义训练集文件夹路径
if self.clear_train_dir and os.path.exists(self.train_dataset_dir):
shutil.rmtree(self.train_dataset_dir, ignore_errors=True)
train_dir = os.path.join(self.train_dataset_dir, 'train')
val_dir = os.path.join(self.train_dataset_dir, 'val')
test_dir = os.path.join(self.train_dataset_dir, 'test')
train_image_dir = os.path.join(train_dir, 'images')
train_label_dir = os.path.join(train_dir, 'labels')
val_image_dir = os.path.join(val_dir, 'images')
val_label_dir = os.path.join(val_dir, 'labels')
test_image_dir = os.path.join(test_dir, 'images')
test_label_dir = os.path.join(test_dir, 'labels')
# 创建训练集输出文件夹
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)
os.makedirs(test_image_dir, exist_ok=True)
os.makedirs(test_label_dir, exist_ok=True)
# 将图像和标签文件按设定的ratio划分到训练集，验证集，测试集中
for i, filename in enumerate(origin_image_filenames):
if i < train_count:
output_image_dir = train_image_dir
output_label_dir = train_label_dir
elif i < train_count + val_count:
output_image_dir = val_image_dir
output_label_dir = val_label_dir
else:
output_image_dir = test_image_dir
output_label_dir = test_label_dir
src_img_name_no_ext = os.path.splitext(filename)[0]
src_image_path = os.path.join(origin_image_dir, filename)
src_label_path = os.path.join(origin_label_dir, src_img_name_no_ext + '.txt')
if os.path.exists(src_label_path):
# 复制图像文件
dst_image_path = os.path.join(output_image_dir, filename)
shutil.copy(src_image_path, dst_image_path)
# 复制标签文件
src_label_path = os.path.join(origin_label_dir, src_img_name_no_ext + '.txt')
dst_label_path = os.path.join(output_label_dir, src_img_name_no_ext + '.txt')
shutil.copy(src_label_path, dst_label_path)
else:
pass
train_dir = os.path.normpath(train_dir)
val_dir = os.path.normpath(val_dir)
test_dir = os.path.normpath(test_dir)
data_dict = {
'train': train_dir,
'val': val_dir,
'test': test_dir,
'nc': len(origin_classes),
'names': origin_classes
}
yaml_file_path = os.path.normpath(os.path.join(self.train_dataset_dir, 'data.yaml'))
with open(yaml_file_path, mode='w') as f:
yaml.safe_dump(data_dict, f, default_flow_style=False, allow_unicode=True)
if __name__ == '__main__':
g_origin_dataset_dir = '原始数据集'
g_train_dataset_dir = '生成的文件地址'
g_train_ratio = 0.7
g_val_ratio = 0.15
g_test_ratio = 0.15
yolo_generator = YOLOTrainDataSetGenerator(g_origin_dataset_dir, g_train_dataset_dir, g_train_ratio, g_val_ratio,
g_test_ratio, True)
yolo_generator.generate()2.data.yaml文件
新建一个data.yaml文件，用于存储训练数据的路径及模型需要进行检测的类别。YOLOv8在进行模型训练时，会读取该文件的信息，用于进行模型的训练与验证。data.yaml的具体内容如下：
train: .\dataset\pcb\train # train images
val: .\dataset\pcb\val # val images
test: # val images (optional)
# number of classes
nc: 6
# Classes
names: ['missing_hole', 'mouse_bite', 'open_circuit', 'short', 'spur', 'spurious_copper']
3.模型训练
使用以下代码进行YOLOv8模型训练，具体内容可以看这篇
【YOLOv8系列】（四）YOLOv8使用自己的数据集进行模型训练，成就感满满_yolov8训练自己的数据集-优快云博客
#coding:utf-8
from ultralytics import YOLO
# 加载预训练模型
model = YOLO("yolov8n.pt")
# Use the model
if __name__ == '__main__':
# Use the model
results = model.train(data='/pcb/data.yaml', epochs=300, batch=4) # 训练模型训练结果指标如下所示：
三、pyqt界面实现
1.ui设计
ui界面设计如下所示：
包含 ：
QLabel，显示图像、序号、类别、置信度等
QTableWidget，显示序号、类别、置信度等识别信息
QPushButton，选择图片、文件夹、视频或者摄像头检测，保存识别csv
2.功能实现
图片检测
文件夹检测
四、总结
由于博主的能力有限，文中提到的方法虽经过实验验证，但难免存在一些不足之处。为不断提升内容的质量与准确性，欢迎您指出任何错误和疏漏。这不仅将帮助我在下次更新时更加完善和严谨，也能让其他读者受益。您的反馈对我至关重要，能够推动我进一步完善相关内容。
此外，如果您有更优秀的实现方案或独到的见解，也非常欢迎分享。这将为大家提供更多思路与选择，促进我们共同的成长与进步。期待您的宝贵建议与经验交流，非常感谢您的支持！
扫一扫
401
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
