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 osimport sysimport cv2import randomimport torchimport numpy as npimport torch.backends.cudnn as cudnnimport qdarkstylefrom PyQt5 import QtCore, QtGui, QtWidgetsfrom PyQt5.QtGui import QIcon, QPixmapfrom models.experimental import attempt_loadfrom utils.general import check_img_size, non_max_suppression, scale_coordsfrom utils.datasets import letterboxfrom utils.plots import plot_one_boxclass 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 = Noneself.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")self.half = self.device.type != 'cpu'# half precision only supported on CUDAcudnn.benchmark = Trueweights = 'weights/yolov5s.pt' # 模型加载路径imgsz = 640# 预测图尺寸大小self.conf_thres = 0.25# NMS置信度self.iou_thres = 0.45# IOU阈值# 载入模型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()# to FP16# 从模型中获取各类别名称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)# MainWindow.setStyleSheet("")self.centralwidget = QtWidgets.QWidget(MainWindow)self.centralwidget.setObjectName("centralwidget")# self.centralwidget.setStyleSheet("border: 1px solid white;")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.setSpacing(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.translateMainWindow.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:returnimg = cv2.imread(img_name)print(img_name)showimg = imgwith torch.no_grad():img = letterbox(img, new_shape=self.imgsz)[0]# Convert# BGR to RGB, to 3x416x416img = 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()# uint8 to fp16/32img /= 255.0# 0 - 255 to 0.0 - 1.0if img.ndimension() == 3:img = img.unsqueeze(0)# Inferencepred = self.model(img)[0]# Apply NMSpred = non_max_suppression(pred, self.conf_thres, self.iou_thres)# Process detectionsfor i, det in enumerate(pred):if det is not None and len(det):# Rescale boxes from img_size to im0 sizedet[:, :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)# print(label.split()[0])# 打印各目标名称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:returnflag = 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():# 默认使用第一个本地cameraflag = 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 = imgwith torch.no_grad():img = letterbox(img, new_shape=self.imgsz)[0]# Convert# BGR to RGB, to 3x416x416img = 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()# uint8 to fp16/32img /= 255.0# 0 - 255 to 0.0 - 1.0if img.ndimension() == 3:img = img.unsqueeze(0)# Inferencepred = self.model(img)[0]# Apply NMSpred = non_max_suppression(pred, self.conf_thres, self.iou_thres)# Process detectionsfor i, det in enumerate(pred):# detections per imageif det is not None and len(det):# Rescale boxes from img_size to im0 sizedet[:, :4] = scale_coords(img.shape[2:], det[:, :4], showimg.shape).round()# Write resultsfor *xyxy, conf, cls in reversed(det):label = '%s %.2f' % (self.names[int(cls)], conf)name_list.append(self.names[int(cls)])# print(label)# 打印各目标+置信度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()# 设置窗口透明度# ui.setWindowOpacity(0.93)# 去除顶部边框# ui.setWindowFlags(Qt.FramelessWindowHint)# 设置窗口图标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.gitpython 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) # type: ignore[attr-defined]

修改为

return _VF.meshgrid(tensors, **kwargs, indexing = ‘ij’) # type: ignore[attr-defined]

问题二:
打包时遇到的错误:

ImportError: ERROR: recursion is detected during loading of “cv2” binary extensions. Check OpenCV installation.

pyinstaller和cv2版本存在兼容问题,卸载已有的opencv-python,安装opencv-python=4.5.3.56