使用PyQt5為YoloV5添加界面
近期因為疫情，無法正常入職上班。所以在家參考相關博文，視頻和代碼等，學習了PyQt5的基礎知識，並嘗試為YOLOV5添加界面。
反正啥也不咋會，在家瞎搗鼓搗鼓，總比閑着強唄~
項目為簡單Demo，僅供自己記錄過程，以及交流學習~
一、項目簡介
使用PyQt5為YoloV5添加一個可視化檢測界面，並實現簡單的界面跳轉，具體情況如下：
特點：
- UI界面與邏輯代碼分離
- 支持自選定模型
- 同時輸出檢測結果與相應相關信息
- 支持圖片，視頻，攝像頭檢測
- 支持視頻暫停與繼續檢測
目的：
- 熟悉QtDesign的使用
- 了解PyQt5基礎控件與布局方法
- 了解界面跳轉
- 了解信號與槽
- 熟悉視頻在PyQt中的處理方法
項目圖片：
二. 項目整體框架與代碼
項目地址：等待上傳
架構介紹：
- 整體為YoloV5的代碼
- ui文件夾中存放ui的py文件和原件，便於使用與更改
- ui_img存放ui使用的圖像文件
- utils中添加了一個用戶賬戶工具id_utils.py
- detect_logical.py是檢測界面的邏輯代碼
- main_logical.py是主界面的邏輯代碼
- userinfo.csv存放用戶賬號id信息
主要是在原始YoloV5-pyqt的基礎上進行修改，具體如下：
- 1.分離了界面和邏輯
- 2.增加了登錄，注冊功能
- 3.重構了部分功能代碼
三、快速開始
環境與相關文件配置：
- 按照 ult-yolov5 中requirement的要求配置環境，自行安裝PyQt5，注意都需要在一個evn環境中進行安裝與配置
- 下載或訓練一個模型，將“.pt”文件放到weights文件夾，（權重文件可以自己選，程序默認打開weights文件夾）
兩種程序使用方式：
- 直接運行detect_logical.py，進入檢測界面
- 運行main_logical.py，先登錄，在進入檢測界面（這是為了學習界面跳轉😂）
四、 核心部分代碼與簡單講解
- UI界面全部都QtDesign設計，然后由pyUIC生成，不做敘述，此部分重點在於如何使用QtDesign設計界面。UI界面可以自行修改，只要對應的控件名與邏輯函數中的對應即可。
- main_logical.py
此部分代碼是負責處理主界面的邏輯，具體包括登錄界面和注冊界面的邏輯，並根據需求實現界面跳轉。
主要思路：
1.導包：導入相關UI
2.創建界面類：每個界面的邏輯獨自為一個類，並在該類中初始化相關UI界面，以及信號槽。
3.信號與槽：使用connect操作，將控件綁定好具體的操作
4.界面跳轉：由於單個界面有具體的類，所以只需在跳轉功能函數中，實例一個具體界面對象，並設置為show；並根據需要決定是否關閉當前界面。
需要說一句的是，參考白日黑羽的課程，在創建新界面的時候，這里沒有直接在當前類中創建一個局部變量，而是使用lib包中的公共信息類shareInfo中的變量來實現的。
# -*- coding: utf-8 -*-
# @Modified by: Ruihao
# @ProjectName:yolov5-pyqt5
import sys
from datetime import datetime
from PyQt5 import QtWidgets
from PyQt5.QtWidgets import *
from utils.id_utils import get_id_info, sava_id_info # 賬號信息工具函數
from lib.share import shareInfo # 公共變量名
# 導入QT-Design生成的UI
from ui.login_ui import Login_Ui_Form
from ui.registe_ui import Ui_Dialog
# 導入設計好的檢測界面
from detect_logical import UI_Logic_Window
# 界面登錄
class win_Login(QMainWindow):
def __init__(self, parent = None):
super(win_Login, self).__init__(parent)
self.ui_login = Login_Ui_Form()
self.ui_login.setupUi(self)
self.init_slots()
self.hidden_pwd()
# 密碼輸入框隱藏
def hidden_pwd(self):
self.ui_login.edit_password.setEchoMode(QLineEdit.Password)
# 綁定信號槽
def init_slots(self):
self.ui_login.btn_login.clicked.connect(self.onSignIn) # 點擊按鈕登錄
self.ui_login.edit_password.returnPressed.connect(self.onSignIn) # 按下回車登錄
self.ui_login.btn_regeist.clicked.connect(self.create_id)
# 跳轉到注冊界面
def create_id(self):
shareInfo.createWin = win_Register()
shareInfo.createWin.show()
# 保存登錄日志
def sava_login_log(self, username):
with open('login_log.txt', 'a', encoding='utf-8') as f:
f.write(username + '\t log in at' + datetime.now().strftimestrftime+ '\r')
# 登錄
def onSignIn(self):
print("You pressed sign in")
# 從登陸界面獲得輸入賬戶名與密碼
username = self.ui_login.edit_username.text().strip()
password = self.ui_login.edit_password.text().strip()
# 獲得賬號信息
USER_PWD = get_id_info()
# print(USER_PWD)
if username not in USER_PWD.keys():
replay = QMessageBox.warning(self,"登陸失敗!", "賬號或密碼輸入錯誤", QMessageBox.Yes)
else:
# 若登陸成功，則跳轉主界面
if USER_PWD.get(username) == password:
print("Jump to main window")
# # 實例化新窗口
# # 寫法1：
# self.ui_new = win_Main()
# # 顯示新窗口
# self.ui_new.show()
# 寫法2：
# 不用self.ui_new,因為這個子窗口不是從屬於當前窗口,寫法不好
# 所以使用公用變量名
shareInfo.mainWin = UI_Logic_Window()
shareInfo.mainWin.show()
# 關閉當前窗口
self.close()
else:
replay = QMessageBox.warning(self, "!", "賬號或密碼輸入錯誤", QMessageBox.Yes)
# 注冊界面
class win_Register(QDialog):
def __init__(self, parent = None):
super(win_Register, self).__init__(parent)
self.ui_register = Ui_Dialog()
self.ui_register.setupUi(self)
self.init_slots()
# 綁定槽信號
def init_slots(self):
self.ui_register.pushButton_regiser.clicked.connect(self.new_account)
self.ui_register.pushButton_cancer.clicked.connect(self.cancel)
# 創建新賬戶
def new_account(self):
print("Create new account")
USER_PWD = get_id_info()
# print(USER_PWD)
new_username = self.ui_register.edit_username.text().strip()
new_password = self.ui_register.edit_password.text().strip()
# 判斷用戶名是否為空
if new_username == "":
replay = QMessageBox.warning(self, "!", "賬號不准為空", QMessageBox.Yes)
else:
# 判斷賬號是否存在
if new_username in USER_PWD.keys():
replay = QMessageBox.warning(self, "!", "賬號已存在", QMessageBox.Yes)
else:
# 判斷密碼是否為空
if new_password == "":
replay = QMessageBox.warning(self, "!", "密碼不能為空", QMessageBox.Yes)
else:
# 注冊成功
print("Successful!")
sava_id_info(new_username, new_password)
replay = QMessageBox.warning(self, "!", "注冊成功！", QMessageBox.Yes)
# 關閉界面
self.close()
# 取消注冊
def cancel(self):
self.close() # 關閉當前界面
if __name__ == "__main__":
app = QApplication(sys.argv)
# 利用共享變量名來實例化對象
shareInfo.loginWin = win_Login() # 登錄界面作為主界面
shareInfo.loginWin.show()
sys.exit(app.exec_())
detect_logical.py
此部分代碼是負責處理檢測的邏輯，具體包括實現模型選擇，初始化，圖片/視頻/攝像頭檢測。
主要思路：
1.導包：導入檢測的UI
2.界面初始化：初始化UI界面，為處理視頻初始化QTimer定時器，並初始化信號槽。
3.視頻檢測部分使用QTimer實現多線程處理。技術介紹見《PyQt5快速開發與實踐》：
4. 重要功能函數簡析：
本項目將目標檢測拆分為了模型加載和檢測兩個部分，model_init負責進行模型加載，而detect負責進行檢測並返回相關檢測信息。
- model_init：主體使用原始的yolov5中的初始化方法，主要參數可以在opt中進行設置。其中，權重默認為yolov5s，界面中可以自己選擇權重，標准的s，m，x模型是支持的。
- detect：考慮到3種檢測模式中都需要使用重復較多的代碼，所以將其抽出為一個函數。輸入為原始圖像，返回的是檢測信息。
- show_video_frame：負責各幀圖像的檢測與顯示。該函數在類初始化過程中，已經和定時器進行綁定，若計時超時，則調用show_video_frame。
- button_video_stop：通過設置num_stop 計數信號量和blockSignals來控制播放與暫停。
# -*- coding: utf-8 -*-
# @Modified by: Ruihao
# @ProjectName:yolov5-pyqt5
import sys
import cv2
import argparse
import random
import torch
import numpy as np
import torch.backends.cudnn as cudnn
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from utils.torch_utils import select_device
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_box2
from ui.detect_ui import Ui_MainWindow # 導入detect_ui的界面
class UI_Logic_Window(QtWidgets.QMainWindow):
def __init__(self, parent=None):
super(UI_Logic_Window, self).__init__(parent)
self.timer_video = QtCore.QTimer() # 創建定時器
self.ui = Ui_MainWindow()
self.ui.setupUi(self)
self.init_slots()
self.cap = cv2.VideoCapture()
self.num_stop = 1 # 暫停與播放輔助信號，note：通過奇偶來控制暫停與播放
# 權重初始文件名
self.openfile_name_model = None
# 控件綁定相關操作
def init_slots(self):
self.ui.pushButton_img.clicked.connect(self.button_image_open)
self.ui.pushButton_video.clicked.connect(self.button_video_open)
self.ui.pushButton_camer.clicked.connect(self.button_camera_open)
self.ui.pushButton_weights.clicked.connect(self.open_model)
self.ui.pushButton_init.clicked.connect(self.model_init)
self.ui.pushButton_stop.clicked.connect(self.button_video_stop)
self.ui.pushButton_finish.clicked.connect(self.finish_detect)
self.timer_video.timeout.connect(self.show_video_frame) # 定時器超時，將槽綁定至show_video_frame
# 打開權重文件
def open_model(self):
self.openfile_name_model, _ = QFileDialog.getOpenFileName(self.ui.pushButton_weights, '選擇weights文件',
'weights/')
if not self.openfile_name_model:
QtWidgets.QMessageBox.warning(self, u"Warning", u"打開權重失敗", buttons=QtWidgets.QMessageBox.Ok,
defaultButton=QtWidgets.QMessageBox.Ok)
else:
print('加載weights文件地址為：' + str(self.openfile_name_model))
# 加載相關參數，並初始化模型
def model_init(self):
# 模型相關參數配置
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='weights/yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--update', action='store_true', help='update all models')
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
parser.add_argument('--name', default='exp', help='save results to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
self.opt = parser.parse_args()
print(self.opt)
# 默認使用opt中的設置（權重等）來對模型進行初始化
source, weights, view_img, save_txt, imgsz = self.opt.source, self.opt.weights, self.opt.view_img, self.opt.save_txt, self.opt.img_size
# 若openfile_name_model不為空，則使用此權重進行初始化
if self.openfile_name_model:
weights = self.openfile_name_model
print("Using button choose model")
self.device = select_device(self.opt.device)
self.half = self.device.type != 'cpu' # half precision only supported on CUDA
cudnn.benchmark = True
# Load model
self.model = attempt_load(weights, map_location=self.device) # load FP32 model
stride = int(self.model.stride.max()) # model stride
self.imgsz = check_img_size(imgsz, s=stride) # check img_size
if self.half:
self.model.half() # to FP16
# Get names and colors
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]
print("model initial done")
# 設置提示框
QtWidgets.QMessageBox.information(self, u"Notice", u"模型加載完成", buttons=QtWidgets.QMessageBox.Ok,
defaultButton=QtWidgets.QMessageBox.Ok)
# 目標檢測
def detect(self, name_list, img):
'''
:param name_list: 文件名列表
:param img: 待檢測圖片
:return: info_show:檢測輸出的文字信息
'''
showimg = img
with torch.no_grad():
img = letterbox(img, new_shape=self.opt.img_size)[0]
# Convert
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
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, augment=self.opt.augment)[0]
# Apply NMS
pred = non_max_suppression(pred, self.opt.conf_thres, self.opt.iou_thres, classes=self.opt.classes,
agnostic=self.opt.agnostic_nms)
info_show = ""
# 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)
name_list.append(self.names[int(cls)])
single_info = plot_one_box2(xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)
# print(single_info)
info_show = info_show + single_info + "\n"
return info_show
# 打開圖片並檢測
def button_image_open(self):
print('button_image_open')
name_list = []
img_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打開圖片", "data/images", "*.jpg;;*.png;;All Files(*)")
# 判斷圖片是否為空
if not img_name:
QtWidgets.QMessageBox.warning(self, u"Warning", u"打開圖片失敗", buttons=QtWidgets.QMessageBox.Ok,
defaultButton=QtWidgets.QMessageBox.Ok)
else:
img = cv2.imread(img_name)
print(img_name)
info_show = self.detect(name_list, img)
print(info_show)
# 檢測信息顯示在界面
self.ui.textBrowser.setText(info_show)
# 檢測結果顯示在界面
self.result = cv2.cvtColor(img, 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.ui.label.setPixmap(QtGui.QPixmap.fromImage(self.QtImg))
self.ui.label.setScaledContents(True) # 設置圖像自適應界面大小
# 打開視頻並檢測
def button_video_open(self):
video_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打開視頻", "data/", "*.mp4;;*.avi;;All Files(*)")
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.timer_video.start(30) # 以30ms為間隔，啟動或重啟定時器
# 進行視頻識別時，關閉其他按鍵點擊功能
self.ui.pushButton_video.setDisabled(True)
self.ui.pushButton_img.setDisabled(True)
self.ui.pushButton_camer.setDisabled(True)
# 打開攝像頭檢測
def button_camera_open(self):
print("Open camera to detect")
# 設置使用的攝像頭序號，系統自帶為0
camera_num = 1
# 打開攝像頭
self.cap = cv2.VideoCapture(camera_num)
# 判斷攝像頭是否處於打開狀態
bool_open = self.cap.isOpened()
if not bool_open:
QtWidgets.QMessageBox.warning(self, u"Warning", u"打開攝像頭失敗", buttons=QtWidgets.QMessageBox.Ok,
defaultButton=QtWidgets.QMessageBox.Ok)
else:
self.timer_video.start(30)
self.ui.pushButton_video.setDisabled(True)
self.ui.pushButton_img.setDisabled(True)
self.ui.pushButton_camer.setDisabled(True)
# 定義視頻幀顯示操作
def show_video_frame(self):
name_list = []
flag, img = self.cap.read()
if img is not None:
info_show = self.detect(name_list, img) # 檢測結果寫入到原始img上
print(info_show)
# 檢測信息顯示在界面
self.ui.textBrowser.setText(info_show)
show = cv2.resize(img, (640, 480)) # 直接將原始img上的檢測結果進行顯示
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.ui.label.setPixmap(QtGui.QPixmap.fromImage(showImage))
self.ui.label.setScaledContents(True) # 設置圖像自適應界面大小
else:
self.timer_video.stop()
self.cap.release()
self.ui.label.clear()
# 視頻幀顯示期間，禁用其他檢測按鍵功能
self.ui.pushButton_video.setDisabled(False)
self.ui.pushButton_img.setDisabled(False)
self.ui.pushButton_camer.setDisabled(False)
# 暫停與繼續檢測
def button_video_stop(self):
self.timer_video.blockSignals(False)
# 暫停檢測
# 若QTimer已經觸發，且激活
if self.timer_video.isActive() == True and self.num_stop%2 == 1:
self.ui.pushButton_stop.setText(u'暫停檢測') # 當前狀態為暫停狀態
self.num_stop = self.num_stop + 1 # 調整標記信號為偶數
self.timer_video.blockSignals(True)
# 繼續檢測
else:
self.num_stop = self.num_stop + 1
self.ui.pushButton_stop.setText(u'繼續檢測')
# 結束視頻檢測
def finish_detect(self):
# self.timer_video.stop()
self.cap.release() # 釋放cap
self.ui.label.clear() # 清空label畫布
# 啟動其他檢測按鍵功能
self.ui.pushButton_video.setDisabled(False)
self.ui.pushButton_img.setDisabled(False)
self.ui.pushButton_camer.setDisabled(False)
# 結束檢測時，查看暫停功能是否復位，將暫停功能恢復至初始狀態
# Note:點擊暫停之后，num_stop為偶數狀態
if(self.num_stop%2 == 0):
print("Reset stop/begin!")
self.ui.pushButton_stop.setText(u'暫停/繼續')
self.num_stop = self.num_stop + 1
self.timer_video.blockSignals(False)
if __name__ == '__main__':
app = QtWidgets.QApplication(sys.argv)
current_ui = UI_Logic_Window()
current_ui.show()
sys.exit(app.exec_())
五、 參考與致謝
- 《PyQt5快速開發與實踐》
- www.python3.vip
- B站白月黑羽的PyQt教程 https://www.bilibili.com/video/BV1cJ411R7bP?from=search&seid=7706040462590056686
- https://xugaoxiang.blog.csdn.net/article/details/118384430 從這個博主的博客中學到了很多知識，感覺博主，博主的代碼框架也很好，也是本文代碼是在其基礎上進行學習和修改的
- Github項目：YOLOv3GUI_Pytorch_PyQt5