News
应广大读者需求，重构了整个仓库，目前适配YOLOv5最新版本。
开源地址：https://github.com/zstar1003/yolov5_pyqt5
最新界面：
目前支持图像/视频/摄像头检测，适配YOLOv5各版本模型。
; 前言
本篇主要利用PyQT5搭建YOLOv5可视化界面，并打包成exe程序。
整体框架参考自：https://xugaoxiang.com/2021/06/30/yolov5-pyqt5
在此基础上，优化了预测逻辑，适配YOLOv5-5.0版本，并使用 qdarkstyle美化了界面，支持图片检测、摄像头检测、视频检测，整体效果如下图所示：
开源仓库：https://github.com/zstar1003/yolov5_pyqt5
可直接运行的exe程序：https://pan.baidu.com/s/16nHvS5tRSeLKB0Ql2-6ZFw?pwd=8888
; 整体框架
项目整体框架如下图所示：
· models：存放模型构建相关程序，直接从yolov5-5.0版本中clone过来
- utils：存放绘图、数据加载等相关工具，直接从yolov5-5.0版本中clone过来
- UI：存放软件图标
- result：存放预测之后的图片或视频
- weights：模型权重，默认使用YOLOv5官方提供的yolov5s.pt
核心代码
main.py
import os
import sys
import cv2
import random
import torch
import numpy as np
import torch.backends.cudnn as cudnn
import qdarkstyle
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtGui import QIcon, QPixmap
from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.datasets import letterbox
from utils.plots import plot_one_box
class Ui_MainWindow(QtWidgets.QMainWindow):
def __init__(self, parent=None):
super(Ui_MainWindow, self).__init__(parent)
self.timer_video = QtCore.QTimer()
self.setupUi(self)
self.init_logo()
self.init_slots()
self.cap = cv2.VideoCapture()
self.out = None
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.half = self.device.type != 'cpu'
cudnn.benchmark = True
weights = 'weights/yolov5s.pt'
imgsz = 640
self.conf_thres = 0.25
self.iou_thres = 0.45
self.model = attempt_load(weights, map_location=self.device)
stride = int(self.model.stride.max())
self.imgsz = check_img_size(imgsz, s=stride)
if self.half:
self.model.half()
self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]
def setupUi(self, MainWindow):
MainWindow.setObjectName("MainWindow")
MainWindow.resize(900, 600)
self.centralwidget = QtWidgets.QWidget(MainWindow)
self.centralwidget.setObjectName("centralwidget")
self.horizontalLayout_2 = QtWidgets.QHBoxLayout(self.centralwidget)
self.horizontalLayout_2.setObjectName("horizontalLayout_2")
self.horizontalLayout = QtWidgets.QHBoxLayout()
self.horizontalLayout.setSizeConstraint(QtWidgets.QLayout.SetNoConstraint)
self.horizontalLayout.setObjectName("horizontalLayout")
self.verticalLayout = QtWidgets.QVBoxLayout()
self.verticalLayout.setContentsMargins(0, 0, 0, 0)
self.verticalLayout.setObjectName("verticalLayout")
self.pushButton_img = QtWidgets.QPushButton(self.centralwidget)
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.MinimumExpanding)
sizePolicy.setHorizontalStretch(0)
sizePolicy.setVerticalStretch(0)
sizePolicy.setHeightForWidth(self.pushButton_img.sizePolicy().hasHeightForWidth())
self.pushButton_img.setSizePolicy(sizePolicy)
self.pushButton_img.setMinimumSize(QtCore.QSize(150, 40))
self.pushButton_img.setMaximumSize(QtCore.QSize(150, 40))
font = QtGui.QFont()
font.setFamily("Agency FB")
font.setPointSize(12)
self.pushButton_img.setFont(font)
self.pushButton_img.setObjectName("pushButton_img")
self.verticalLayout.addWidget(self.pushButton_img, 0, QtCore.Qt.AlignHCenter)
self.verticalLayout.addStretch(5)
self.pushButton_camera = QtWidgets.QPushButton(self.centralwidget)
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
sizePolicy.setHorizontalStretch(0)
sizePolicy.setVerticalStretch(0)
sizePolicy.setHeightForWidth(self.pushButton_camera.sizePolicy().hasHeightForWidth())
self.pushButton_camera.setSizePolicy(sizePolicy)
self.pushButton_camera.setMinimumSize(QtCore.QSize(150, 40))
self.pushButton_camera.setMaximumSize(QtCore.QSize(150, 40))
self.pushButton_camera.setFont(font)
self.pushButton_camera.setObjectName("pushButton_camera")
self.verticalLayout.addWidget(self.pushButton_camera, 0, QtCore.Qt.AlignHCenter)
self.verticalLayout.addStretch(5)
self.pushButton_video = QtWidgets.QPushButton(self.centralwidget)
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
sizePolicy.setHorizontalStretch(0)
sizePolicy.setVerticalStretch(0)
sizePolicy.setHeightForWidth(self.pushButton_video.sizePolicy().hasHeightForWidth())
self.pushButton_video.setSizePolicy(sizePolicy)
self.pushButton_video.setMinimumSize(QtCore.QSize(150, 40))
self.pushButton_video.setMaximumSize(QtCore.QSize(150, 40))
self.pushButton_video.setFont(font)
self.pushButton_video.setObjectName("pushButton_video")
self.verticalLayout.addWidget(self.pushButton_video, 0, QtCore.Qt.AlignHCenter)
self.verticalLayout.addStretch(50)
self.pushButton_showdir = QtWidgets.QPushButton(self.centralwidget)
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
sizePolicy.setHorizontalStretch(0)
sizePolicy.setVerticalStretch(0)
sizePolicy.setHeightForWidth(self.pushButton_showdir.sizePolicy().hasHeightForWidth())
self.pushButton_showdir.setSizePolicy(sizePolicy)
self.pushButton_showdir.setMinimumSize(QtCore.QSize(150, 50))
self.pushButton_showdir.setMaximumSize(QtCore.QSize(150, 50))
self.pushButton_showdir.setFont(font)
self.pushButton_showdir.setObjectName("pushButton_showdir")
self.verticalLayout.addWidget(self.pushButton_showdir, 0, QtCore.Qt.AlignHCenter)
self.verticalLayout.setStretch(2, 1)
self.horizontalLayout.addLayout(self.verticalLayout)
self.label = QtWidgets.QLabel(self.centralwidget)
self.label.setObjectName("label")
self.horizontalLayout.addWidget(self.label)
self.horizontalLayout.setStretch(0, 1)
self.horizontalLayout.setStretch(1, 3)
self.horizontalLayout_2.addLayout(self.horizontalLayout)
self.label.setStyleSheet("border: 1px solid white;")
MainWindow.setCentralWidget(self.centralwidget)
self.menubar = QtWidgets.QMenuBar(MainWindow)
self.menubar.setGeometry(QtCore.QRect(0, 0, 800, 23))
self.menubar.setObjectName("menubar")
MainWindow.setMenuBar(self.menubar)
self.statusbar = QtWidgets.QStatusBar(MainWindow)
self.statusbar.setObjectName("statusbar")
MainWindow.setStatusBar(self.statusbar)
self.retranslateUi(MainWindow)
QtCore.QMetaObject.connectSlotsByName(MainWindow)
def retranslateUi(self, MainWindow):
_translate = QtCore.QCoreApplication.translate
MainWindow.setWindowTitle(_translate("MainWindow", "YOLOv5目标检测平台"))
self.pushButton_img.setText(_translate("MainWindow", "图片检测"))
self.pushButton_camera.setText(_translate("MainWindow", "摄像头检测"))
self.pushButton_video.setText(_translate("MainWindow", "视频检测"))
self.pushButton_showdir.setText(_translate("MainWindow", "打开输出文件夹"))
self.label.setText(_translate("MainWindow", "TextLabel"))
def init_slots(self):
self.pushButton_img.clicked.connect(self.button_image_open)
self.pushButton_video.clicked.connect(self.button_video_open)
self.pushButton_camera.clicked.connect(self.button_camera_open)
self.pushButton_showdir.clicked.connect(self.button_show_dir)
self.timer_video.timeout.connect(self.show_video_frame)
def init_logo(self):
pix = QtGui.QPixmap('')
self.label.setScaledContents(True)
self.label.setPixmap(pix)
def button_image_open(self):
print('打开图片')
name_list = []
img_name, _ = QtWidgets.QFileDialog.getOpenFileName(
self, "打开图片", "", "*.jpg;;*.png;;All Files(*)")
if not img_name:
return
img = cv2.imread(img_name)
print(img_name)
showimg = img
with torch.no_grad():
img = letterbox(img, new_shape=self.imgsz)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half() if self.half else img.float()
img /= 255.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = self.model(img)[0]
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
for i, det in enumerate(pred):
if det is not None and len(det):
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], showimg.shape).round()
for *xyxy, conf, cls in reversed(det):
label = '%s %.2f' % (self.names[int(cls)], conf)
name_list.append(self.names[int(cls)])
plot_one_box(xyxy, showimg, label=label,
color=self.colors[int(cls)], line_thickness=2)
cv2.imwrite('result/prediction.jpg', showimg)
self.result = cv2.cvtColor(showimg, cv2.COLOR_BGR2BGRA)
self.result = cv2.resize(self.result, (640, 480), interpolation=cv2.INTER_AREA)
self.QtImg = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0], QtGui.QImage.Format_RGB32)
self.label.setPixmap(QtGui.QPixmap.fromImage(self.QtImg))
def button_video_open(self):
video_name, _ = QtWidgets.QFileDialog.getOpenFileName(
self, "打开视频", "", "*.mp4;;*.avi;;All Files(*)")
if not video_name:
return
flag = self.cap.open(video_name)
if flag == False:
QtWidgets.QMessageBox.warning(
self, u"Warning", u"打开视频失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
else:
self.out = cv2.VideoWriter('result/vedio_prediction.avi', cv2.VideoWriter_fourcc(
*'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
self.timer_video.start(30)
self.pushButton_video.setDisabled(True)
self.pushButton_img.setDisabled(True)
self.pushButton_camera.setDisabled(True)
def button_camera_open(self):
if not self.timer_video.isActive():
flag = self.cap.open(0)
if flag == False:
QtWidgets.QMessageBox.warning(
self, u"Warning", u"打开摄像头失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
else:
self.out = cv2.VideoWriter('result/camera_prediction.avi', cv2.VideoWriter_fourcc(
*'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
self.timer_video.start(30)
self.pushButton_video.setDisabled(True)
self.pushButton_img.setDisabled(True)
self.pushButton_camera.setText(u"关闭摄像头")
else:
self.timer_video.stop()
self.cap.release()
self.out.release()
self.label.clear()
self.init_logo()
self.pushButton_video.setDisabled(False)
self.pushButton_img.setDisabled(False)
self.pushButton_camera.setText(u"摄像头检测")
def show_video_frame(self):
name_list = []
flag, img = self.cap.read()
if img is not None:
showimg = img
with torch.no_grad():
img = letterbox(img, new_shape=self.imgsz)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half() if self.half else img.float()
img /= 255.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = self.model(img)[0]
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres)
for i, det in enumerate(pred):
if det is not None and len(det):
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], showimg.shape).round()
for *xyxy, conf, cls in reversed(det):
label = '%s %.2f' % (self.names[int(cls)], conf)
name_list.append(self.names[int(cls)])
plot_one_box(
xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)
self.out.write(showimg)
show = cv2.resize(showimg, (640, 480))
self.result = cv2.cvtColor(show, cv2.COLOR_BGR2RGB)
showImage = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0],
QtGui.QImage.Format_RGB888)
self.label.setPixmap(QtGui.QPixmap.fromImage(showImage))
else:
self.timer_video.stop()
self.cap.release()
self.out.release()
self.label.clear()
self.pushButton_video.setDisabled(False)
self.pushButton_img.setDisabled(False)
self.pushButton_camera.setDisabled(False)
self.init_logo()
def button_show_dir(self):
path = os.getcwd() + '\\' + 'result'
os.system(f"start explorer {path}")
if __name__ == '__main__':
app = QtWidgets.QApplication(sys.argv)
app.setStyleSheet(qdarkstyle.load_stylesheet_pyqt5())
ui = Ui_MainWindow()
icon = QIcon()
icon.addPixmap(QPixmap("./UI/icon.ico"), QIcon.Normal, QIcon.Off)
ui.setWindowIcon(icon)
ui.show()
sys.exit(app.exec_())
整体逻辑是软件已启动就开始载入模型，然后利用 槽函数去响应按钮信息。
打包exe
为了尽可能减少打包之后的体积，在打包之前，先使用Anaconda新建一个虚拟环境并安装好pytorch等YOLOv5所需必要库。
打包通常采用的是Pyinstaller这个工具库，本次打包使用一个新的工具叫 Auto Py to Exe，该工具仍是调用Pyinstaller进行打包，不过对选项进行了可视化，操作更加便捷。
安装方式：
git clone https://github.com/brentvollebregt/auto-py-to-exe.git
python setup.py install
注意安装时可能会提示缺少一些包，依次pip安装即可， geventwebsocket库需要这样进行安装。
pip install gevent-websocket
安装好之后，在终端输入auto-py-to-exe，会在浏览器中默认打开如下界面：
脚本位置选择 main.py，选择单目录模式，隐藏控制台，并选择图标和输出路径，然后就可以一键进行打包。
打包完成之后，会在输出文件夹下输入一个 main文件夹。
运行之前，需要将原始工程中的几个文件夹拷贝进去，否则会提示找不到文件，如下图所示：
双击
main.exe，即可看到可视化界面。
报错解决
在调式时，遇到一些小问题，这里也记录下。
问题一：遇到警告：
UserWarning: torch.meshgrid: in an upcoming release, it will be required to …
在报错的文件中将
return _VF.meshgrid(tensors, **kwargs)
修改为
return _VF.meshgrid(tensors, **kwargs, indexing = 'ij')
问题二：
打包时遇到的错误：
ImportError: ERROR: recursion is detected during loading of “cv2” binary extensions. Check OpenCV installation.
pyinstaller和cv2版本存在兼容问题，卸载已有的opencv-python，安装opencv-python=4.5.3.56
Original: https://blog.csdn.net/qq1198768105/article/details/126463741
Author: zstar-_
Title: 【目标检测】利用PyQT5搭建YOLOv5可视化界面
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/727217/
转载文章受原作者版权保护。转载请注明原作者出处！