学习记录08：使用yolov3搭建检测界面重制版

前言

这次的代码更加简单，也解决了一些报错问题，基本上拿来就可以用了，大致也就实现了一个检测图片的功能，如果要改成自己的，只需将
run_detect（path）函数重写即可，因为传入的是图片路径，所以只需在函数里面接受路径，获取图片，然后调用你的yolo检测函数对图片检测，然后获得返回后的检测图片，然后将这个图片返回出这个函数即可实现，非常简单，如果连这个都无法实现还是建议找个人帮忙看看吧！

import cv2
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QIcon, QPixmap, QImage
import os
import sys
from yolo import YOLO  # 这里面就是我模型的检测函数，你只需要调用类似的，就可以实现任何一个模型的检测图片并展示
from PIL import Image

# 这里调用的是定义模型的函数，他在里面加载了模型的权重，因为不想每次检测都去加载一下，所以直接全局静态加载，然后就可以只需在模型
# 启动时加载一次，就可以一直调用了，一般的原始模型检测函数都是单张检测，然后每次调用这个单张检测都会重新加载一次模型，所以这里
# 直接提前加载，然后重复利用，减少预测的时间，同样的这里如果是自己的模型，也要改成自己模型预测时加载模型的样子
yolo = YOLO()
crop = False
count = False


class MyClassImage(QWidget):
    def __init__(self):
        super().__init__()
        self.label2 = None
        self.label3 = None
        self.initUI()   # 初始化界面并并展示
        self.openfile_name_image = ''   # 选择文件的路径
        self.image = None

    # 定义初始化组件
    def initUI(self):
        # 设置窗口大小
        self.resize(850, 400)
        # 设置窗口名称
        self.setWindowTitle("检测图片")
        # 创建按钮
        btn5 = QPushButton("退出检测图片", self)
        btn5.clicked.connect(self.close)
        btn1 = QPushButton("选择检测图片", self)

        # 关键点击事件，点击止呕启动self.select_image函数
        btn1.clicked.connect(self.select_image)
        btn2 = QPushButton("开始检测", self)

        # 关键点击事件，点击止呕启动self.detect函数
        btn2.clicked.connect(self.detect)

        # 创建标签，可以放文本或图片或空
        self.label2 = QLabel("", self)
        self.label2.resize(400, 300)
        self.label3 = QLabel("", self)
        self.label3.resize(400, 300)
        label4 = QLabel("              原始图片", self)
        label5 = QLabel("              检测图片", self)
        # 定义网格布局
        grid1 = QGridLayout()
        grid1.addWidget(label4, 0, 0)
        # 网格布局 第一行第二列防止label5
        grid1.addWidget(label5, 0, 1)

        hlo = QHBoxLayout()
        hlo.addStretch()
        grid = QGridLayout()
        # vlo.addStretch(0)
        grid.addWidget(btn1, 0, 0)
        grid.addWidget(btn2, 0, 1)
        grid.addWidget(btn5, 0, 2)

        hlo1 = QHBoxLayout()
        hlo1.addWidget(self.label2)
        hlo1.addWidget(self.label3)

        vlo = QVBoxLayout(self)
        vlo.addLayout(grid)
        vlo.addLayout(hlo1)
        vlo.addStretch(1)
        vlo.addLayout(grid1)
        vlo.addStretch(0)
        self.show()

    def closeEvent(self, event):
        result = QMessageBox.question(self, "提示：", "您真的要退出程序吗",
                                      QMessageBox.Yes | QMessageBox.No, QMessageBox.Yes)
        if result == QMessageBox.Yes:
            event.accept()
        else:
            event.ignore()

    def detect(self):
        if self.image is None:
            print('没有选择图片')
        elif self.image is not None:
            # 检测图片
            run_detect(self.openfile_name_image)
            # 读取检测之后的图片
            img = cv2.imread('img/result/'+self.openfile_name_image.split('/')[-1])
            img = cv2.resize(img, (400, 300), interpolation=cv2.INTER_AREA)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            # cv2.imshow('test', img)
            # cv2.waitKey(20)
            # 将图片放在标签self.label3中
            a = QImage(img.data, img.shape[1], img.shape[0], img.shape[1] * 3, QImage.Format_RGB888)
            self.label3.setPixmap(QPixmap.fromImage(a))
        pass

    def select_image(self):
        # temp为选择文件的路径  这里打开的是这个main.py函数的同级目录下的img文件夹
        temp, _ = QFileDialog.getOpenFileName(self, "选择照片文件", r"./img/")
        if temp is not None:
            self.openfile_name_image = temp
        #     读取选择的图片
        self.image = cv2.imread(self.openfile_name_image)
        # print(self.openfile_name_image)
        # 将路径中的图片读取之后放在self.label2
        self.label2.setPixmap(QPixmap(str(self.openfile_name_image)))
        self.label2.setScaledContents(True)
        # 读取收缩放至(400, 300)
        self.label2.setMaximumSize(400, 300)
        self.label2.setScaledContents(True)


# 重写这个函数
def run_detect(path):
    try:
        image = Image.open(path)
    except:
        print('Open Error! Try again!')
    else:
    ##这里是我模型检测函数，替换成自己的即可，这个函数返回的就是检测好的图片，然后保存在本地的同级目录下的img/result
        r_image = yolo.detect_image(image, crop=crop, count=count)
        r_image.save('img/result/' + path.split('/')[-1])


if __name__ == '__main__':
    app = QApplication(sys.argv)
    mc = MyClassImage()
    sys.exit(app.exec_())