使用PyQt5为YoloV5添加界面

近期因为疫情，无法正常入职上班。所以在家参考相关博文，视频和代码等，学习了PyQt5的基础知识，并尝试为YOLOV5添加界面。
反正啥也不咋会，在家瞎捣鼓捣鼓，总比闲着强呗~
项目为简单Demo，仅供自己记录过程，以及交流学习~

一、项目简介

使用PyQt5为YoloV5添加一个可视化检测界面，并实现简单的界面跳转，具体情况如下：
特点：

1. UI界面与逻辑代码分离
2. 支持自选定模型
3. 同时输出检测结果与相应相关信息
4. 支持图片，视频，摄像头检测
5. 支持视频暂停与继续检测

目的：

1. 熟悉QtDesign的使用
2. 了解PyQt5基础控件与布局方法
3. 了解界面跳转
4. 了解信号与槽
5. 熟悉视频在PyQt中的处理方法

项目图片：

二. 项目整体框架与代码

项目地址：https://github.com/BonesCat/YoloV5_PyQt5/tree/main
B站视频：https://www.bilibili.com/video/BV1ZU4y1E7at/

更新：代码更新至百度云，自取！！！
链接：https://pan.baidu.com/s/1qF6JMIl8qoTVtCTnBfIaZw
提取码：3gdn
可以保存视频和图片啊，上传了截图给大家了。还是不行自己在动手改改吧，不是很难。

架构介绍：

1. 整体为YoloV5的代码
2. ui文件夹中存放ui的py文件和原件，便于使用与更改
3. ui_img存放ui使用的图像文件
4. utils中添加了一个用户账户工具id_utils.py
5. detect_logical.py是检测界面的逻辑代码
6. main_logical.py是主界面的逻辑代码
7. userinfo.csv存放用户账号id信息

主要是在原始YoloV5-pyqt的基础上进行修改，具体如下：

• 1.分离了界面和逻辑
• 2.增加了登录，注册功能
• 3.重构了部分功能代码

三、快速开始

环境与相关文件配置：

• 按照 ult-yolov5 中requirement的要求配置环境，自行安装PyQt5，注意都需要在一个evn环境中进行安装与配置
• 下载或训练一个模型，将“.pt”文件放到weights文件夹，（权重文件可以自己选，程序默认打开weights文件夹）
• 当前设置的为cpu运行模式，无奈当前木有小钱钱，可以自己设置为gpu，在opt里面设置就可了，检测速度会提升

两种程序使用方式：

• 直接运行detect_logical.py，进入检测界面
• 运行main_logical.py，先登录，在进入检测界面（这是为了学习界面跳转😂）

四、 核心部分代码与简单讲解

• UI界面全部都QtDesign设计，然后由pyUIC生成，不做叙述，此部分重点在于如何使用QtDesign设计界面。UI界面可以自行修改，只要对应的控件名与逻辑函数中的对应即可。
  我是喜欢先摆放控件，再调整布局。
  1.关键是先大概设计好各个部件的位置，然后按照水平和竖直的布局方式，对各个局部控件进行布局调整（如上中下，左中右这种布局）。
  2.局部调整好了，对整体在进行调整。全局都有布局的情况下，页面是可以进行自动缩放的。
  注：白夜黑羽的课程里面，有一节说这个的。可以去看视频。整体的设计后的结果，可以使用QtDesigner看原始的UI文件，在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

** 码字真的超级费时间，请转载请注明出处哦，谢谢**