YOLOv8应用于产品缺陷检测实例
- 项目概况
- 项目实现
- YOLOv8安装及模型训练
- 关键代码展示
- 动态效果展示
项目概况
本项目是应用YOLOv8框架实现训练自定义模型实现单一零件的缺陷检测,软件界面由PyQt5实现。
功能已正式使用,识别效果达到预期。
项目实现
项目使用了以下几个要素:
- 全新的界面设计 ,PyQt5结合QtDesigner自定义界面设计,快速构建想要的UI;
- 相机选型和光源 ,项目使用了迈德威视工业相机GYD-GE130M-T,对照官方给出的文档很容易就可以实现相机的使用,光源方面使用12V常规灯带打光,去除其他光源的干扰;
- IO触发 功能,实现串口控制IO,以快速输出检测识别结果,便于对接工业传送装置;
YOLOv8安装及模型训练
可参考该大佬文章: 传送门
项目展示图:
关键代码展示
迈德威视相机使用线程代码片
.
import timeimport cv2
from ultralytics import YOLO
import numpy as np
from PyQt5 import QtGui
from PyQt5.QtCore import pyqtSignal, QThread
from MindVisionSdk import mvsdk
import warningswarnings.filterwarnings('ignore')class MyThread(QThread):camera_signal = pyqtSignal(object) # 输出相机数据流info_signal = pyqtSignal(str) # 输出信息状态def __init__(self, parent=None):super().__init__(parent)self.drawLine = Falseself.monoCamera = Noneself.FrameBufferSize = Noneself.isPause = Falseself.mCamera = Noneself.pFrameBuffer = Noneself.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 timeimport 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 myYamlToolclass 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 = Falseself.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 = Trueself.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 = Trueself.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 = Falseself.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.shapebytes_per_line = 3 * widthqimage = QtGui.QImage(annotated_frame.data, width, height, bytes_per_line,QtGui.QImage.Format_RGB888)# 将QImage转换为QPixmappixmap = 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 = Trueself.pushButton_addline.setText("去除网格线")else:self.camera_thread.drawLine = Falseself.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视觉框架动态检测产品缺陷