目标检测利用PyQT5搭建YOLOv5可视化界面
Posted zstar-_
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了目标检测利用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' # 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)以上是关于目标检测利用PyQT5搭建YOLOv5可视化界面的主要内容，如果未能解决你的问题，请参考以下文章
智能火焰与烟雾检测系统（Python+YOLOv5深度学习模型+清新界面）
睿智的目标检测56——Pytorch搭建YoloV5目标检测平台