- 数据集准备:收集并标注PCB缺陷的图像。
- 模型训练:使用YOLO v8框架训练一个模型来识别这些缺陷。
- GUI开发:利用PyQt5创建一个用户友好的图形界面。
- 模型部署:在GUI中集成训练好的模型,使用户能够上传PCB图像并得到缺陷检测的结果。
下面是对这个项目的简要介绍以及一个简单的代码示例,包括了训练代码和GUI界面的基本结构。
项目简介
1. 数据集准备
- 数据集:收集带有不同PCB缺陷类型的图像,并使用如LabelImg等工具进行标注。
- 数据集划分:将数据集分为训练集、验证集和测试集。
2. 模型训练
- 使用YOLO v8框架(如
ultralytics/yolov8
)进行训练。 - 配置文件定义模型架构和训练参数。
- 使用GPU加速训练过程。
3. GUI开发
- 使用PyQt5设计一个简洁直观的界面,允许用户上传图片、查看结果等。
- 实现模型加载和推理功能。
4. 模型部署
- 将训练好的模型部署到GUI中,使得用户可以方便地使用该系统进行PCB缺陷检测。
训练代码示例
下面是一个简单的YOLO v8训练脚本示例,用于训练PCB缺陷检测模型。
1# train.py
2
3from ultralytics import YOLO
4
5# 加载YOLO v8模型
6model = YOLO('yolov8n.pt') # 使用预训练模型作为基础
7
8# 设置训练参数
9data_config = 'data.yaml' # 数据集配置文件
10epochs = 100 # 训练轮数
11batch = 16 # 批量大小
12
13# 开始训练
14results = model.train(data=data_config, epochs=epochs, batch=batch)
GUI代码示例
下面是一个使用PyQt5创建的基本GUI界面示例,用于展示如何集成YOLO v8模型进行实时检测。
1# gui.py
2
3import sys
4from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QVBoxLayout, QLabel, QFileDialog
5from PyQt5.QtGui import QPixmap
6import cv2
7import numpy as np
8from ultralytics import YOLO
9
10class PCBDefectDetector(QWidget):
11 def __init__(self):
12 super().__init__()
13 self.initUI()
14
15 def initUI(self):
16 self.setWindowTitle('PCB Defect Detection System')
17 self.setGeometry(300, 300, 600, 400)
18
19 self.image_label = QLabel(self)
20 self.image_label.resize(400, 300)
21
22 self.load_button = QPushButton('Load Image', self)
23 self.load_button.clicked.connect(self.loadImage)
24
25 layout = QVBoxLayout()
26 layout.addWidget(self.image_label)
27 layout.addWidget(self.load_button)
28 self.setLayout(layout)
29
30 def loadImage(self):
31 options = QFileDialog.Options()
32 options |= QFileDialog.ReadOnly
33 file_name, _ = QFileDialog.getOpenFileName(self, "Open Image", "", "Image Files (*.png *.jpg *.jpeg)", options=options)
34 if file_name:
35 image = cv2.imread(file_name)
36 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
37 self.detect_defects(image)
38 height, width, channel = image.shape
39 bytes_per_line = 3 * width
40 q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888)
41 pixmap = QPixmap.fromImage(q_image)
42 self.image_label.setPixmap(pixmap)
43
44 def detect_defects(self, image):
45 # 加载模型
46 model = YOLO('path/to/best.pt') # 替换为你的模型路径
47
48 # 进行推理
49 results = model.predict(source=image, save=False)
50
51 # 处理结果
52 for r in results:
53 boxes = r.boxes
54 for box in boxes:
55 b = box.xyxy[0] # 获取边界框
56 c = box.cls # 获取分类
57 # 绘制边界框
58 cv2.rectangle(image, (int(b[0]), int(b[1])), (int(b[2]), int(b[3])), (0, 255, 0), 2)
59
60app = QApplication(sys.argv)
61ex = PCBDefectDetector()
62ex.show()
63sys.exit(app.exec_())
项目结构
train.py
:用于训练模型。gui.py
:用于运行GUI应用程序。data.yaml
:数据集配置文件。images/
和labels/
:存放训练图像和标注文件的目录。
请确保你已经安装了所有必要的库,并且正确设置了YOLO v8模型的路径。此外,你需要准备一个适当的数据集,并根据实际情况修改训练和GUI代码中的相关路径。