YOLOv8应用于产品缺陷检测实例
- 项目概况
- 项目实现
- YOLOv8安装及模型训练
- 关键代码展示
- 动态效果展示
项目概况
本项目是应用YOLOv8框架实现训练自定义模型实现单一零件的缺陷检测，软件界面由PyQt5实现。
功能已正式使用，识别效果达到预期。
项目实现
项目使用了以下几个要素：
- 全新的界面设计 ，PyQt5结合QtDesigner自定义界面设计，快速构建想要的UI；
- 相机选型和光源 ，项目使用了迈德威视工业相机GYD-GE130M-T，对照官方给出的文档很容易就可以实现相机的使用，光源方面使用12V常规灯带打光，去除其他光源的干扰；
- IO触发 功能，实现串口控制IO，以快速输出检测识别结果，便于对接工业传送装置；
YOLOv8安装及模型训练
可参考该大佬文章: 传送门
项目展示图: 
关键代码展示
迈德威视相机使用线程代码片.
import time
import cv2
from ultralytics import YOLO
import numpy as np
from PyQt5 import QtGui
from PyQt5.QtCore import pyqtSignal, QThread
from MindVisionSdk import mvsdk
import warnings
warnings.filterwarnings('ignore')
class MyThread(QThread):
camera_signal = pyqtSignal(object) # 输出相机数据流
info_signal = pyqtSignal(str) # 输出信息状态
def __init__(self, parent=None):
super().__init__(parent)
self.drawLine = False
self.monoCamera = None
self.FrameBufferSize = None
self.isPause = False
self.mCamera = None
self.pFrameBuffer = None
self.model = YOLO(model="./YoloModel/best.pt")
def __del__(self):
mvsdk.CameraUnInit(self.mCamera)
mvsdk.CameraAlignFree(self.pFrameBuffer)
def run(self):
# 枚举相机设备列表
DevList = mvsdk.CameraEnumerateDevice()
if len(DevList) >= 1:
DevInfo = DevList[0] # 选取设备列表第一个相机
try:
self.mCamera = mvsdk.CameraInit(DevInfo, -1, -1)
self.info_signal.emit("【{}】 初始化相机成功".format(time.strftime("%Y-%m-%d %H:%M:%S")))
except mvsdk.CameraException as e:
self.info_signal.emit("【{}】 初始化相机异常:{}".format(time.strftime("%Y-%m-%d %H:%M:%S"), e))
cap = mvsdk.CameraGetCapability(self.mCamera)
self.monoCamera = (cap.sIspCapacity.bMonoSensor != 0)
if self.monoCamera:
mvsdk.CameraSetIspOutFormat(self.mCamera, mvsdk.CAMERA_MEDIA_TYPE_MONO8)
else:
mvsdk.CameraSetIspOutFormat(self.mCamera, mvsdk.CAMERA_MEDIA_TYPE_BGR8)
mvsdk.CameraSetTriggerMode(self.mCamera, 0)
mvsdk.CameraSetAeState(self.mCamera,True) # True为自动曝光 False为手动曝光
mvsdk.CameraSetRotate(self.mCamera, 0)
mvsdk.CameraSetMirror(self.mCamera, 0, True)
mvsdk.CameraPlay(self.mCamera)
# 计算RGB buffer所需的大小，这里直接按照相机的最大分辨率来分配
self.FrameBufferSize = cap.sResolutionRange.iWidthMax * cap.sResolutionRange.iHeightMax * (
1 if self.monoCamera else 3)
# 分配RGB buffer，用来存放ISP输出的图像
# 备注：从相机传输到PC端的是RAW数据，在PC端通过软件ISP转为RGB数据（如果是黑白相机就不需要转换格式，但是ISP还有其它处理，所以也需要分配这个buffer）
self.pFrameBuffer = mvsdk.CameraAlignMalloc(self.FrameBufferSize, 16)
while True:
try:
if not self.isPause:
pRawData, FrameHead = mvsdk.CameraGetImageBuffer(self.mCamera, 10)
mvsdk.CameraImageProcess(self.mCamera, pRawData, self.pFrameBuffer, FrameHead)
if self.drawLine:
mvsdk.CameraImageOverlay(self.mCamera, self.pFrameBuffer, FrameHead)
mvsdk.CameraReleaseImageBuffer(self.mCamera, pRawData)
# 此时图片已经存储在pFrameBuffer中，对于彩色相机pFrameBuffer=RGB数据，黑白相机pFrameBuffer=8位灰度数据
# 把pFrameBuffer转换成opencv的图像格式以进行后续算法处理
frame_data = (mvsdk.c_ubyte * FrameHead.uBytes).from_address(self.pFrameBuffer)
frame = np.frombuffer(frame_data, dtype=np.uint8)
frame = frame.reshape((FrameHead.iHeight, FrameHead.iWidth,
1 if FrameHead.uiMediaType == mvsdk.CAMERA_MEDIA_TYPE_MONO8 else 3))
frame = cv2.resize(frame,(960,540), interpolation=cv2.INTER_LINEAR)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = self.model.predict(source=frame,verbose=False)
# 提取检测结果
# for result in results:
# boxes = result.boxes.xyxy # 边界框坐标
# scores = result.boxes.conf # 置信度分数
# classes = result.boxes.cls # 类别索引
# # 如果有类别名称，可以通过类别索引获取
# class_names = [self.model.names[int(cls)] for cls in classes]
# # 打印检测结果
# for box, score, class_name in zip(boxes, scores, class_names):
# print(f"Class: {class_name}, Score: {score:.2f}, Box: {box}")
annotated_frame = results[0].plot()
showImage = QtGui.QImage(annotated_frame.data, frame.shape[1], frame.shape[0], QtGui.QImage.Format_RGB888)
self.camera_signal.emit(showImage)
self.info_signal.emit("【{}】 实时显示中......".format(time.strftime("%Y-%m-%d %H:%M:%S")))
except mvsdk.CameraException as e:
if e.error_code != mvsdk.CAMERA_STATUS_TIME_OUT:
self.info_signal.emit("【{}】 实时显示线程报错:{}".format(time.strftime("%Y-%m-%d %H:%M:%S"), e))
finally:
self.msleep(20)
主界面逻辑代码片
import os.path
import sys
import time
import cv2
from PyQt5 import QtWidgets
from PyQt5 import QtGui
from PyQt5.QtWidgets import QMessageBox
from WindowUI import Ui_MainWindow
from CameraThread import MyThread
from MindVisionSdk import mvsdk
from Yaml_Tool import myYamlTool
class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
def __init__(self):
super(MainWindow, self).__init__()
self.setupUi(self)
self.camera_thread = MyThread()
self.camera_thread.camera_signal.connect(self.showVideo)
self.camera_thread.info_signal.connect(self.showStatus)
self.pushButton_open_close.clicked.connect(self.camera_start_stop)
self.pushButton_play.clicked.connect(self.camera_play_pause)
self.pushButton_play.setEnabled(False)
self.pushButton_catch.clicked.connect(self.catchPicture)
self.pushButton_addline.clicked.connect(self.drawLine)
self.radioButton_run.setStyleSheet("QRadioButton::indicator:!checked {border: 2px solid #8f8f91;}"
"QRadioButton::indicator:checked {background-color: #55ff7f;}"
"QRadioButton::indicator:checked:hover {border: 2px solid #8f8f91;}")
self.radioButton_mark.setStyleSheet("QRadioButton::indicator:!checked {border: 2px solid #8f8f91;}"
"QRadioButton::indicator:checked {background-color: #55ff7f;}"
"QRadioButton::indicator:checked:hover {border: 2px solid #8f8f91;}")
self.yamlTool = myYamlTool()
self.yamlData = self.yamlTool.read_yaml('config.yaml')
if self.yamlData['runMode'] == 'mark':
self.radioButton_mark.setChecked(True)
else:
self.radioButton_run.setChecked(True)
self.radioButton_mark.clicked.connect(lambda: self.selectMode(0))
self.radioButton_run.clicked.connect(lambda: self.selectMode(1))
self.doubleSpinBox_threshold.setValue(self.yamlData['threshold'])
self.doubleSpinBox_threshold.valueChanged.connect(self.thresholdChange)
def closeEvent(self, event):
"""
重写关闭按钮事件
:param event:
:return:
"""
reply = QMessageBox.question(self, '警告', '<font color=red><b>确定退出工具?</b></font>',
QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
if reply == QMessageBox.Yes:
if self.camera_thread.isRunning():
self.camera_thread.__del__()
self.camera_thread.terminate()
event.accept()
else:
event.ignore()
def camera_start_stop(self):
"""
打开关闭相机按钮事件
:return:
"""
if self.pushButton_open_close.text() == "打开相机":
self.camera_thread.start()
self.camera_thread.isPause = False
self.pushButton_open_close.setText("关闭相机")
self.statusbar.showMessage("【{}】 已打开相机".format(time.strftime("%Y-%m-%d %H:%M:%S")))
self.pushButton_play.setEnabled(True)
self.pushButton_play.setText("暂停")
self.pushButton_catch.setEnabled(True)
else:
self.camera_thread.__del__()
self.camera_thread.terminate()
self.camera_thread.isPause = True
self.pushButton_open_close.setText("打开相机")
self.statusbar.showMessage("【{}】 已关闭相机".format(time.strftime("%Y-%m-%d %H:%M:%S")))
self.label_display.setText("相机已关闭")
self.pushButton_play.setEnabled(False)
self.pushButton_play.setText("播放")
self.pushButton_catch.setEnabled(False)
def camera_play_pause(self):
"""
相机暂停和继续
:return:
"""
if self.camera_thread.isRunning():
if self.pushButton_play.text() == "暂停":
mvsdk.CameraPause(self.camera_thread.mCamera)
self.camera_thread.isPause = True
self.statusbar.showMessage("【{}】 实时显示已暂停".format(time.strftime("%Y-%m-%d %H:%M:%S")))
self.pushButton_play.setText("播放")
else:
mvsdk.CameraPlay(self.camera_thread.mCamera)
self.camera_thread.isPause = False
self.statusbar.showMessage("【{}】 实时显示已恢复".format(time.strftime("%Y-%m-%d %H:%M:%S")))
self.pushButton_play.setText("暂停")
def catchPicture(self):
"""
抓取图片
:return:
"""
if self.camera_thread.isRunning():
try:
mvsdk.CameraSetMirror(self.camera_thread.mCamera, 1, True)
pFrameBuffer = mvsdk.CameraAlignMalloc(self.camera_thread.FrameBufferSize, 16)
pRawData, FrameHead = mvsdk.CameraGetImageBuffer(self.camera_thread.mCamera, 1000)
mvsdk.CameraImageProcess(self.camera_thread.mCamera, pRawData, pFrameBuffer, FrameHead)
mvsdk.CameraReleaseImageBuffer(self.camera_thread.mCamera, pRawData)
mvsdk.CameraSetMirror(self.camera_thread.mCamera, 1, False)
# 此时图片已经存储在pFrameBuffer中，对于彩色相机pFrameBuffer=RGB数据，黑白相机pFrameBuffer=8位灰度数据
# 该示例中我们只是把图片保存到硬盘文件中
if self.yamlData['runMode'] == 'mark':
if not os.path.exists("CatchImage"):
os.mkdir("CatchImage")
dirName = 'CatchImage'
else:
if not os.path.exists("RunImage"):
os.mkdir("RunImage")
dirName = 'RunImage'
catchTime = time.strftime("%Y%m%d%H%M%S")
status = mvsdk.CameraSaveImage(self.camera_thread.mCamera,
"./{}/grab_{}.jpg".format(dirName, catchTime),
pFrameBuffer,
FrameHead, mvsdk.FILE_JPG, 100)
if os.path.exists("./{}/grab_{}.jpg".format(dirName, catchTime)):
# 进行目标检测
results = self.camera_thread.model.predict('./{}/grab_{}.jpg'.format(dirName, catchTime))
annotated_frame = results[0].plot()
# 保存处理后带标签的图片
cv2.imwrite("./{}/process_{}.jpg".format(dirName, catchTime), annotated_frame)
# 将图像数据转换为QImage格式
height, width, channel = annotated_frame.shape
bytes_per_line = 3 * width
qimage = QtGui.QImage(annotated_frame.data, width, height, bytes_per_line,
QtGui.QImage.Format_RGB888)
# 将QImage转换为QPixmap
pixmap = QtGui.QPixmap.fromImage(qimage)
self.label_show.setPixmap(pixmap)
self.label_show.setScaledContents(True)
# 提取检测结果
self.textBrowser.append(
"*********************************{}*************************************".format(
time.strftime("%Y-%m-%d %H:%M:%S")))
find_labels = []
for result in results:
boxes = result.boxes.xyxy # 边界框坐标
scores = result.boxes.conf # 置信度分数
classes = result.boxes.cls # 类别索引
# 如果有类别名称，可以通过类别索引获取
class_names = [self.camera_thread.model.names[int(cls)] for cls in classes]
# 打印检测结果
for box, score, class_name in zip(boxes, scores, class_names):
self.textBrowser.append(f"检测标签: {class_name}, 置信度: {score:.2f}")
if score >= self.yamlData['threshold']:
find_labels.append(class_name)
if len(find_labels) == 2 and 'pin-left-pass' in find_labels and 'pin-right-pass' in find_labels:
self.label_result.setText("PASS")
self.label_result.setStyleSheet("background-color: #55ff7f")
# TODO 给出PASS信号
else:
self.label_result.setText("FAIL")
self.label_result.setStyleSheet("background-color: #d80000")
# TODO 给出FAIL信号
self.textBrowser.append(
"*****************************************************************************************")
except mvsdk.CameraException as e:
QMessageBox.critical(self, "错误", "抓图失败({}): {}".format(e.error_code, e.message), QMessageBox.Yes,
QMessageBox.Yes)
def drawLine(self):
"""
添加网格线
:return:
"""
if self.pushButton_addline.text() == "添加网格线":
self.camera_thread.drawLine = True
self.pushButton_addline.setText("去除网格线")
else:
self.camera_thread.drawLine = False
self.pushButton_addline.setText("添加网格线")
def showVideo(self, showImage):
"""
显示视频
:param showImage:
:return:
"""
self.label_display.setPixmap(QtGui.QPixmap.fromImage(showImage))
def showStatus(self, info):
"""
显示状态
:param info:
:return:
"""
self.statusbar.showMessage(info)
def selectMode(self, mode):
"""
选择抓图模式
:param mode:
:return:
"""
if mode == 0:
self.yamlTool.update_yaml('config.yaml', 'runMode', 'mark')
else:
self.yamlTool.update_yaml('config.yaml', 'runMode', 'run')
self.yamlData = self.yamlTool.read_yaml('config.yaml')
def thresholdChange(self, value):
"""
置信阈值改变事件
:param value:
:return:
"""
self.yamlTool.update_yaml('config.yaml', 'threshold', value)
self.yamlData = self.yamlTool.read_yaml('config.yaml')
if __name__ == '__main__':
app = QtWidgets.QApplication(sys.argv)
main_window = MainWindow()
main_window.show()
sys.exit(app.exec_())
动态效果展示
yolov8视觉框架动态检测产品缺陷
