识别效果视频
Yolov5-5.0行人闯红灯检测+简易GUI
目录结构
本文程序目录结构如图，在Yolov5根目录下创建GUI文件夹用于存储GUI界面相关代码文件。下面左图为源代码根目录，右图为GUI文件目录，在GUI文件夹下只有红色方框框选的文件有用，其他文件是前期摸索的实验性代码。本文将一一列举框选的Py文件。
 |
|
1.Yolov5版本选择
本文Github下载Yolov5-5.0版本，下载网址：GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite
相信诸位来学习的道友应该都知道如何使用Github，小的就不赘述了。
2.PT模型推理（detect_pt.py）
本文将Detect.py文件当作模块调用，删除了Detect中的计时代码、预测结果保存代码摄像头调用代码：保留了置信度与重合度阈值设定和指定预测类别参数class，所以本文代码目前只能用于预测视频。另外为了不让文件每次调用时都加载一次模型，本文将模型加载部分放在GUI主程序里面，将加载后的模型当作参数进行传递。
import logging
import time
import torch.backends.cudnn as cudnn
import cv2
import numpy as np
from pathlib import Path
from utils.general import check_imshow,check_img_size,non_max_suppression,increment_path\
,scale_coords
from utils.plots import plot_one_box
from utils.datasets import LoadStreams,LoadImages
from utils.torch_utils import select_device,time_synchronized,load_classifier
import torch
from models.experimental import attempt_load,Ensemble
from utils.datasets import letterbox
logger = logging.getLogger(__name__)#日志
def pred_image(model,im0s,save_img,imgsz,conf_thres, iou_thres, classes, agnostic_nms,stride,names,color,half,device):
t0 = time.time()#计时
img = letterbox(im0s, imgsz, stride=stride)[0]# 图片调整大小填充
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device)
img = img.half() if 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)
pred = model(img)[0] #输入填充调整后的图片推理
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)
s, img0 = '', im0s
classes_dict = {'person':"0", '红灯':"0", '绿灯':"0", '闯红灯':"0", '黄灯':"0", 'car':"0", 'bicycle':"0", 'motorcycle':"0"}
'''classes_dict是用来统计预测结果中的类别个数的，这里可以修改为诸位模型的预测类别'''
for i, det in enumerate(pred): # 枚举每个图像目标检测结果
s += '%gx%g ' % img.shape[2:] # 打印字符串
if len(det):
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()
# 统计整理检测结果
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # 每个类的检测数
classes_dict[names[int(c)]] = f"{n}"
for *xyxy, conf, cls in reversed(det): #reversed() 函数用于将结果反转，以便按照从高到低的置信度顺序遍历。
# print(*xyxy)
if save_img : # 添加检测框到图片
color_class = color[names[int(cls)]]
label = f'{names[int(cls)]} {conf:.2f}'
img0 = plot_one_box(xyxy, img0, label=label, color=color_class, line_thickness=1)
return img0,classes_dict
def pt_detection(opencv_imread_image,model,classes,device,names,stride,imgsz,color,half):#GUI预测使用
conf_thres, iou_thres = 0.25, 0.45
# classes = [0, 1, 2, 3, 4, 5, 6, 7]
agnostic_nms = False # 进行NMS去除不同类别之间的框, 默认False
project = './Gui_file' # 保存到文件夹
name = "run_detect" # 保存名字
save_txt = False # 是否保存预测txt文件，默认False
exist_ok = True # 是否重新创建日志文件, False时重新创建文件,默认Ture
save_conf = False # 保存置信度到TXT文件
save_img = True # 保存bundingbox到图像
img0,classes_dict=pred_image(model, opencv_imread_image, save_img, imgsz, conf_thres, iou_thres, classes, agnostic_nms, stride, names, color, half,
device)
return img0,classes_dict
3.ONNX模型推理（detect_onnx.py）
本文还进行ONNX模型推理可视化，本部分代码使用了大佬的文章，下面是链接：
https://blog.csdn.net/qq_22487889/article/details/128011883import random
import sys
import time
import matplotlib
import onnx
import onnxruntime
import cv2
import numpy as np
import torch
import torchvision
from PIL import Image, ImageDraw, ImageFont
from matplotlib import pyplot as plt
# Settings
matplotlib.rc('font', **{'size': 11})
matplotlib.use('Agg') # for writing to files only
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
CLASSES = ['person', '红灯', '绿灯', '闯红灯','黄灯','car','bicycle','motorcycle']
'''CLASSES列表用于存储需要进行预测的类别'''
class Yolov5ONNX(object):
def __init__(self, onnxpath):
self.onnx_session = onnxruntime.InferenceSession(onnxpath)
self.input_name = self.get_input_name()
self.output_name = self.get_output_name()
def get_input_name(self):
input_name = []
for node in self.onnx_session.get_inputs():
input_name.append(node.name)
return input_name
def get_output_name(self):
output_name = []
for node in self.onnx_session.get_outputs():
output_name.append(node.name)
return output_name
def get_input_feed(self, img_tensor):
input_feed = {}
for name in self.input_name:
input_feed[name] = img_tensor
return input_feed
def inference(self, img,input_image_size):
if type(img)==str:
img = cv2.imread(img)
else:
img = img
or_img = cv2.resize(img, input_image_size)
img = or_img[:, :, ::-1].transpose(2, 0, 1) # BGR2RGB和HWC2CHW
img = img.astype(dtype=np.float32)
img /= 255.0
img = np.expand_dims(img, axis=0)
input_feed = self.get_input_feed(img)
pred = self.onnx_session.run(None, input_feed)[0]
return pred, or_img
def nms(dets, thresh):
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
scores = dets[:, 4]
keep = []
index = scores.argsort()[::-1]
while index.size > 0:
i = index[0]
keep.append(i)
x11 = np.maximum(x1[i], x1[index[1:]])
y11 = np.maximum(y1[i], y1[index[1:]])
x22 = np.minimum(x2[i], x2[index[1:]])
y22 = np.minimum(y2[i], y2[index[1:]])
w = np.maximum(0, x22 - x11 + 1)
h = np.maximum(0, y22 - y11 + 1)
overlaps = w * h
ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
idx = np.where(ious <= thresh)[0]
index = index[idx + 1]
return keep
def xywh2xyxy(x):
# [x, y, w, h] to [x1, y1, x2, y2]
y = np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2
y[:, 1] = x[:, 1] - x[:, 3] / 2
y[:, 2] = x[:, 0] + x[:, 2] / 2
y[:, 3] = x[:, 1] + x[:, 3] / 2
return y
def filter_box(org_box, conf_thres, iou_thres): # 过滤掉无用的框
org_box = np.squeeze(org_box)
conf = org_box[..., 4] > conf_thres
box = org_box[conf == True]
cls_cinf = box[..., 5:] # 左闭右开（5 6 7 8），就只剩下了每个grid cell中各类别的概率
cls = []
for i in range(len(cls_cinf)):
cls.append(int(np.argmax(cls_cinf[i]))) # 剩下的objecctness score比较大的grid cell，分别对应的预测类别列表
all_cls = list(set(cls)) # 去重，找出图中都有哪些类别
output = []
for i in range(len(all_cls)):
curr_cls = all_cls[i]
curr_cls_box = []
curr_out_box = []
for j in range(len(cls)):
if cls[j] == curr_cls:
box[j][5] = curr_cls
curr_cls_box.append(box[j][:6])
curr_cls_box = np.array(curr_cls_box)
# curr_cls_box_old = np.copy(curr_cls_box)
curr_cls_box = xywh2xyxy(curr_cls_box)
curr_out_box = nms(curr_cls_box, iou_thres)
for k in curr_out_box:
output.append(curr_cls_box[k])
output = np.array(output)
return output
def draw(image, box_data):
boxes = box_data[..., :4].astype(np.int32) # x1 x2 y1 y2
scores = box_data[..., 4]#置信度
classes = box_data[..., 5].astype(np.int32)#classfy
for box, score, cl in zip(boxes, scores, classes):
label = CLASSES[int(cl)]
color = color_bounding_box(label,cl)
plot_one_box(box,image,color=color,label=label,line_thickness=1)
return image
def color_bounding_box(names,cls):
color= {
'person':[255, 255, 255], '红灯':[0, 0, 255], '绿灯':[0, 255, 0],
'闯红灯':[0, 0, 255], '黄灯':[0, 255, 255], 'car':[0, 255, 0],
'bicycle':[0, 255, 0], 'motorcycle':[0, 255, 0]
}
return color[names]
def plot_one_box(x, img, color=None, label=None, line_thickness=3):
tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
if color not in [[0,255,0],[0,0,255],[255,255,255],[0,0,255],[0,255,255]]:
color = color or [random.randint(0, 255) for _ in range(3)]
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
if label:
tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
font_size = t_size[1]
font = ImageFont.truetype(r'.\font_ttf\SimHei.ttf', font_size)
t_size = font.getsize(label)
c2 = c1[0] + t_size[0], c1[1] - t_size[1]
cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
img_PIL = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
draw = ImageDraw.Draw(img_PIL)
draw.text((c1[0], c2[1] - 2), label, fill=(255, 255, 255), font=font)
return cv2.cvtColor(np.array(img_PIL), cv2.COLOR_RGB2BGR)
def box_iou(box1, box2):
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
"""
Return intersection-over-union (Jaccard index) of boxes.
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
Arguments:
box1 (Tensor[N, 4])
box2 (Tensor[M, 4])
Returns:
iou (Tensor[N, M]): the NxM matrix containing the pairwise
IoU values for every element in boxes1 and boxes2
"""
def box_area(box):
# box = 4xn
return (box[2] - box[0]) * (box[3] - box[1])
area1 = box_area(box1.T)
area2 = box_area(box2.T)
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
def detect_image(onnx_path,image):
# onnx_path = 'weights/sim_best20221027.onnx'
# onnx_path_file = r'D:\yolov5-5.0\runs\train\2023-3-16 epoch 300 15000\weight\best-cpu.onnx'
model = Yolov5ONNX(onnx_path)
input_image_size = (640,384)
# output, or_img = model.inference('data/images/img.png')
output, or_img = model.inference(image,input_image_size)
conf_thres, iou_thres = 0.5,0.5
outbox = filter_box(output, conf_thres, iou_thres )
if len(outbox) == 0:
print('没有发现物体')
sys.exit(0)
or_img = draw(or_img, outbox)
or_img=cv2.resize(or_img,(640,360))
return or_img,outbox4.使用PYQT5搭建GUI
本文使用QT产品Pyqt5搭建GUI，界面效果如下：
本课题搭建的GUI界面，功能较为简单，左侧只有Open Video按钮可以使用，右侧Open Model用于加载ONNX模型或者PT模型。
4.GUI主程序代码（GUI_main.py）
# -*- coding: utf-8 -*-
import cv2
import numpy as np
from utils.torch_utils import select_device
from utils.general import check_img_size
from detect_pt import pt_detection
import torch
import torch.backends.cudnn as cudnn
from models.experimental import attempt_load
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import Qt, QCoreApplication, QThread, pyqtSignal
from PyQt5.QtGui import QPixmap, QImage, QFont,QMovie
from PyQt5.QtWidgets import QFileDialog,QProgressDialog,QDialog,QLabel,QVBoxLayout
from PyQt5.QtCore import QSettings
from detect_onnx import detect_image
import time
class LoadingDialog(QDialog):
def __init__(self, parent=None):
super().__init__(parent)
self.setWindowModality(Qt.WindowModal)
self.setFixedSize(200, 100)
self.label = QLabel(self)
self.movie = QMovie("loading.gif")
self.label.setMovie(self.movie)
self.movie.start()
self.layout = QVBoxLayout()
self.layout.addWidget(self.label)
self.setLayout(self.layout)
def closeEvent(self, event):
Ui_MainWindow.closeCat()
super().closeEvent(event)
class Ui_MainWindow(object):
def setupUi(self, MainWindow):
MainWindow.setObjectName("MainWindow")
MainWindow.resize(1038, 765)
MainWindow.setStyleSheet("""
MainWindow {
padding: 10px;
border-radius: 5px;
}
""")
self.centralwidget = QtWidgets.QWidget(MainWindow)
self.centralwidget.setObjectName("centralwidget")
self.gridLayout = QtWidgets.QGridLayout(self.centralwidget)
self.gridLayout.setObjectName("gridLayout")
self.start_button = QtWidgets.QPushButton(self.centralwidget)
self.start_button.setObjectName("start_button")
self.start_button.setDisabled(True)
self.gridLayout.addWidget(self.start_button, 0, 9, 1, 1)
self.open_onnx_button = QtWidgets.QPushButton(self.centralwidget)
self.open_onnx_button.setObjectName("open_onnx_button")
self.gridLayout.addWidget(self.open_onnx_button, 0, 10, 1, 1)
self.person_button = QtWidgets.QPushButton(self.centralwidget)
self.person_button.setObjectName("person_button")
self.gridLayout.addWidget(self.person_button, 3, 3, 1, 1)
self.open_camera_button = QtWidgets.QPushButton(self.centralwidget)
self.open_camera_button.setObjectName("open_camera_button")
self.gridLayout.addWidget(self.open_camera_button, 0, 3, 1, 1)
self.Open_Video_button = QtWidgets.QPushButton(self.centralwidget)
self.Open_Video_button.setObjectName("Open_Video_button")
self.gridLayout.addWidget(self.Open_Video_button, 0, 2, 1, 1)
#左右图片投放区域
self.horizontalLayout = QtWidgets.QHBoxLayout()
self.horizontalLayout.setObjectName("horizontalLayout")
self.left_video = QtWidgets.QLabel(self.centralwidget)
self.left_video.setText("")
# self.size = self.left_video.size()
self.left_video.setObjectName("left_video")
self.horizontalLayout.addWidget(self.left_video)
self.right_video = QtWidgets.QLabel(self.centralwidget)
self.right_video.setText("")
self.right_video.setObjectName("right_video")
# self.right_video.resize(self.size)
# 设置两个控件的尺寸相同
self.horizontalLayout.addWidget(self.right_video)
self.gridLayout.addLayout(self.horizontalLayout, 1, 1, 1, 10)
# self.horizontalLayout.setStretchFactor(self.left_video, 10)
# self.horizontalLayout.setStretchFactor(self.right_video, 10)
#建立按钮
self.open_image_button = QtWidgets.QPushButton(self.centralwidget)
self.open_image_button.setObjectName("open_image_button")
self.gridLayout.addWidget(self.open_image_button, 0, 1, 1, 1)
self.red_person_button = QtWidgets.QPushButton(self.centralwidget)
self.red_person_button.setObjectName("red_person_button")
self.gridLayout.addWidget(self.red_person_button, 3, 1, 1, 1)
self.Bicycle_button = QtWidgets.QPushButton(self.centralwidget)
self.Bicycle_button.setObjectName("Bicycle_button")
self.gridLayout.addWidget(self.Bicycle_button, 3, 7, 1, 1)
self.truffick_button = QtWidgets.QPushButton(self.centralwidget)
self.truffick_button.setObjectName("truffick_button")
self.gridLayout.addWidget(self.truffick_button, 0, 6, 1, 1)
self.Motorcycle_button = QtWidgets.QPushButton(self.centralwidget)
self.Motorcycle_button.setObjectName("Motorcycle_button")
self.gridLayout.addWidget(self.Motorcycle_button, 3, 5, 1, 1)
self.car_button = QtWidgets.QPushButton(self.centralwidget)
self.car_button.setObjectName("car_button")
self.gridLayout.addWidget(self.car_button, 3, 9, 1, 1)
self.red_green_label = QtWidgets.QLabel(self.centralwidget)
self.red_green_label.setText("")
self.red_green_label.setObjectName("red_green_label")
self.gridLayout.addWidget(self.red_green_label, 0, 7, 1, 1)
#预测类显示标签
self.label_Motorcycle = QtWidgets.QLabel(self.centralwidget)
self.label_Motorcycle.setText("0")
self.label_Motorcycle.setObjectName("label_Motorcycle")
self.gridLayout.addWidget(self.label_Motorcycle, 3, 6, 1, 1)
self.label_Car = QtWidgets.QLabel(self.centralwidget)
self.label_Car.setText("0")
self.label_Car.setObjectName("label_Car")
self.gridLayout.addWidget(self.label_Car, 3, 10, 1, 1)
self.label_Person = QtWidgets.QLabel(self.centralwidget)
self.label_Person.setText("0")
self.label_Person.setObjectName("label_Person")
self.gridLayout.addWidget(self.label_Person, 3, 4, 1, 1)
self.label_Bicycle = QtWidgets.QLabel(self.centralwidget)
self.label_Bicycle.setText("0")
self.label_Bicycle.setObjectName("label_Bicycle")
self.gridLayout.addWidget(self.label_Bicycle, 3, 8, 1, 1)
self.label_red_person = QtWidgets.QLabel(self.centralwidget)
self.label_red_person.setText("0")
self.label_red_person.setObjectName("label_red_person")
self.gridLayout.addWidget(self.label_red_person, 3, 2, 1, 1)
#应用label显示样式
self.qlabel(self.label_red_person,self.label_Person,self.label_Bicycle,self.label_Motorcycle,self.label_Car)
MainWindow.setCentralWidget(self.centralwidget)
self.menubar = QtWidgets.QMenuBar(MainWindow)
self.menubar.setGeometry(QtCore.QRect(0, 0, 1038, 26))
self.menubar.setObjectName("menubar")
self.menu = QtWidgets.QMenu(self.menubar)
self.menu.setObjectName("menu")
self.menuAuto_labels = QtWidgets.QMenu(self.menubar)
self.menuAuto_labels.setObjectName("menuAuto_labels")
self.menuLabel_Tools = QtWidgets.QMenu(self.menubar)
self.menuLabel_Tools.setObjectName("menuLabel_Tools")
self.menu_4 = QtWidgets.QMenu(self.menubar)
font = QtGui.QFont()
font.setFamily("黑体")
self.menu_4.setFont(font)
self.menu_4.setObjectName("menu_4")
self.menu_2 = QtWidgets.QMenu(self.menu_4)
font = QtGui.QFont()
font.setFamily("黑体")
self.menu_2.setFont(font)
self.menu_2.setObjectName("menu_2")
MainWindow.setMenuBar(self.menubar)
self.statusbar = QtWidgets.QStatusBar(MainWindow)
self.statusbar.setObjectName("statusbar")
MainWindow.setStatusBar(self.statusbar)
self.actionOpen_Video = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionOpen_Video.setFont(font)
self.actionOpen_Video.setObjectName("actionOpen_Video")
self.actionQuit = QtWidgets.QAction(MainWindow)
self.actionQuit.setObjectName("actionQuit")
self.actionLabel_Image = QtWidgets.QAction(MainWindow)
self.actionLabel_Image.setObjectName("actionLabel_Image")
self.actionVoTT = QtWidgets.QAction(MainWindow)
self.actionVoTT.setObjectName("actionVoTT")
self.actionLabelme = QtWidgets.QAction(MainWindow)
self.actionLabelme.setObjectName("actionLabelme")
self.actionLabel_sprites = QtWidgets.QAction(MainWindow)
self.actionLabel_sprites.setObjectName("actionLabel_sprites")
self.actionData_resource = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionData_resource.setFont(font)
self.actionData_resource.setObjectName("actionData_resource")
self.actionData_Resource = QtWidgets.QAction(MainWindow)
self.actionData_Resource.setObjectName("actionData_Resource")
self.actionChange_image_size = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionChange_image_size.setFont(font)
self.actionChange_image_size.setObjectName("actionChange_image_size")
self.actionAug_image = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionAug_image.setFont(font)
self.actionAug_image.setObjectName("actionAug_image")
self.actionXML_TXT = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionXML_TXT.setFont(font)
self.actionXML_TXT.setObjectName("actionXML_TXT")
self.actionTXT_XML = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionTXT_XML.setFont(font)
self.actionTXT_XML.setObjectName("actionTXT_XML")
self.actionTrain_test_val = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionTrain_test_val.setFont(font)
self.actionTrain_test_val.setObjectName("actionTrain_test_val")
self.actionchoose_model = QtWidgets.QAction(MainWindow)
font = QtGui.QFont()
font.setFamily("黑体")
self.actionchoose_model.setFont(font)
self.actionchoose_model.setObjectName("actionchoose_model")
self.menu.addAction(self.actionOpen_Video)
self.menu.addSeparator()
self.menu.addAction(self.actionQuit)
self.menuAuto_labels.addAction(self.actionchoose_model)
self.menuLabel_Tools.addAction(self.actionLabel_Image)
self.menuLabel_Tools.addAction(self.actionLabelme)
self.menuLabel_Tools.addAction(self.actionVoTT)
self.menuLabel_Tools.addAction(self.actionLabel_sprites)
self.menu_2.addAction(self.actionXML_TXT)
self.menu_2.addAction(self.actionTXT_XML)
self.menu_4.addAction(self.actionChange_image_size)
self.menu_4.addAction(self.actionAug_image)
self.menu_4.addAction(self.menu_2.menuAction())
self.menu_4.addAction(self.actionTrain_test_val)
self.menubar.addAction(self.menu.menuAction())
self.menubar.addAction(self.menuAuto_labels.menuAction())
self.menubar.addAction(self.menuLabel_Tools.menuAction())
self.menubar.addAction(self.menu_4.menuAction())
# self.finished_signal = pyqtSignal()#设置pt模型加载初始化信号
# self.loading_dialog = LoadingDialog()
# self.finished_signal.connect(self.loading_dialog.close)
#创建文件浏览记录
self.settings = QSettings('MyCompany', 'MyApp')
self.open_onnx_button_clicked =False
self.Open_Video_button_clicked =False
self.retranslateUi(MainWindow)
self.actionOpen_Video.triggered.connect(self.open_video)
self.actionQuit.triggered.connect(QCoreApplication.quit)#退出按钮
# self.truffick_button.clicked.connect(self.show_light)
self.actionQuit.setShortcut('Ctrl+Q')
self.Open_Video_button.clicked.connect(self.open_video)
self.start_button.clicked.connect(self.start_video)
self.open_image_button.clicked.connect(self.open_image)
self.open_onnx_button.clicked.connect(self.open_model)
self.car_button.clicked.connect(self.detection_car)
self.person_button.clicked.connect(self.detection_person)
self.Bicycle_button.clicked.connect(self.detection_bicycle)
self.Motorcycle_button.clicked.connect(self.detection_motorcycle)
self.red_person_button.clicked.connect(self.detection_runlight)
QtCore.QMetaObject.connectSlotsByName(MainWindow)
# 定时器
# self.timer = QTimer()
# self.timer.timeout.connect(self.show_left_image)
# # 是否暂停
# self.paused = False
# classes_dict = {'person': "0", '红灯': "0", '绿灯': "0", '闯红灯': "0", '黄灯': "0", 'car': "0", 'bicycle': "0",
# 'motorcycle': "0"}
def closeCat(self):
if self.video_cap:
self.video_cap.release()
def detection_car(self):
self.classes = [1,2,4,5]
def detection_person(self):
self.classes = [0,1,2,3,4]
def detection_runlight(self):
self.classes = [1,2,3,4]
def detection_bicycle(self):
self.classes = [1,2,4,6]
def detection_motorcycle(self):
self.classes = [1,2,4,7]
def open_model(self):
# 获取上一次保存的目录
# last_dir = self.settings.value('last_dir', '.')
options = QFileDialog.Options()
options |= QFileDialog.DontUseNativeDialog
# if last_dir:
# filename, _ = QFileDialog.getOpenFileName(self.centralwidget, 'Open model', '', 'Model files (*.onnx)',
# options=options)
filename, _ = QFileDialog.getOpenFileName(self.centralwidget, 'Open model', '', 'Model files (*.onnx *.pt)',
options=options)
if filename:
self.end_name =filename.split('/')[-1].split('.')[-1]
self.open_onnx_button_clicked = True
if self.open_onnx_button_clicked and self.Open_Video_button_clicked:
self.start_button.setDisabled(False)
elif self.end_name == "onnx":
self.model_path = filename
# self.settings.setValue('last_dir', filename)
elif self.end_name == "pt":
self.model_path = filename
self.agument_pt = self.pt_model_init(self.model_path)#pt模型初始化
def qlabel(self,*labels):#标签样式
# # label标签样式设置
for label in labels:
label.setStyleSheet("""
QLabel {
font-size: 40px;
color: #00FF00;
background-color: black;
padding: 10px;
border-radius: 5px;
}
""")
label.setAlignment(Qt.AlignCenter)#居中
# font = QFont()
# font.setPointSize(20)
# label.setFont(font) # 字体
# QtWidgets.QLabel(self.centralwidget).setStyleSheet('background-color : black; color: #00FF00;;') # 颜色
# QtWidgets.QLabel(self.centralwidget).setAlignment(Qt.AlignCenter) # 居中
def show_light(self,red_green): # 显示交通灯状态
if red_green == "红灯":
image = cv2.imread(r'.\image\red.png')
else:
image = cv2.imread(r'.\image\green.png')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
height, width, channel = image.shape
bytes_per_line = 3 * width
q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888)
# 计算缩放比例
width_ratio = self.red_green_label.width() / q_image.width()
height_ratio = self.red_green_label.height() / q_image.height()
ratio = min(width_ratio, height_ratio)
# 缩放图片
scaled_image = q_image.scaled(q_image.width() * ratio, q_image.height() * ratio, Qt.KeepAspectRatio)
# jpg = QtGui.QPixmap(q_image).scaled(self.left_video.width(),
# self.left_video.height()) # 通过文件路径获取图片文件，并设置图片长宽为label控件的长宽
self.red_green_label.setPixmap(QPixmap.fromImage(scaled_image))
def open_image(self): # 打开图片
filename, _ = QFileDialog.getOpenFileName(self.centralwidget, 'Open Image', '',
'Image files (*.jpg *.png *.jpeg)')
if filename:
# self.start_button.setDisabled(False)
frame = cv2.imread(filename)
self.show_left_image(frame)
def open_video(self): # 打开文件
options = QFileDialog.Options()
options |= QFileDialog.DontUseNativeDialog
# QFileDialog.setDirectory(r"\detect-images")
filename, _ = QFileDialog.getOpenFileName(self.centralwidget, 'Open Video', '', 'Video files (*.mp4 *.avi)',
options=options)
if filename:
self.Open_Video_button_clicked = True
if self.open_onnx_button_clicked and self.Open_Video_button_clicked:
self.start_button.setDisabled(False)
# self.start_button.setDisabled(False)
self.video_cap = cv2.VideoCapture(filename)
ret, frame = self.video_cap.read()
if ret:
self.show_left_image(frame)
def start_video(self): # 开始按钮，开始预测
while True:
ret, frame = self.video_cap.read()
if ret:
self.frame = frame #用于预测
self.show_left_image(frame)
if self.end_name == "onnx":
self.detect_onnx_model()
else:
self.detect_pt_model()
cv2.waitKey(0)
else:
break
def detect_pt_model(self):
show_image, locate_data = pt_detection(model=self.model,opencv_imread_image=self.frame,classes = self.classes,device=self.device,
names=self.names,stride=self.stride,imgsz=self.imgsz,color=self.color,half=self.half) # 预测
self.show_right_image(show_image)
self.display_num_label(locate_data) # onnx显示数字
def detect_onnx_model(self):
show_image, locate_data = detect_image(self.model_path, self.frame) # 预测
self.show_right_image(show_image)
self.display_num_label(locate_data) # onnx显示数字
def show_left_image(self, image): # 显示原始图片
# if self.paused:
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) #BGR->RGB
height, width, channel = image.shape
bytes_per_line = 3 * width
#将array转化成Qlabel能显示的数据格式
q_image = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888)
# 计算缩放比例
width_ratio = self.left_video.width() / q_image.width()
height_ratio = self.left_video.height() / q_image.height()
self.width_ratio = q_image.width() *min(width_ratio, height_ratio)
self.height_ratio = q_image.height() * min(width_ratio, height_ratio)
# 缩放图片
scaled_image = q_image.scaled(self.width_ratio, self.height_ratio, Qt.KeepAspectRatio)
# jpg = QtGui.QPixmap(q_image).scaled(self.left_video.width(),
# self.left_video.height()) # 通过文件路径获取图片文件，并设置图片长宽为label控件的长宽
self.left_video.setPixmap(QPixmap.fromImage(scaled_image))
# else:
#
# pass
def show_right_image(self, show_image): # 显示预测图片
# if self.model_path.split("/")[-1].split(".")[-1] == "onnx":
#转换图片数据格式
show_image = cv2.cvtColor(show_image, cv2.COLOR_BGR2RGB) # BGR->RGB
height, width, channel = show_image.shape
bytes_per_line = 3 * width
# 将array转化成Qlabel能显示的数据格式
show_image = QImage(show_image.data, width, height, bytes_per_line, QImage.Format_RGB888)
# 缩放图片
scaled_image = show_image.scaled(self.width_ratio, self.height_ratio, Qt.KeepAspectRatio)
self.right_video.setPixmap(QPixmap.fromImage(scaled_image))
def display_num_label(self,label_data): # LCD显示数字
# print(label_data)
if self.end_name == "onnx":
labels = list(label_data[..., 5].astype(np.int32))
# print(labels)
self.label_Car.setText(str(labels.count(5)))
self.label_Person.setText(str(labels.count(0)))
self.label_Bicycle.setText(str(labels.count(6)))
self.label_Motorcycle.setText(str(labels.count(7)))
self.label_red_person.setText(str(labels.count(3)))
else:
# classes_dict = {'person': 0, '红灯': 0, '绿灯': 0, '闯红灯': 0, '黄灯': 0, 'car': 0, 'bicycle': 0, 'motorcycle': 0}
self.label_Car.setText(label_data["car"])
self.label_Person.setText(label_data["person"])
self.label_Bicycle.setText(label_data["bicycle"])
self.label_Motorcycle.setText(label_data["motorcycle"])
self.label_red_person.setText(label_data["闯红灯"])
if label_data["红灯"] != "0":
self.show_light(red_green="红灯")
else:
self.show_light(red_green="绿灯")
# pass
def label_image(self, label):
print(label)
def pt_model_init(self,path):
imgsz = 640
self.classes = [0, 1, 2, 3, 4, 5, 6, 7]
self.device = select_device("") # 选择设备GPU,CPU
weights_path = path
self.half = self.device.type != 'cpu' # 仅在 CUDA 上支持半精度
self.model = attempt_load(weights_path, map_location=self.device) # load FP32 model
if self.half:
self.model.half() # to FP16
cudnn.benchmark = True # 设置 True 以加快恒定图像大小推理
# 通过claee设置颜色
self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
# print(names)
# colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
self.color = {'person': [255, 191, 0], '红灯': [0, 0, 255], '绿灯': [0, 255, 0],
'闯红灯': [0, 0, 255], '黄灯': [0, 255, 255], 'car': [180, 105, 255],
'bicycle': [255, 0, 255], 'motorcycle': [0, 100, 255]
}
self.stride = int(self.model.stride.max()) # model stride模型步幅
self.imgsz = check_img_size(imgsz, s=self.stride) # check img_size
if self.device.type != 'cpu':
self.model(torch.zeros(1, 3, self.imgsz, self.imgsz).to(self.device).type_as(next(self.model.parameters()))) # run once
def retranslateUi(self, MainWindow):
_translate = QtCore.QCoreApplication.translate
MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
self.start_button.setText(_translate("MainWindow", "开始"))
self.open_onnx_button.setText(_translate("MainWindow", "open model"))
self.person_button.setText(_translate("MainWindow", "Person"))
self.open_camera_button.setText(_translate("MainWindow", "Open camera"))
self.Open_Video_button.setText(_translate("MainWindow", "Open Video"))
self.open_image_button.setText(_translate("MainWindow", "Open image"))
self.red_person_button.setText(_translate("MainWindow", "闯红灯"))
self.Bicycle_button.setText(_translate("MainWindow", "Bicycle"))
self.truffick_button.setText(_translate("MainWindow", "交通状态"))
self.Motorcycle_button.setText(_translate("MainWindow", "Motorcycle"))
self.car_button.setText(_translate("MainWindow", "Car"))
self.menu.setTitle(_translate("MainWindow", "菜单"))
self.menuAuto_labels.setTitle(_translate("MainWindow", "半自动标注"))
self.menuLabel_Tools.setTitle(_translate("MainWindow", "标注工具"))
self.menu_4.setTitle(_translate("MainWindow", "数据集预处理"))
self.menu_2.setTitle(_translate("MainWindow", "标注文件转换"))
self.actionOpen_Video.setText(_translate("MainWindow", "打开文件"))
self.actionQuit.setText(_translate("MainWindow", "Quit"))
self.actionLabel_Image.setText(_translate("MainWindow", "Label Image"))
self.actionVoTT.setText(_translate("MainWindow", "VoTT"))
self.actionLabelme.setText(_translate("MainWindow", "Labelme"))
self.actionLabel_sprites.setText(_translate("MainWindow", "Label Sprites"))
self.actionData_resource.setText(_translate("MainWindow", "打开数据集文件"))
self.actionData_Resource.setText(_translate("MainWindow", "Data Resource"))
self.actionChange_image_size.setText(_translate("MainWindow", "图片尺寸归一"))
self.actionAug_image.setText(_translate("MainWindow", "数据增强"))
self.actionXML_TXT.setText(_translate("MainWindow", "XML->TXT"))
self.actionTXT_XML.setText(_translate("MainWindow", "TXT->XML"))
self.actionTrain_test_val.setText(_translate("MainWindow", "数据集Train Val Test划分"))
self.actionchoose_model.setText(_translate("MainWindow", "选择已有模型"))
5.GUI启动程序（main.py）
import sys
import pyqt5_plugins.examplebutton
from GUI_main import Ui_MainWindow
from PyQt5.QtWidgets import *
if __name__ == '__main__':
app = QApplication(sys.argv)
MainWindow = QMainWindow()
ui = Ui_MainWindow() # 这是类函数的名称
ui.setupUi(MainWindow) # 运行类函数里的setupUi
# ui.show_light()
MainWindow.show() # 显示窗口
sys.exit(app.exec())
5279
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
