参考:b站up:白月黑雨编程
https://www.bilibili.com/video/BV1Wj421d78y/?p=3&spm_id_from=333.880.my_history.page.click
改进:增加图像检测,增加错误弹窗
2024.4.23
改进2:
增加控制台显示
解决自定义权重无法绘制问题:增加draw_detection和draw_detections函数,且用apply()代替itrrows改善速度
全代码:
#
################################################################################
## Form generated from reading UI file 'WindowsTest.ui'
##
## Created by: Qt User Interface Compiler version 6.6.3
##
## WARNING! All changes made in this file will be lost when recompiling UI file!
################################################################################
#
from PySide6 import QtWidgets, QtCore, QtGui
from PySide6.QtCore import (QCoreApplication, QDate, QDateTime, QLocale,
QMetaObject, QObject, QPoint, QRect,
QSize, QTime, QUrl, Qt
# 控制台
,QEventLoop,QTimer
)
from PySide6.QtGui import (QBrush, QColor, QConicalGradient, QCursor,
QFont, QFontDatabase, QGradient, QIcon,
QImage, QKeySequence, QLinearGradient, QPainter,
QPalette, QPixmap, QRadialGradient, QTransform)
# 窗口控件
from PySide6.QtWidgets import (QApplication, QFrame, QLabel, QMainWindow,
QSizePolicy, QStatusBar, QToolButton, QWidget
,QMessageBox,QPushButton
,QTextBrowser,QGroupBox
)
#
# from PySide6.QtWidgets import QApplication, QMessageBox,QPushButton
# Thread类允许你创建并启动新的线程,以便同时执行多个任务。
from threading import Thread
import cv2
import time
# 系统模块
import sys
# 模型检测绘制
# 官方yolov5s权重,可以使用官方库进行绘制
# YOLO() m模型导出函数
# result.plot()绘制结果
from ultralytics import YOLO
# 自定义模型权重:使用pytorch
import torch
from PIL import Image
# 控制台输出 重定向信号
class EmittingStr(QtCore.QObject):
textWritten = QtCore.Signal(str) # 定义一个发送str的信号,这里用的方法名与PyQt5不一样
def write(self, text):
self.textWritten.emit(str(text))
loop = QEventLoop()
QTimer.singleShot(1000, loop.quit)
loop.exec()
# class Ui_MainWindow(object):
class MWindow(QtWidgets.QMainWindow):
# 初始化操作
def __init__(self):
# 调用父类的构造函数。以完成初始化
super().__init__()
# 设置界面
self.setupUI()
# 控件绑定
self.BttCamDect.clicked.connect(self.openCam)
self.BttVidDect.clicked.connect(self.openVideoFile)
self.BttImgDect.clicked.connect(self.openImgFile)
self.BttStop.clicked.connect(self.stop)
# 加载 YOLOv5s模型(初始化)
# self.model = YOLO("yolov5s.pt")
self.model = torch.hub.load("./", "custom" ,path= "runs/train/exp15/weights/best.pt" ,source= "local")
# self.model = torch.hub.load ("./", "custom" ,path= "./yolov5s.pt" ,source= "local")
sys.stdout = EmittingStr()
self.TBCon.connect(sys.stdout, QtCore.SIGNAL("textWritten(QString)"), self.outputWritten)
sys.stderr = EmittingStr()
self.TBCon.connect(sys.stderr, QtCore.SIGNAL("textWritten(QString)"), self.outputWritten)
# 视频处理
self.timer_camera = QtCore.QTimer()
# 定时到了,回调 self.show_camera
# 定义定时器,用于控制显示视频的帧率
self.timer_videoFile = QtCore.QTimer()
# 定时到了,回调 self.show_camera
self.timer_videoFile.timeout.connect(self.showVideoFile)
# 要处理的视频帧图片队列,目前就放1帧图片。
self.frameToAnalyze = []
# 启动处理视频帧独立线程
Thread(target=self.frameAnalyzeThreadFunc, daemon=True).start()
self.timer_videoFile = QtCore.QTimer()
# 定时到了,回调 self.show_camera
self.timer_videoFile.timeout.connect(self.showVideoFile)
# 当前要播放的视频帧号
self.vframeIdx = 0
# cv2.VideoCapture 实例
self.cap = None
self.stopFlag = False
self.imgPath = ""
self. path = ""
# 界面设置
def setupUI(self):
# if not MainWindow.objectName():
# MainWindow.setObjectName(u"MainWindow")
# MainWindow.resize(800, 600)
centralWidget = QtWidgets.QWidget(self)
self.setCentralWidget(centralWidget)
self.setWindowTitle('YOLO-Qt 演示')
centralWidget = QtWidgets.QWidget(self)
self.setCentralWidget(centralWidget)
# # central Widget 里面的 主 layout:垂直布布局
# # HLayout:水平布局
# # VLayout:垂直布局
# mainLayout = QtWidgets.QVBoxLayout(centralWidget)
# HLayout = QtWidgets.QHBoxLayout()
# VLayout = QtWidgets.QVBoxLayout()
# Label
# 布局方式:坐标法
# 原始数据
self.LabOri = QtWidgets.QLabel('原始图片' ,self ) # 创建控件
self.LabOri.setObjectName(u"LabWild" ) # 控件命名
self.LabOri.setAlignment(Qt.AlignCenter ) # 字体对齐
self.LabOri.setStyleSheet('border:1px solid #D7E2F9;' ) # 边框
# 检测数据
self.LabDect = QtWidgets.QLabel('检测控件' ,self ) # 创建控件
self.LabDect.setObjectName(u"LabDect" ) # 控件命名
self.LabDect.setAlignment(Qt.AlignCenter ) # 字体对齐
self.LabDect.setStyleSheet('border:1px solid #D7E2F9;' ) # 边框
# Button
# 图像检测
self.BttImgDect = QtWidgets.QPushButton('🖼图片文件' ,self)
self.BttImgDect.setObjectName(u"BttPicDect")
# 录像检测
self.BttVidDect = QtWidgets.QPushButton('🎞️视频文件' ,self)
self.BttVidDect.setObjectName(u"BttVidDect")
# 摄像头
self.BttCamDect = QtWidgets.QPushButton('📹摄像头' ,self)
self.BttCamDect.setObjectName(u"BttCamDect")
# 停止
self.BttStop = QtWidgets.QPushButton('🛑停止' ,self)
self.BttStop.setObjectName(u"BttCamDect")
# 警告弹窗
self.MsgWarn = QMessageBox()
self.MsgWarn.setIcon(QMessageBox.Warning)
self.MsgWarn.setStandardButtons(QMessageBox.Ok)
# self.MsgWarn.setDefaultButton(QMessageBox.Ok)
self.MsgWarn.setText("error")
self.error_occurred = False
# 打印台控件
self.TBCon = QTextBrowser(self)
self.TBCon.setObjectName(u"textBrowser")
self.GBCon = QGroupBox(self)
self.GBCon.setObjectName(u"groupBox")
# 布局方法:坐标法
self.resize(1300, 750)
self.LabOri.setGeometry(QRect(10, 70, 640, 400) ) # 设置大小
self.LabDect.setGeometry(QRect(650, 70, 640, 400) ) # 设置大小,位置
self.BttImgDect.setGeometry(QRect(1120, 490, 160, 40))
self.BttVidDect.setGeometry(QRect(1120, 560, 160, 40))
self.BttCamDect.setGeometry(QRect(1120, 630, 160, 40))
self.BttStop.setGeometry(QRect(1120, 700, 160, 40))
self.TBCon.setGeometry(QRect(30, 520, 1000,200))
self.GBCon.setGeometry(QRect(20, 490, 1020, 230))
# MainWindow.setCentralWidget(self)
# self.statusbar = QStatusBar(MainWindow)
# self.statusbar.setObjectName(u"statusbar")
# MainWindow.setStatusBar(self.statusbar)
# self.retranslateUi(MainWindow)
# QMetaObject.connectSlotsByName(MainWindow)
# setupUi
# def retranslateUi(self, MainWindow):
# def retranslateUi(self):
# #MainWindow.setWindowTitle(QCoreApplication.translate("MainWindow", u"MainWindow", None))
# self.LabWild.setText(QCoreApplication.translate("MainWindow", u"\u539f\u59cb\u753b\u9762", None))
# self.LabDect.setText(QCoreApplication.translate("MainWindow", u"\u68c0\u6d4b\u753b\u9762", None))
# self.BttPicDect.setText(QCoreApplication.translate("MainWindow", u"\u56fe\u7247\u8bc6\u522b", None))
# self.BttVidDect.setText(QCoreApplication.translate("MainWindow", u"\u89c6\u9891\u8bc6\u522b", None))
# self.BttCamDect.setText(QCoreApplication.translate("MainWindow", u"\u6444\u50cf\u5934\u8bc6\u522b", None))
# # retranslateUi
# 打开照片文件
def outputWritten(self, text):
# self.TBCon.clear()
cursor = self.TBCon.textCursor()
cursor.movePosition(QtGui.QTextCursor.End)
cursor.insertText(text)
self.TBCon.setTextCursor(cursor)
self.TBCon.ensureCursorVisible()
def openImgFile(self):
# 先关闭原来打开的
self.stop()
self.imgPath, _ = QtWidgets.QFileDialog.getOpenFileName(
self, # 父窗口对象
"选择图片文件", # 标题
dir="./Test/testvideo_img/", # 起始目录
filter="图片类型 (*.jpg *.png *jpeg)" # 选择类型过滤项,过滤内容在括号中
)
print('imgPath is', self.imgPath)
if not self.imgPath:
self.MsgWarn.setText("Not Image Path")
self.error_occurred = True
self.showErrorDialog(self.error_occurred)
return
else:
image = cv2.imread(self.imgPath)
self.setImgToOriLabel(image)
# # imgPath = imgPath[0]
# # pixmap = QPixmap(self.imgPath)
# # pixmap = pixmap.scaled(640,400)
# # self.LabOri.setPixmap(pixmap)
# self.LabOri.setPixmap((QPixmap(imgPath).scaled(640,400)))
# # qImage = self.ImgAnalyze(imgPath)
# self.ImgAnalyze()
# #self.LabDect.setPixmap(QPixmap.fromImage(self.ImgAnalyzeThreadFunc(self.imgPath)))
def setImgToOriLabel(self,image):
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
qImage = QtGui.QImage(image.data, image.shape[1], image.shape[0],
QtGui.QImage.Format_RGB888) # 变成QImage形式
self.LabOri.setPixmap((QtGui.QPixmap.fromImage(qImage).scaled(640, 400)))
self.ImgToAnalyze(image)
def ImgToAnalyze(self,image):
#detection
results = self.model(image)
print(results)
results.show(image)
# self.setImgToDectLabel(results)
# self.JIANXIU(results,image)
self.setImgToDectLabel(results,image)
# def JIANXIU(self,results,image):
# # results
# # image 1/1: 540x960 13 cars, 1 van
# # Speed: 4.0ms pre-process, 119.0ms inference, 0.0ms NMS per image at shape (1, 3, 384, 640)
# # <class 'models.common.Detections'>
# # print(results)
# # results.print(image)
# boxes = results.xyxy[0].cpu().numpy() # 获取边界框坐标信息,xyxy 格式表示(left, top, right, bottom)
# detections = results.pandas().xyxy[0]
# self.draw_detections(detections, image, boxes)
# image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# qImage = QtGui.QImage(image_rgb.data, image.shape[1], image.shape[0],
# QtGui.QImage.Format_RGB888) # 变成QImage形式
# self.LabDect.setPixmap((QtGui.QPixmap.fromImage(qImage).scaled(640,400))) # 往显示Label里 显示QImage
#
#
def setImgToDectLabel(self,results,image):
# boxes = results.xyxy[0].cpu().numpy() # 获取边界框坐标信息,xyxy 格式表示(left, top, right, bottom)
detections = results.pandas().xyxy[0]
self.draw_detections(detections, image)
# image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
qImage = QtGui.QImage(image.data, image.shape[1], image.shape[0],
QtGui.QImage.Format_RGB888) # 变成QImage形式
self.LabDect.setPixmap((QtGui.QPixmap.fromImage(qImage).scaled(640,400))) # 往显示Label里 显示QImage
def openVideoFile(self):
# 先关闭原来打开的
self.stop()
videoPath, _ = QtWidgets.QFileDialog.getOpenFileName(
self, # 父窗口对象
"选择视频文件", # 标题
".", # 起始目录
"图片类型 (*.mp4 *.avi)" # 选择类型过滤项,过滤内容在括号中
)
print('videoPath is', videoPath)
if not videoPath:
self.MsgWarn.setText("Not Video Path")
self.error_occurred = True
self.showErrorDialog(self.error_occurred)
return
self.cap = cv2.VideoCapture(videoPath)
if not self.cap.isOpened(): # 确保视频文件被正确打开
print("打开文件失败")
self.MsgWarn.setText("Fail to open Video")
self.error_occurred = True
self.showErrorDialog(self.error_occurred)
return
self.timer_videoFile.start(30)
self.stopFlag = False
print("ok")
def showVideoFile(self):
if self.cap: # 确保self.cap存在
# 选取视频帧位置,
self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.vframeIdx)
self.vframeIdx += 1
ret, frame = self.cap.read() # 从视频流中读取
if ret:
frame = cv2.resize(frame, (640, 400))
else:
self.MsgWarn.setText("Not Video")
self.error_occurred = True
self.showErrorDialog(self.error_occurred)
return
self.setFrameToOriLabel(frame)
def openCam(self):
self.stop()
self.cap = cv2.VideoCapture(1, cv2.CAP_DSHOW)
if not self.cap.isOpened():
print("1号摄像头不能打开")
self.MsgWarn.setText("Cam error")
self.error_occurred = True
self.showErrorDialog(self.error_occurred)
return ()
else:
print("ok")
if not self.timer_camera.isActive(): # 若定时器未启动
#看相机具体参数
self.timer_camera.start(60)
def showCam(self):
print("showCam")
ret, frame = self.cap.read() # 从视频流中读取
if not ret:
return
self.setFrameToOriLabel(frame)
# 如果当前没有处理任务 开始处理
if not self.frameToAnalyze:
self.frameToAnalyze.append(frame)
def setFrameToOriLabel(self ,frame):
# 视频色彩转换回RGB,OpenCV images as BGR
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
qImage = QtGui.QImage(frame.data, frame.shape[1], frame.shape[0],
QtGui.QImage.Format_RGB888) # 变成QImage形式
# 往显示视频的Label里 显示QImage
self.LabOri.setPixmap(QtGui.QPixmap.fromImage(qImage))
# 如果当前没有处理任务
if not self.frameToAnalyze:
self.frameToAnalyze.append(frame)
def frameAnalyzeThreadFunc(self):
while True:
if not self.frameToAnalyze:
time.sleep(0.01)
continue
frame = self.frameToAnalyze.pop(0)
# results = self.model(frame)[0]
results = self.model(frame)
# img = results.plot(line_width=1)
# boxes = results.xyxy[0].cpu().numpy() # 获取边界框坐标信息,xyxy 格式表示(left, top, right, bottom)
detections = results.pandas().xyxy[0]
self.draw_detections(detections, frame)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
qImage = QtGui.QImage(frame.data, frame.shape[1],frame.shape[0],
QtGui.QImage.Format_RGB888) # 变成QImage形式
if self.stopFlag == False:
self.LabDect.setPixmap(QtGui.QPixmap.fromImage(qImage)) # 往显示Label里 显示QImage
# time.sleep(0.5)
time.sleep(0.05)
# 绘制检测结果
def draw_detection(self, row, image):
x1, y1, x2, y2, conf, cls = row[:6] # 获取边界框坐标和类别
cls_tran = {
0: ("cars",(0,0,255)),
1: ("bus",(0,255,0)),
2: ("van",(255,0,0)),
}
cls,Anch_color = cls_tran.get(cls, ("others",(255,255,255)))
line_width = 1.5 # 非整数线条宽度
num_rectangles = 2 # 需要绘制的矩形数量
for i in range(num_rectangles):
# 每个矩形的左上角和右下角坐标
offset = i - (line_width - 1) / 2
new_x1 = x1 - offset
new_y1 = y1 - offset
new_x2 = x2 + offset
new_y2 = y2 + offset
image = cv2.rectangle(image, (int(new_x1), int(new_y1)), (int(new_x2), int(new_y2)), Anch_color, 2) # 绘制矩形
image = cv2.putText(image, f'{cls} {conf:.2f}', (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
Anch_color, 2) # 绘制类别和置信度
return image
def draw_detections(self, detections, image):
print(detections)
# itrrows 过慢 -》apply
# for _,det in detections.iterrows(): # 假设检测结果在pred属性中,需要根据实际情况调整
# x1, y1, x2, y2, conf, cls = det[:6] # 获取边界框坐标和类别
# cls_tran = {
# 0: "cars",
# 1: "bus",
# 2: "van"
# }
# cls = cls_tran.get(cls,"others")
# line_width = 1.5 # 非整数线条宽度
# num_rectangles = 8 # 需要绘制的矩形数量
#
# for i in range(num_rectangles):
# # 每个矩形的左上角和右下角坐标
# offset = i - (line_width - 1) / 2
# new_x1 = x1 - offset
# new_y1 = y1 - offset
# new_x2 = x2 + offset
# new_y2 = y2 + offset
# image = cv2.rectangle(image, (int(new_x1), int(new_y1)), (int(new_x2), int(new_y2)), (0, 255, 0), 2) # 绘制矩形
# image = cv2.putText(image, f'{cls} {conf:.2f}', (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
# (36, 255, 12), 2) # 绘制类别和置信度
processed_image = detections.apply(lambda row: self.draw_detection(row, image), axis=1)
return processed_image
def set(self):
pass
def stop(self):
self.timer_camera.stop() # 关闭定时器
if self.cap is not None: # 检查self.cap是否为None,防止release出错
self.cap.release() # 释放视频流
self.LabDect.clear() # 清空视频显示区域
self.LabOri.clear() # 清空视频显示区域
self.TBCon.clear()
def showErrorDialog(self ,condition): # 错误弹窗
# 停止弹窗
response = self.MsgWarn.exec()
if response == QMessageBox.Ok:
self.stopFlag = True # 让 frameAnalyzeThreadFunc 不要再设置 labdect
self.timer_camera.stop() # 关闭定时器
self.timer_videoFile.stop() # 关闭定时器
if self.cap is not None:
self.stop()
self.MsgWarn.setText("error")
self.error_occurred = False
# #使用
# self.MsgWarn.setText("(reason)")
# self.error_occurred = True
# self.showErrorDialog(self.error_occurred)
# 创建一个 PyQt 程序,并显示一个窗口,然后运行应用程序的事件循环
# 创建一个应用程序对象 app。
app = QtWidgets.QApplication()
# 创建了一个窗口对象
window = MWindow()
# 显示窗口
window.show()
# 启动应用程序的事件循环
app.exec()