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过来
核心代码
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 = None self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.half = self.device.type != 'cpu' # half precision only supported on CUDA cudnn.benchmark = True weights = '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.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] # Convert # BGR to RGB, to 3x416x416 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() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference pred = self.model(img)[0] # Apply NMS pred = non_max_suppression(pred, self.conf_thres, self.iou_thres) # Process detections for i, det in enumerate(pred): if det is not None and len(det): # Rescale boxes from img_size to im0 size 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) # 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: 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(): # 默认使用第一个本地camera 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] # Convert # BGR to RGB, to 3x416x416 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() # uint8 to fp16/32 img /= 255.0 # 0 - 255 to 0.0 - 1.0 if img.ndimension() == 3: img = img.unsqueeze(0) # Inference pred = self.model(img)[0] # Apply NMS pred = non_max_suppression(pred, self.conf_thres, self.iou_thres) # Process detections for i, det in enumerate(pred): # detections per image if det is not None and len(det): # Rescale boxes from img_size to im0 size det[:, :4] = scale_coords( img.shape[2:], det[:, :4], showimg.shape).round() # Write results for *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