章毅恒 陈育罕 张雪秋
摘 要:目标检测数据集中的样本数量及差异性对网络检测性能具有重要影响。通过人工标注图像的方式生成数据集耗时又费力,因此设计一款目标检测数据集自动采集软件。该软件可以检测RTSP视频流、图片文件及视频文件,通过预训练好的轻量化检测网络得到目标位置类别及定位信息,生成标注文件并保存到本地。实验结果表明,该文软件在不同的平台上均可运行,其数据集采集速度相比人工标注有显著优势。
关键词:数据集制作;PyQt;目标检测;多线程;YOLOv5
中图分类号:TP311.5 文献标识码:A 文章编号:2096-4706(2023)15-0012-05
Design and Implementation of Target Detection Dataset Automatic Collection Software Based on PyQt
ZHANG Yiheng, CHEN Yuhan, ZHANG Xueqiu
(College of Mechanical Engineering, Chongqing University of Technology, Chongqing 400054, China)
Abstract: The number and diversity of samples in the target detection dataset have a significant impact on network detection performance. Generating a dataset by manually annotating images is time-consuming and laborious, so that a software for automatic collection of object detection datasets is designed. This software can detect RTSP video streams, image files, and video files, obtain target location categories and positioning information through a pre trained lightweight detection network, generate annotation files, and save them locally. The experimental results show that the software can run on different platforms, and its dataset collection speed has significant advantages compared to manual annotation.
Keywords: dataset making; PyQt; target detection; multithreading; YOLOv5
0 引 言
卷积神经网络本质上是一种输入到输出的映射,训练时神经网络通过大量输入与输出的关系不断修正参数,使得网络可以得到输入与输出之间的映射关系。样本数量越多,网络得到的映射关系越精确。样本之间的差异越大,网络的泛化能力就越强。因此,输入与输出的样本数据集合是决定目标检测网络精度的重要因素。
人工制作数据集需要在图片中找出目标的种类与位置,并标注出详细数据,这个过程需要消耗巨大的人力成本。使用训练好的网络检测并生成数据集大大减少了标注图片形成数据集的人工成本,有助于提升深度学习的研究效率。
为此,本文设计一款实时目标检测与数据集采集标注系统,使用PyQt平台进行开发,采用PyTorch、和PyTorch等库实现功能。通过预训练好的目标检测网络自动检测图像,同时采集图像目标位置类别与定位信息,生成txt格式的标注文件。为了能够在没有安装GPU或配置深度学习环境的计算机上仍可使用本软件,采用预训练好的轻量化YOLOv5网络作为检测网络,可以检测COCO数据集包含的80个目标种类。为方便操作人员制作数据集,将数据集自动采集程序设计为可以根据选定的采样频率采集检测网络输出的图像与预测信息,并保存到自定义的文件夹中。为了更加便捷地增加样本的差异性,将软件设计为可以读取RTSP视频流、本地图片和本地视频三种图像。
实验证明,本文设计开发的数据集采集软件运行稳定,检测过程流畅,可在操作系统为Win 10及以上版本、CPU为x86_64架构的计算机中运行。相比于人工标注数据集,使用本软件标注数据集可节约90%的人力成本,便于相关研究人员快速制作数据集。
1 软件界面设计与功能设计
软件界面如图1所示,界面分为RTSP流操作区域、本地文件操作区域、数据集采样操作区域、检测操作区域以及图像显示区域六个区域。RTSP流操作区域提供输入RTSP流地址、获取RTSP流图像和停止获取RTSP流图像等功能。本地文件区域提供选取本地文件并显示文件类型和地址等功能。采样操作区域提供单帧手动采集、选择样本保存地址、调节采集间隔帧数、启动自动采集、停止自动采集等功能。检测区域提供置信度调节、开始停止检测等功能。
为了增加软件的灵活性,设定六种输入软件的内容,分别为RTSP流地址、视频文件地址、图片文件地址、检测网络置信度、数据集采样频率、数据集保存位置。RTSP流为输入的字符串,由使用人员从输入框自行输入。可从0至1的范围内调节置信度。数据集自动采样频率由间隔帧数确定,可自行输入,也可通过按钮以每10帧为一单位进行调节。除了自动采集数据集,还可通过按钮手动采集单帧图像及其标注信息。检测网络输入图像有三种来源,分别为RTSP流实时画面、本地图片文件、本地视频文件。多种输入可以丰富数据集的来源,增加数据集的多样性和规模。设定四种队列(分别为视频流队列、检测-输入队列、检测-输出队列、检测-目标信息队列)来完成不同模块线程间的信息交互,使程序运行更加流畅。设定4种工作状态,分别为初始狀态、播放视频流状态、显示本地文件状态、检测状态。执行不同操作前先判断软件的工作状态,避免出现数据混乱和程序卡死的问题。
2 功能模块开发
采用vscode编辑器+PyQt-Integration插件+Python插件作为集成开发环境。使用PyQt-Integration插件将.ui文件转换为.py格式的UI声明文件,通过UI声明文件给UI中的按钮、标签等槽函数绑定各种功能。各模块类继承PyQt的QTherad线程类以实现多线程,线程间通过队列通信,队列继承自Python内部的queue类。软件结构如图2所示。
2.1 UI模块
使用Qt Designer开发程序界面,使用Qlabel元素作为展示视频流和图片/视频文件的承载元素,使用QspinBox元素调节参数,使用QListWidget元素展示统计结果,操作按钮为PushButton。图像读取路径和保存地址通过QlineEdit输入。
2.2 视频流读取模块
调用中的videocapture函数来读取RTSP视频流的图像。由于VideoCapture函数读取的图像会自行开辟缓存,输出图像数据时不一定是最新帧,造成掉帧和花屏的现象。因此,单独开辟一个线程,将videocapture每次读取的图像输入到视频流队列中。后续模块调用时从队列中读取图像,避免卡顿现象的发生。
2.3 本地文件读取模块
如果本地文件地址为图片,读取完后会在原图区域显示图片。如果本地文件地址为视频,则显示视频的第一帧画面。新编写文件读取函数,解决了原生图片及视频读取函数无法读取中文路径以及只能读取8位深图像的问题。
2.4 图像检测模块
图像检测模块负责处理目标检测网络的输入输出。如果软件正在读取RTSP流,那么检测网络的输入为RTSP流最新帧图像。如果未获取或停止获取RTSP视频流数据,则检测网络输入为本地文件,若本地文件地址为图片或视频,则检测网络输入为选定的图片或视频。若本地文件地址为文件夹,则检测该文件夹下所有图片。检测时,复制输入图像并将其推入检测-输入队列,将检测结果推入检测-输出队列,将预测框信息、目标统计信息等推入检测-目标信息队列,从而保证检测网络的输入输出与目标信息同步,避免播放与采集时出现数据不匹配的问题。
2.5 数据显示模块
数据显示模块负责将原始图像、检测网络输出图像和检测统计结果展示在界面上。当软件处于检测文件或检测视频流状态时,数据显示线程从检测-输入队列读取原图并填入UI界面的原图显示区域,从检测-输出队列读取检测完成的图片填入检测结果图片显示区域,从检测-目标信息队列读取目标种类与数量统计结果填入统计区域。为了使图像以原比例保存,编写了图像尺寸自适应函数,通过填充边缘的方式使图像有效部分的比例不变。由于读取图像的色彩顺序为BGR,重新编写图像读取函数,将图像色彩顺序调整为RGB。
2.6 数据集采集模块
数据集采集模块负责将被检测图像及其检测出的目标信息保存到本地,以供制作数据集。采集线程设置一个计数器,当计数器数量与采集间隔帧数相同时,采集线程读取目标信息队列中最新的目标信息,将当前所检测图像帧中每一个目标的中心点位置、高度、宽度和种类等标注信息保存为本地txt文件,保存位置由使用者自行选择。同时,读取当前检测结果图像队列的最新一帧图像为本地图片文件,保存的图像与标注信息命名相同。采集完毕后,计数器清零,重新开始计数。
3 检测程序
为了快速且准确地检测图像,采用预训练好的基于PyQt的目标检测网络获取图像中的目标信息。
3.1 轻量化检测网络
YOLO系列网络是目标检测领域常用的一种端到端深度网络,该网络将检测问题转化为回归问题,可以在预测目标中心位置与宽高的同时,预测目标的类别。与R-CNN系列的两级网络相比,YOLO网络无需大规模生成先验框,从而大幅减少参数,提升运行速度。本文所设计的软件需要检测视频流,对检测速度提出了更高的要求,因此YOLO网络更适合作为本文检测网络的基础。
深度学习网络训练与检测时,往往使用GPU硬件进行计算以加快速度。为了确保本文所设计的系统在缺乏GPU平台的计算机上能够移植且保持检测的实时性,使用YOLOv5网络的轻量化版本YOLOv5-lite作为目标检测网络,使得网络在只有CPU的计算机上仍能实时检测目标。
YOLOv5-lite网络结构如图3所示。YOLOv5-lite继承了YOLOv5的核心思想,但相比于对性能要求更高的YOLOv5s,YOLOv5-lite做了如下轻量化处理,使得网络在CPU设备上仍可流畅运行。
首先,YOLOv5-lite去除了Focus层,以Conv-BatchNorm-Relu结构取而代之,从而避免多次slice操作导致的缓存过度占用;使用ShuffleNetv2网络取代原YOLOv5s网络的Backbone层,相比于原Backbone层大量使用的C3模块,Shuffle_Block操作对内存访问量要求更低。为进一步降低内存占用与运算量,将原始YOLOv5s里PAN结构中的Concat操作更改为add操作。
3.2 网络训练
训练所用计算机的配置:CPU为Intel(R) Xeon(R) Silver 4210R,运行内存为32 GB,GPU为GTX3060ti,显存为8 GB。训练环境为Python 3.8.16、PyTorch 1.12、cuda 11.7。为提高预训练网络的特征提取能力,在COCO数据集中选取9 000张图片作为训练集。训练参数设置:batch大小為30,epoch数量为300,初始学习率为0.01。
3.3 网络实验与验证
为了验证网络在较低性能CPU平台上的运行速度与精度,在Intel i7-6700k处理器上测试网络的速度与精度。速度指标为FPS,检测精度指标为mAP@0.5,表示各类别检测精度的平均值。
验证数据集为COCO数据集中未经训练的1 000张图片,同时横向对比其他若干检测网络。实验结果如表1所示。
实验结果表明,YOLOv5-lite网络在CPU上的检测速度为28.1 f/s,远远高于原始YOLOv5s网络和fast-rcnn网络,满足实时性要求,同时又远高于YOLOv4-tiny的检测精度。故YOLOv5-lite网络适合作为本软件的目标检测网络。
4 程序打包与功能验证
4.1 程序打包
运行本软件所需的依赖库如表2所示。
为了避免对此软件无用的库也被打包进可执行程序而带来程序体积增大的问题,使用Anconda新建一个虚拟环境,安装表2所列出的依赖库。在编写程序时,尽量按需引入功能模块,减小打包体积。为了使程序在无依赖库与Python环境的计算中仍可运行,使用Pyinstaller将Python文件打包成Windows平台可执行程序(.exe文件)。
4.2 跨平台运行验证
为了验证软件的可移植性,在不同的平台上测试软件的运行。从使用场景的角度出发,只考虑软件在操作系统为Win10及以上、CPU位数为64位x86_64架构的计算机上运行。测试用视频流来自海康威视摄像头GH19L,輸出图像尺寸为1 920×1 080,码率为25 f/s。在部分计算机上测试各功能是否能正常运行,实验结果如表3所示。
验证结果表明,本文软件具有良好的可移植性,在用于测试的计算机上均能流畅运行。
4.3 标注效率验证
在COCO数据集中随机选取未经检测网络训练的500张图片,使用摄像机自行采集并人工标注500张图片,共计7 360个目标,构成验证数据集,用于对比使用软件标注的优劣。软件正确标注了6 356个目标,错误标注了371个目标,漏检了633个目标,漏检目标多为成像尺寸较小的目标。相较于人工标注,使用本文软件标注能够节约86%的人工成本,充分验证了本文软件的实用价值。
5 结 论
本文提出一款可快速制作目标检测数据集的软件。通过预训练网络检测RTSP视频流或本地文件,得到图片中的目标种类与位置信息,并保存到本地文件夹中,从而自动生成数据集。软件可以自选输入,检测置信度,采样频率和数据集保存位置等信息。测试结果表明,本文软件具有良好的可移植性,可大幅减轻数据集制作的工作量。
参考文献:
[1] 高洪波.煤矿应急救援指挥信息系统的设计与实现 [J].现代信息科技,2022,6(21):105-108+111.
[2] 吴春波,陈伟,赵振华,等.基于PyQt的SHPB试验数据处理程序开发 [J].计算机应用与软件,2021,38(12):11-13+80.
[3] 刘宏伟.基于Qt的多通道振动信号采集仪上位机软件设计 [J].现代信息科技,2023,7(3):24-28+32.
[4] 陆安琴,秦婵婵,胡圣波,等.基于PyQt5的数字图像处理实验平台设计 [J].电子元器件与信息技术,2021,5(2):158-161.
[5] 陈富荣,肖明明.基于YOLOv5的改进小目标检测算法研究 [J].现代信息科技,2023,7(3):55-60+65.
[6] ZHANG X Y,ZHOU X Y,LIN M X,et al. ShuffleNet: An Extr-emely Efficient Convolutional Neural Network for Mobile Devices [C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.Salt Lake City:IEEE,2018:6848-6856.
[7] MA N N,ZHANG X Y,ZHENG H T,et al. Shufflenetv2: Practical guidelines for efficient CNN architecture design [C]//Proceedings of the European Conference on Computer Vision. Munich:Springer,2018:122-138.