赞
踩
全流程教程,从数据采集到模型使用到最终展示。 支持图片检测、视频检测、摄像头实时检测,还支持视频的暂停、结束等功能。
本文使用YOLOV5来进行车辆和车牌的定位检测,定位车牌后使用CRNN进行车牌类信息的识别即做到OCR的效果。使用爬虫技术爬取百度上的图片从而制作数据集,使用PyQt5来制作我们的可视化展示软件。
若有任何疑问和建议欢迎评论区讨论。
先放上最终UI实现效果
已制作数据集 数据集
共有三个数据集:车型识别数据集、车牌定位数据集、单一车牌字符识别数据集,均为coco格式或voc格式
车型识别数据集包含car、truck、bus、taix、suv5种车
车牌定位数据集只进行了标注,未进行分类
单一车牌字符识别数据集包含0-9、A-Z、各省简称
自己制作数据集可以参考如下步骤
可以通过爬虫爬取百度上的图片。
这里直接上代码,如果想详细了解可以参考我的另外一篇文章python爬取百度图片,可以大量批量爬取(仅供学习,很详细)
import requests#导入请求库 import time import re #设定爬取的总数 total=90 batch_size=30 all_success=0 for i in range(total//batch_size): url='https://image.baidu.com/search/acjson?tn=resultjson_com&logid=10371129381236677678&ipn=rj&ct=201326592&is=&fp=result&fr=&word=车&queryWord=车&cl=2&lm=-1&ie=utf-8&oe=utf-8&adpicid=&st=&z=&ic=&hd=&latest=©right=&s=&se=&tab=&width=&height=&face=&istype=&qc=&nc=1&expermode=&nojc=&isAsync=&pn={0}&rn=30&gsm=3c&1682846532783='.format((i+1)*30) #添加请求头,模拟浏览器,有些网站可以不加这个,不过最好是加上 headers={'User-Agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.63 Safari/537.36'} res=requests.get(url,headers=headers)#发送请求,返回数据 html=res.text#把返回的内容解析 # 使用正则表达式匹配图片url img_url_list=re.findall('"thumbURL":"(.*?)"',html) #print(img_url_list) for j in range(len(img_url_list)): res_img=requests.get(img_url_list[j],headers=headers) img=res_img.content#这个里是图片,我们需要返回二进制数据 # 图片保存的路径 with open('D:\\code\\person\\mask\\'+str(j)+'mask_img.jpg','wb')as f: f.write(img) time.sleep(3)#每当保存一张图片,先暂停一下,不然太频繁容易发现是机器爬虫,导致无法获取 all_success=all_success+1 print("爬取{}张图片成功".format(all_success)) print("爬取{}张图片成功".format(all_success))
一些爬取的效果图如下。
labelme是图形图像注释工具,它是用Python编写的,并将Qt用于其图形界面。说直白点,它是有界面的, 像软件一样,可以交互,但是它又是由命令行启动的,比软件的使用稍微麻烦点。其界面如下图:
github链接: labelme https://github.com/wkentaro/labelme
它的功能很多,包括:
YOLO框架(You Only Look Once)与RCNN系列算法不一样,是以不同的方式处理对象检测。它将整个图像放在一个实例中,并预测这些框的边界框坐标和及所属类别概率。使用YOLO算法最大优的点是速度极快,每秒可处理45帧,也能够理解一般的对象表示。
在本节中,将介绍YOLO用于检测给定图像中的对象的处理步骤。
首先,输入图像:
然后,YOLO将输入图像划分为网格形式(例如3 X 3):
最后,对每个网格应用图像分类和定位处理,获得预测对象的边界框及其对应的类概率。
整个过程是不是很清晰,下面逐一详细介绍。首先需要将标记数据传递给模型以进行训练。假设已将图像划分为大小为3 X 3的网格,且总共只有3个类别,分别是行人(c1)、汽车(c2)和摩托车(c3)。因此,对于每个单元格,标签y将是一个八维向量:
其中:
pc定义对象是否存在于网格中(存在的概率);
bx、by、bh、bw指定边界框;
c1、c2、c3代表类别。如果检测对象是汽车,则c2位置处的值将为1,c1和c3处的值将为0;
假设从上面的例子中选择第一个网格:
由于此网格中没有对象,因此pc将为零,此网格的y标签将为:
?意味着其它值是什么并不重要,因为网格中没有对象。下面举例另一个有车的网格(c2=1):
在为此网格编写y标签之前,首先要了解YOLO如何确定网格中是否存在实际对象。大图中有两个物体(两辆车),因此YOLO将取这两个物体的中心点,物体将被分配到包含这些物体中心的网格中。中心点左侧网格的y标签会是这样的:
由于此网格中存在对象,因此pc将等于1,bx、by、bh、bw将相对于正在处理的特定网格单元计算。由于检测出的对象是汽车,所以c2=1,c1和c3均为0。对于9个网格中的每一个单元格,都具有八维输出向量。最终的输出形状为3X3X8。
使用上面的例子(输入图像:100X100X3,输出:3X3X8),模型将按如下方式进行训练:
使用经典的CNN网络构建模型,并进行模型训练。在测试阶段,将图像传递给模型,经过一次前向传播就得到输出y。为了简单起见,使用3X3网格解释这一点,但通常在实际场景中会采用更大的网格(比如19X19)。
即使一个对象跨越多个网格,它也只会被分配到其中点所在的单个网格。可以通过增加更多网格来减少多个对象出现在同一网格单元中的几率。
如果你有git,则使用git clone
git clone https://github.com/ultralytics/yolov5 # clone
如果你没有git,你可以使用Dwonload ZIP下载代码项目。
yolov5代码地址:yolov5
注意:yolov5的代码是最新的v8.0版本
可以通过这个链接下载6.0版本https://github.com/ultralytics/yolov5/tree/v6.0
import torch
if __name__ == '__main__':
print(torch.zeros(1))
pip install -r requirement.txt
如下图所示,等待安装完成
我们需要下载其预训练的权重文件然后再此基础上进行调整训练,这样在数据集数量较小时也能取得不错的检测准确率。
供选择的有yolov5s,yolov5m,yolov5l,yolov5x。模型的大小逐渐增大,训练时间更长,准确率提高。
这里我们以yolov5s为例训练。下载地址为yolov5s.pt
所有权重下载地址可在https://github.com/ultralytics/yolov5/releases/tag/v6.0界面找到
配置models/yolov5s_mask.yaml 可以直接复制yolov5s.yaml文件,然后在nc即类别出进行修改,其数量为5。其中anchors参数表示锚框的大小,可以通过对数据集进行knn聚类得到,这里直接使用默认即对COCO数据集进行聚类的结果。
配置mask.yaml 。这里train指定训练数据集所在位置,val测试数据集所在位置,nc类别数,names类别的名称(注意顺序,我们txt文件中索引0代表未戴口罩,1代表戴口罩,所以这里顺序是‘no_mask’,‘mask’)
train: /home/cmy/newdataset/mask/train/ #注意修改为自己的数据集所在位置
val: /home/cmy/newdataset/mask/val/
nc: 5
names: ['car','truck','bus','taix','suv']
python3 train.py --img 640 --batch 8 --epochs 50 --data mask.yaml --cfg yolov5s_mask.yaml --weights "yolov5s.pt"
如果出现显卡空间不足的情况可以改小–bath参数
会在/runs/train/exp/weights/best.pt下生成最终的权重文件。
将此文件复制到yolov5目录下后续使用
同理做出对车牌的检测。
CRNN基本网络结构
整个CRNN网络可以分为三个部分:Convlutional Layers,Recurrent Layers,Transcription Layers
这里的卷积层就是一个普通的CNN网络,用于提取输入图像的Convolutional feature maps,即将图像(32,100,3)转换为
(1,25,512)卷积特征矩阵,网络细节请参考本文给出的实现代码。
这里的循环网络层是一个深层双向LSTM网络,在卷积特征的基础上继续提取文字序列特征。
所谓深层RNN网络,是指超过两层的RNN网络。对于单层双向RNN网络,结构如下:
深层双向RNN网络
由于CNN输出的Feature map大小是(1,25,512),所以对于RNN最大时间长度 T=25(即有25个时间输入,每个输入
列向量有512维度)。
将RNN输出做softmax后,为字符输出。
可以通过github下载获取CRNN_Chinese_Characters_Rec
我们通过pyqt来制作展示的界面
PyQt5 是 Digia的一套 Qt5 应用框架与 python 的结合,同时支持 python2.x和 python3.x。
这里使用的是Python 3.x。Qt库由 Riverbank Computing开发,是最强大的GUI库之一 。
PyQt5 是由一系列 Python 模块组成。超过 620 个类,6000 函数和方法。能在诸如 Unix、Windows 和Mac OS 等主流操作系统上运行。PyQt5 有两种证书,GPL和 商业证书。
PyQt5 有两种安装方式,一种是从官网下载源码安装,另外一种是使用 pip 安装。
这里我推荐大家使用pip 安装。因为它会自动根据你的Python 版本来选择合适的 PyQt5 版本,如果是手动下载源码安装,难免会选择出错。建议使用比较稳妥的安装方式。
pip3 install PyQt5
另外,如果你的网络访问外网不是很好的话建议使用豆瓣的镜像下载,不然会很很慢或者直接安装失败。
pip install PyQt5 -i https://pypi.douban.com/simple
执行以下代码:
import sys
from PyQt5.QtWidgets import QWidget, QApplication
app = QApplication(sys.argv)
widget = QWidget()
widget.resize(640, 480)
widget.setWindowTitle("Hello, PyQt5!")
widget.show()
sys.exit(app.exec())
如果没有报错,弹出了一个标题为"Hello, PyQt5!"的窗口,则说明安装成功。
若pip安装pyqt5报错 error: Microsoft Visual C++ 14.0 or greater is required. Get it with “Microsoft C++ Build Tools”: https://visualstudio.microsoft.com/visual-cpp-build-tools/
error: Microsoft Visual C++ 14.0 or greater is required. Get it with "Microsoft C++ Build Tools": https://visualstudio.microsoft.com/visual-cpp-build-tools/
可参考我另外一篇文章 已解决(pip安装pyqt5报错) error: Microsoft Visual C++ 14.0 or greater is required. Get it with “Microsoft
对于运行项目做到此即可,若对pyqt感兴趣如图形界面开发工具Qt Designer等可以参考文章https://zhuanlan.zhihu.com/p/162866700
OpenCV(开源的计算机视觉库)是基于BSD协议,因此它可免费用于学术和商业用途。其提供C++,C,Python和Java接口,支持Windows,Linux,Mac OS,iOS和Android。
我们使用OpenCV来处理图片和视频,以便于将图片转为Yolov5模型需要的输入。
首先我们得先安装另一个第三方库numpy,这是opencv的依赖库,没有它无法进行python-opencv开发。
#安装numpy:
pip install numpy
#安装opencv-python:
pip install opencv-python
使用pyqt设计我们界面的布局,新建布局代码mainwindow_ui .py。pyqt原始提供的组件并不美观,为了美化组件可以使用qt_material。其安装和使用都比较简单。
安装qt_material
pip install qt_material
布局代码mainwindow_ui .py
# -*- coding: utf-8 -*- # Form implementation generated from reading ui file 'mainwindow_ui.ui' # # Created by: PyQt5 UI code generator 5.15.7 # # WARNING: Any manual changes made to this file will be lost when pyuic5 is # run again. Do not edit this file unless you know what you are doing. from PyQt5 import QtCore, QtGui, QtWidgets class Ui_MainWindow(object): def setupUi(self, MainWindow): MainWindow.setObjectName("目标检测") MainWindow.resize(1367, 850) font = QtGui.QFont() font.setFamily("Adobe 宋体 Std L") MainWindow.setFont(font) self.centralwidget = QtWidgets.QWidget(MainWindow) self.centralwidget.setObjectName("centralwidget") self.groupBox = QtWidgets.QGroupBox(self.centralwidget) self.groupBox.setGeometry(QtCore.QRect(10, 50, 141, 771)) font = QtGui.QFont() font.setFamily("Adobe 黑体 Std R") self.groupBox.setFont(font) self.groupBox.setObjectName("groupBox") self.btn_loadweight = QtWidgets.QPushButton(self.groupBox) self.btn_loadweight.setGeometry(QtCore.QRect(10, 20, 121, 31)) self.btn_loadweight.setObjectName("btn_loadweight") self.btn_loadimg = QtWidgets.QPushButton(self.groupBox) self.btn_loadimg.setGeometry(QtCore.QRect(10, 60, 121, 31)) self.btn_loadimg.setObjectName("btn_loadimg") self.btn_loadvideo = QtWidgets.QPushButton(self.groupBox) self.btn_loadvideo.setGeometry(QtCore.QRect(10, 100, 121, 31)) self.btn_loadvideo.setObjectName("btn_loadvideo") self.btn_opencamera = QtWidgets.QPushButton(self.groupBox) self.btn_opencamera.setGeometry(QtCore.QRect(10, 140, 121, 31)) self.btn_opencamera.setObjectName("btn_opencamera") self.btn_stop = QtWidgets.QPushButton(self.groupBox) self.btn_stop.setGeometry(QtCore.QRect(10, 500, 121, 31)) self.btn_stop.setObjectName("btn_going") self.btn_over = QtWidgets.QPushButton(self.groupBox) self.btn_over.setGeometry(QtCore.QRect(10, 540, 121, 31)) self.btn_over.setObjectName("btn_stop") self.fps_label = QtWidgets.QLabel(self.groupBox) self.fps_label.setGeometry(QtCore.QRect(50, 630, 121, 31)) self.fps_label.setObjectName("fps_label") font.setPointSize(15) # font.setFamily("Adobe Devanagari") self.fps_label.setFont(font) self.btn_takephoto = QtWidgets.QPushButton(self.groupBox) self.btn_takephoto.setGeometry(QtCore.QRect(10, 380, 121, 31)) self.btn_takephoto.setObjectName("btn_takephoto") self.btn_clear = QtWidgets.QPushButton(self.groupBox) self.btn_clear.setGeometry(QtCore.QRect(10, 460, 121, 31)) self.btn_clear.setObjectName("btn_clear") self.btn_labelimg = QtWidgets.QPushButton(self.groupBox) self.btn_labelimg.setGeometry(QtCore.QRect(10, 420, 121, 31)) self.btn_labelimg.setObjectName("btn_labelimg") self.btn_closecamera = QtWidgets.QPushButton(self.groupBox) self.btn_closecamera.setGeometry(QtCore.QRect(10, 220, 121, 31)) self.btn_closecamera.setObjectName("btn_closecamera") self.layoutWidget = QtWidgets.QWidget(self.groupBox) self.layoutWidget.setGeometry(QtCore.QRect(10, 300, 121, 23)) self.layoutWidget.setObjectName("layoutWidget") self.horizontalLayout_2 = QtWidgets.QHBoxLayout(self.layoutWidget) self.horizontalLayout_2.setContentsMargins(0, 0, 0, 0) self.horizontalLayout_2.setObjectName("horizontalLayout_2") # self.label_2 = QtWidgets.QLabel(self.layoutWidget) # self.label_2.setObjectName("label_2") # self.horizontalLayout_2.addWidget(self.label_2) # self.confSpinBox = QtWidgets.QDoubleSpinBox(self.layoutWidget) # self.confSpinBox.setObjectName("confSpinBox") # self.horizontalLayout_2.addWidget(self.confSpinBox) self.layoutWidget_2 = QtWidgets.QWidget(self.groupBox) self.layoutWidget_2.setGeometry(QtCore.QRect(10, 340, 121, 23)) self.layoutWidget_2.setObjectName("layoutWidget_2") self.horizontalLayout_3 = QtWidgets.QHBoxLayout(self.layoutWidget_2) self.horizontalLayout_3.setContentsMargins(0, 0, 0, 0) self.horizontalLayout_3.setObjectName("horizontalLayout_3") # self.label_3 = QtWidgets.QLabel(self.layoutWidget_2) # self.label_3.setObjectName("label_3") # self.horizontalLayout_3.addWidget(self.label_3) # self.latencySpinBox = QtWidgets.QDoubleSpinBox(self.layoutWidget_2) # self.latencySpinBox.setObjectName("latencySpinBox") # self.horizontalLayout_3.addWidget(self.latencySpinBox) self.layoutWidget1 = QtWidgets.QWidget(self.groupBox) self.layoutWidget1.setGeometry(QtCore.QRect(10, 260, 121, 23)) self.layoutWidget1.setObjectName("layoutWidget1") self.horizontalLayout = QtWidgets.QHBoxLayout(self.layoutWidget1) self.horizontalLayout.setContentsMargins(0, 0, 0, 0) self.horizontalLayout.setObjectName("horizontalLayout") # self.label = QtWidgets.QLabel(self.layoutWidget1) # self.label.setObjectName("label") # self.horizontalLayout.addWidget(self.label) # self.iouSpinBox = QtWidgets.QDoubleSpinBox(self.layoutWidget1) # self.iouSpinBox.setObjectName("iouSpinBox") # self.horizontalLayout.addWidget(self.iouSpinBox) self.btn_camera_detect = QtWidgets.QPushButton(self.groupBox) self.btn_camera_detect.setGeometry(QtCore.QRect(10, 180, 121, 31)) self.btn_camera_detect.setObjectName("btn_camera_detect") self.groupBox_2 = QtWidgets.QGroupBox(self.centralwidget) self.groupBox_2.setGeometry(QtCore.QRect(170, 50, 1181, 591)) self.groupBox_2.setObjectName("groupBox_2") self.label_origin = QtWidgets.QLabel(self.groupBox_2) self.label_origin.setGeometry(QtCore.QRect(10, 30, 571, 551)) font = QtGui.QFont() font.setPointSize(15) self.label_origin.setFont(font) self.label_origin.setStyleSheet("background-color: rgb(85, 170, 255);") self.label_origin.setAlignment(QtCore.Qt.AlignCenter) self.label_origin.setObjectName("label_origin") self.label_detect = QtWidgets.QLabel(self.groupBox_2) self.label_detect.setGeometry(QtCore.QRect(600, 30, 571, 551)) font = QtGui.QFont() font.setPointSize(15) self.label_detect.setFont(font) self.label_detect.setStyleSheet("background-color: rgb(170, 255, 127);") self.label_detect.setAlignment(QtCore.Qt.AlignCenter) self.label_detect.setObjectName("label_detect") self.label_mian_titlle = QtWidgets.QLabel(self.centralwidget) self.label_mian_titlle.setGeometry(QtCore.QRect(530, 10, 281, 41)) font = QtGui.QFont() font.setFamily("Adobe Devanagari") font.setPointSize(20) self.label_mian_titlle.setFont(font) self.label_mian_titlle.setLayoutDirection(QtCore.Qt.LeftToRight) self.label_mian_titlle.setAlignment(QtCore.Qt.AlignCenter) self.label_mian_titlle.setObjectName("label_mian_titlle") self.lcdNumber = QtWidgets.QLCDNumber(self.centralwidget) self.lcdNumber.setGeometry(QtCore.QRect(1050, 20, 301, 31)) font = QtGui.QFont() font.setFamily("Adobe 黑体 Std R") self.lcdNumber.setFont(font) self.lcdNumber.setObjectName("lcdNumber") self.lcdNumber.setDigitCount(20) self.groupBox_3 = QtWidgets.QGroupBox(self.centralwidget) self.groupBox_3.setGeometry(QtCore.QRect(170, 650, 581, 171)) self.groupBox_3.setObjectName("groupBox_3") self.textBrowser_print = QtWidgets.QTextBrowser(self.groupBox_3) self.textBrowser_print.setGeometry(QtCore.QRect(10, 30, 561, 141)) self.textBrowser_print.setObjectName("textBrowser_print") self.groupBox_4 = QtWidgets.QGroupBox(self.centralwidget) self.groupBox_4.setGeometry(QtCore.QRect(770, 650, 581, 171)) self.groupBox_4.setObjectName("groupBox_4") self.textBrowser_detect = QtWidgets.QTextBrowser(self.groupBox_4) self.textBrowser_detect.setGeometry(QtCore.QRect(10, 30, 561, 141)) self.textBrowser_detect.setObjectName("textBrowser_detect") self.label_author = QtWidgets.QLabel(self.centralwidget) self.label_author.setGeometry(QtCore.QRect(20, 20, 111, 16)) font = QtGui.QFont() font.setFamily("Adobe 黑体 Std R") font.setPointSize(12) self.label_author.setFont(font) self.label_author.setObjectName("label_author") MainWindow.setCentralWidget(self.centralwidget) self.menubar = QtWidgets.QMenuBar(MainWindow) self.menubar.setGeometry(QtCore.QRect(0, 0, 1367, 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", "MainWindow")) self.groupBox.setTitle(_translate("MainWindow", "")) self.btn_loadweight.setText(_translate("MainWindow", "加载Weights")) self.btn_loadimg.setText(_translate("MainWindow", "图片检测")) self.btn_loadvideo.setText(_translate("MainWindow", "视频检测")) self.btn_opencamera.setText(_translate("MainWindow", "打开摄像头")) self.btn_over.setText(_translate("MainWindow", "结束")) self.fps_label.setText(_translate("MainWindow", "帧率")) self.btn_stop.setText(_translate("MainWindow", "暂停")) self.btn_takephoto.setText(_translate("MainWindow", "拍照")) self.btn_clear.setText(_translate("MainWindow", "清除输出")) self.btn_labelimg.setText(_translate("MainWindow", "标注工具")) self.btn_closecamera.setText(_translate("MainWindow", "关闭摄像头")) # self.label_2.setText(_translate("MainWindow", "conf:")) # self.label_3.setText(_translate("MainWindow", "latency:")) # self.label.setText(_translate("MainWindow", "IoU:")) self.btn_camera_detect.setText(_translate("MainWindow", "摄像头检测")) self.groupBox_2.setTitle(_translate("MainWindow", "显示区域")) self.label_origin.setText(_translate("MainWindow", "原图区域")) self.label_detect.setText(_translate("MainWindow", "检测区域")) self.label_mian_titlle.setText(_translate("MainWindow", "YoLo检测界面")) self.groupBox_3.setTitle(_translate("MainWindow", "打印输出")) self.groupBox_4.setTitle(_translate("MainWindow", "检测输出")) self.label_author.setText(_translate("MainWindow", " "))
注意这个代码要放在yolov5项目代码中一起使用以及上一节的布局代码mainwindow_ui.py。yolov5代码的获取和配置在第2节YOLOV5有描述。yolov5代码获取链接yolo。单独运行这个是不行的。
import argparse import datetime import os import random import sys import time import cv2 import numpy as np import torch from PyQt5 import QtWidgets, QtGui, QtCore from PyQt5.QtCore import QTimer, QDateTime, QDate, QTime, QThread, pyqtSignal from PyQt5.QtGui import * from PyQt5.QtWidgets import * from torch.backends import cudnn from models.experimental import attempt_load from mainwindow_ui import Ui_MainWindow # from utils.datasets import letterbox from utils.general import check_img_size, non_max_suppression, scale_coords # from utils.plots import plot_one_box2 def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32): # Resize and pad image while meeting stride-multiple constraints shape = img.shape[:2] # current shape [height, width] if isinstance(new_shape, int): new_shape = (new_shape, new_shape) # Scale ratio (new / old) r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) if not scaleup: # only scale down, do not scale up (for better test mAP) r = min(r, 1.0) # Compute padding ratio = r, r # width, height ratios new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding if auto: # minimum rectangle dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding elif scaleFill: # stretch dw, dh = 0.0, 0.0 new_unpad = (new_shape[1], new_shape[0]) ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios dw /= 2 # divide padding into 2 sides dh /= 2 if shape[::-1] != new_unpad: # resize img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border return img, ratio, (dw, dh) def plot_one_box2(x, img, color=None, label=None, line_thickness=3): # Plots one bounding box on image img tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness 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] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) single_info = "Position: (%d, %d), (%d, %d), Obj and Confidence: %s"%(c1[0], c1[1], c2[0], c2[1], label) return single_info class UI_Logic_Window(QtWidgets.QMainWindow): def __init__(self, parent=None): super(UI_Logic_Window, self).__init__(parent) self.initUI() # 初始化界面 def initUI(self): self.setWindowIcon(QIcon("./icon/yolo.png")) # 创建一个窗口对象 self.ui = Ui_MainWindow() self.ui.setupUi(self) self.timer_video = QtCore.QTimer(self) # 创建定时器 self.timer_photo = QtCore.QTimer(self) # 创建定时器 self.output_folder = 'output/' self.cap = cv2.VideoCapture() self.vid_writer = None self.camera_detect = False self.num_stop = 1 # 暂停与播放辅助信号,note:通过奇偶来控制暂停与播放 self.openfile_name_model = None # 权重初始文件名 self.count = 0 self.start_time = time.time() # 打开线程 self.stop_going = 0 # 刷新lcd时间 self.lcd_time = QTimer(self) self.lcd_time.setInterval(1000) self.lcd_time.timeout.connect(self.refresh) self.lcd_time.start() self.ui.textBrowser_print.append("特别说明:如需启动检测,请先加载weights文件!!!") self.init_slots() # 刷新时间 def refresh(self): now_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") self.ui.lcdNumber.display(now_time) # 初始化槽函数 def init_slots(self): self.ui.btn_loadweight.clicked.connect(self.load_model) self.ui.btn_loadimg.clicked.connect(self.button_image_open) self.ui.btn_loadvideo.clicked.connect(self.button_video_open) self.ui.btn_opencamera.clicked.connect(self.button_camera_open) self.ui.btn_camera_detect.clicked.connect(self.button_camera_detect) self.ui.btn_stop.clicked.connect(self.button_stop) self.ui.btn_over.clicked.connect(self.button_over) self.ui.btn_closecamera.clicked.connect(self.button_closecamera) self.ui.btn_clear.clicked.connect(self.button_clear) self.ui.btn_takephoto.clicked.connect(self.button_takephoto) self.ui.btn_labelimg.clicked.connect(self.button_labelimg) self.timer_video.timeout.connect(self.show_video_frame) # 定时器超时,将槽绑定至show_video_frame self.timer_photo.timeout.connect(self.show_image) # 定时器超时,将槽绑定至show_video_frame # 加载模型 def load_model(self): self.openfile_name_model, _ = QtWidgets.QFileDialog.getOpenFileName(self.ui.btn_loadweight, '选择weights文件', 'weights/', "*.pt;;*.pth") if not self.openfile_name_model: QtWidgets.QMessageBox.warning(self, u"Warning", u"打开权重失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok) self.ui.textBrowser_print.append("打开权重失败") else: self.ui.textBrowser_print.append("加载weights文件地址为:" + str(self.openfile_name_model)) self.model_init() #初始化权重 # 初始化权重 def model_init(self): # 模型相关参数配置 parser = argparse.ArgumentParser() parser.add_argument('--weights', nargs='+', type=str, default='weights/yolov5s6.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='cuda', 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 # 改变权重文件 if self.openfile_name_model: weights = self.openfile_name_model print(weights) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 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 # Second-stage classifier # 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] # 设置提示框 QtWidgets.QMessageBox.information(self, u"Notice", u"模型加载完成", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok) self.ui.textBrowser_print.append("模型加载完成") # 打开图片 def button_image_open(self): # 打印信息显示在界面 name_list = [] try: img_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打开图片", "data/images", "*.jpg;;*.png;;All Files(*)") except OSError as reason: print('文件打开出错啦!核对路径是否正确'+ str(reason)) self.ui.textBrowser_print.append("文件打开出错啦!核对路径是否正确") else: # 判断图片是否为空 if not img_name: QtWidgets.QMessageBox.warning(self, u"Warning", u"打开图片失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok) self.ui.textBrowser_print.append("打开图片失败") else: self.ui.textBrowser_print.append("打开图片成功") img = cv2.imread(img_name) print("img_name:", img_name) self.origin = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA) self.origin = cv2.resize(self.origin, (640, 480), interpolation=cv2.INTER_AREA) self.QtImg_origin = QtGui.QImage(self.origin.data, self.origin.shape[1], self.origin.shape[0], QtGui.QImage.Format_RGB32) self.ui.label_origin.setPixmap(QtGui.QPixmap.fromImage(self.QtImg_origin)) self.ui.label_origin.setScaledContents(True) # 设置图像自适应界面大小 info_show = self.detect(name_list, img) # print(info_show) # 获取当前系统时间,作为img文件名 now = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime(time.time())) file_extension = img_name.split('.')[-1] new_filename = now + '.' + file_extension # 获得文件后缀名 file_path = self.output_folder + 'img_output/' + new_filename cv2.imwrite(file_path, img) # 检测信息显示在界面 self.ui.textBrowser_detect.append(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_detect.setPixmap(QtGui.QPixmap.fromImage(self.QtImg)) self.ui.label_detect.setScaledContents(True) # 设置图像自适应界面大小 # 目标检测 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) info_show = info_show + single_info + "\n" return info_show # 打开视频并检测 def button_video_open(self): video_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打开视频", "data/video/", "*.mp4;;*.avi;;All Files(*)") flag = self.cap.open(video_name) # 判断摄像头是否打开 if not flag: QtWidgets.QMessageBox.warning(self, u"Warning", u"打开视频失败", buttons=QtWidgets.QMessageBox.Ok,defaultButton=QtWidgets.QMessageBox.Ok) else: # -------------------------写入视频----------------------------------# self.ui.textBrowser_print.append("打开视频检测") fps, w, h, save_path = self.set_video_name_and_path() self.vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) self.timer_video.start(30) # 以30ms为间隔,启动或重启定时器 # 进行视频识别时,关闭其他按键点击功能 self.ui.btn_loadvideo.setDisabled(True) self.ui.btn_loadimg.setDisabled(True) self.ui.btn_opencamera.setDisabled(True) def set_video_name_and_path(self): # 获取当前系统时间,作为img和video的文件名 now = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime(time.time())) # if vid_cap: # video fps = self.cap.get(cv2.CAP_PROP_FPS) w = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 视频检测结果存储位置 save_path = self.output_folder + 'video_output/' + now + '.mp4' return fps, w, h, save_path # 定义视频帧显示操作 def show_video_frame(self): name_list = [] flag, img = self.cap.read() # 显示视频数据的帧数 self.count += 1 if self.count % 10 == 0: self.count = 0 fps = int(30 / (time.time() - self.start_time)) self.ui.fps_label.setText('fps:' + str(fps)) self.start_time = time.time() if img is not None: # 原始数据的显示 self.origin = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA) self.origin = cv2.resize(self.origin, (640, 480), interpolation=cv2.INTER_AREA) self.QtImg_origin = QtGui.QImage(self.origin.data, self.origin.shape[1], self.origin.shape[0], QtGui.QImage.Format_RGB32) self.ui.label_origin.setPixmap(QtGui.QPixmap.fromImage(self.QtImg_origin)) self.ui.label_origin.setScaledContents(True) # 设置图像自适应界面大小 # 检测数据的显示 info_show = self.detect(name_list, img) # 检测结果写入到原始img上 self.vid_writer.write(img) # 检测结果写入视频 print(info_show) # 检测信息显示在界面 self.ui.textBrowser_detect.append(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_detect.setPixmap(QtGui.QPixmap.fromImage(showImage)) self.ui.label_detect.setScaledContents(True) # 设置图像自适应界面大小 else: self.timer_video.stop() # 读写结束,释放资源 self.cap.release() # 释放video_capture资源 self.vid_writer.release() # 释放video_writer资源 self.ui.label.clear() # 视频帧显示期间,禁用其他检测按键功能 self.ui.btn_loadvideo.setDisabled(True) self.ui.btn_loadimg.setDisabled(True) self.ui.btn_opencamera.setDisabled(True) '''显示图片''' def show_image(self): flag, self.image = self.cap.read() # 从视频流中读取图片 image_show = cv2.resize(self.image, (620, 420)) # 把读到的帧的大小重新设置为显示的窗口大小 width, height = image_show.shape[:2] # 行:宽,列:高 image_show = cv2.cvtColor(image_show, cv2.COLOR_BGR2RGB) # opencv读的通道是BGR,要转成RGB image_show = cv2.flip(image_show, 1) # 水平翻转,因为摄像头拍的是镜像的。 # 把读取到的视频数据变成QImage形式(图片数据、高、宽、RGB颜色空间,三个通道各有2**8=256种颜色) self.photo= QtGui.QImage(image_show.data, height, width, QImage.Format_RGB888) self.ui.label_origin.setPixmap(QPixmap.fromImage(self.photo)) # 往显示视频的Label里显示QImage self.ui.label_origin.setScaledContents(True) # 图片自适应 # 使用摄像头检测 def button_camera_open(self): self.camera_detect = True self.ui.textBrowser_print.append("打开摄像头") # 设置使用的摄像头序号,系统自带为0 camera_num = 0 # 打开摄像头 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: QtWidgets.QMessageBox.information(self, u"Warning", u"打开摄像头成功", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok) self.ui.btn_loadvideo.setDisabled(True) self.ui.btn_loadimg.setDisabled(True) # 启动摄像头检测 def button_camera_detect(self): self.ui.textBrowser_print.append("启动摄像头检测") fps, w, h, save_path = self.set_video_name_and_path() fps = 5 # 控制摄像头检测下的fps,Note:保存的视频,播放速度有点快,我只是粗暴的调整了FPS self.vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) self.timer_video.start(30) self.ui.btn_loadvideo.setDisabled(True) self.ui.btn_loadimg.setDisabled(True) self.ui.btn_opencamera.setDisabled(True) # 视频暂停按钮 def button_stop(self): self.timer_video.blockSignals(False) # 暂停检测 # 若QTimer已经触发,且激活 if self.timer_video.isActive() == True and self.num_stop % 2 == 1: self.ui.btn_stop.setText('继续') self.ui.textBrowser_print.append("视频暂停播放") self.num_stop = self.num_stop + 1 # 调整标记信号为偶数 self.timer_video.blockSignals(True) # 继续检测 else: self.num_stop = self.num_stop + 1 self.ui.btn_stop.setText('暂停') self.ui.textBrowser_print.append("视频继续播放") # 停止视频播放 def button_over(self): self.ui.textBrowser_print.append("视频结束播放") self.cap.release() # 释放video_capture资源 self.timer_video.stop() # 停止读取 self.timer_photo.stop() # 停止读取 if self.vid_writer != None: self.vid_writer.release() # 释放video_writer资源 self.ui.label_origin.clear() # 清空label画布 self.ui.label_detect.clear() # 清空label画布 # 启动其他检测按键功能 self.ui.btn_loadvideo.setDisabled(False) self.ui.btn_loadimg.setDisabled(False) self.ui.btn_opencamera.setDisabled(False) # 结束检测时,查看暂停功能是否复位,将暂停功能恢复至初始状态 # Note:点击暂停之后,num_stop为偶数状态 if self.num_stop % 2 == 0: print("Reset stop/begin!") self.ui.btn_stop.setText(u'暂停') self.num_stop = self.num_stop + 1 self.timer_video.blockSignals(False) # 关闭摄像头 def button_closecamera(self): self.ui.textBrowser_print.append("关闭摄像头") self.ui.fps_label.setText("帧率") self.timer_video.stop() # 停止读取 self.timer_photo.stop() # 停止读取 self.cap.release() # 释放摄像头 self.ui.label_origin.clear() # 清空label画布 self.ui.label_detect.clear() # 清空label画布 self.cap = cv2.VideoCapture(0, cv2.CAP_DSHOW) # 摄像头 self.ui.btn_loadvideo.setDisabled(False) self.ui.btn_loadimg.setDisabled(False) self.ui.btn_opencamera.setDisabled(False) # 拍照 def button_takephoto(self): self.ui.textBrowser_print.append("启动拍照") self.timer_photo.start(30) self.show_image() if self.cap.isOpened(): FName = "data/images" + fr"/img{time.strftime('%Y%m%d%H%M%S', time.localtime())}" print(FName) # 原始数据的显示 flag, self.image = self.cap.read() # 从视频流中读取图片 image_show = cv2.resize(self.image, (640, 480)) # 把读到的帧的大小重新设置为显示的窗口大小 image_show = cv2.cvtColor(image_show, cv2.COLOR_BGR2RGB) # opencv读的通道是BGR,要转成RGB image_show = cv2.flip(image_show, 1) # 水平翻转,因为摄像头拍的是镜像的。 # 把读取到的视频数据变成QImage形式(图片数据、高、宽、RGB颜色空间,三个通道各有2**8=256种颜色) self.showImage = QtGui.QImage(image_show.data, image_show.shape[1], image_show.shape[0], QImage.Format_RGB888) self.ui.label_detect.setPixmap(QtGui.QPixmap.fromImage(self.photo)) self.ui.label_detect.setScaledContents(True) # 设置图像自适应界面大小 self.showImage.save(FName + ".jpg", "JPG", 300) else: QMessageBox.critical(self, '错误', '摄像头未打开!') return None # 调用lablelimg批注工具 def button_labelimg(self): self.ui.textBrowser_print.append("启动标注工具") os.system("labelimg") # 清除显示区域 def button_clear(self): self.ui.textBrowser_print.append("清除显示区域") self.ui.textBrowser_print.clear() self.ui.textBrowser_detect.clear() # 窗口居中 def center(self): qr = self.frameGeometry() cp = QDesktopWidget().availableGeometry().center() qr.moveCenter(cp) # 设置窗口大小 self.move(qr.topLeft()) # 关闭事件 def closeEvent(self, event) -> None: reply = QMessageBox.question(self, 'Message', "Are you sure to quit?", QMessageBox.Yes | QMessageBox.No, QMessageBox.No) if reply == QMessageBox.Yes: if self.cap.isOpened(): self.cap.release() if self.timer_video.isActive(): self.timer_video.stop() if self.timer_photo.isActive(): self.timer_photo.stop() event.accept() else: event.ignore() from qt_material import apply_stylesheet if __name__ == '__main__': # QtCore.QCoreApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling) # 自适应分辨率 app = QtWidgets.QApplication(sys.argv) current_ui = UI_Logic_Window() current_ui.show() extra = { # Font 'font_family': 'Adobe 华文仿宋 Std L', 'font_size': 14, } apply_stylesheet(app, theme='light_blue.xml',extra=extra) # apply_stylesheet(app, theme='dark_teal.xml') sys.exit(app.exec_())
需要先选择权重
图片检测效果图
视频检测效果图
摄像头实时检测效果图
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。