使用AidLite推理Yolov5人体检测案例演示

使用AidLite推理Yolov5人体检测完整项目代码包,点此下载

1. 人体检测算法模型简介

Yolov5人体检测模型,即Yolov5/models文件夹中的Yolov5n_person.pt。

人体检测模型,主要采用人体检测数据集训练的Yolov5检测模型,数据集中类别数为1。

使用Labelimg打开的标注信息

2. pt模型转换成tflite模型

TensorFlow Lite 支持多种硬件加速器,AidLux系统中模型推理支持tflite_gpu加速的方式,因此需要在PC端,将Yolov5人体检测的pt模型转换成tflite的格式。

因为export.py模型转换需要tensorflow,所以在终端输入:

pip3 install tensorflow -i https://pypi.tuna.tsinghua.edu.cn/simple,下载tensorflow库。

模型转换的代码,位于Yolov5/export.py,而模型转换的过程可以在PC端本地进行。

主要修改export.py中的三个地方:

data:即Yolov5中人体的yaml配置文件。

weights:即训练好的人体检测pt模型,代码中已提供,也可根据情况替换。

include:因为转换成tflite格式,所以这一步修改成tflite。

运行export.py文件,在Yolov5/models文件夹下面,可以看到生成的Yolov5n_person-fp16.tflite文件。

通过网站页面的方式,将转换好的Yolov5n_person-fp16.tflite文件上传到AidLux系统的/Yolov5/models文件夹中(代码中已上传,如有新的模型,请自行更新)。

3. 查看tflite模型相关信息

在AidLux对Yolov5模型推理时,因为需要使用模型输入输出的参数,所以可以通过Netron查看相关的模型参数。

Netron的网站:https://netron.app/

点击“Open Model”,打开前面转换好的Yolov5n_person-fp16.tflite文件。

点击最下方的输出单元,可以看到输出的信息。

输入的参数:1,640,640,3

输出的参数:1,25200,6

4. Yolov5人体检测代码理解

Yolov5人体检测的代码,即Yolov5/aidlux文件夹的yolov5_person.py。

其中包含了很多AidLux专属的函数接口,开发者可以在AidLite编程接口(AI),查看相关的函数说明。

整体的代码主要分为三个部分:

1 )加载相关的函数库

2 )模型初始化及加载

其中主要用到两个函数接口,一个是aidlite_gpu.aidlite()和aidlite.ANNMode()。

其中的in_shape,out_shape,即前面通过netron查看的模型参数。

需要注意的是,两行代码的最后都有一个4,因为模型参数格式为float格式,所占的字节数为4。

3 )模型输入输出和推理

使用setInput_Float32()函数输入数据,然后通过invoke()函数开启推理,最后使用getOutput_Float()获取推理结果。

5. Yolov5视频推理测试

运行yolov5_person.py进行视频推理测试,在手机版本的Aidlux和PC端网页的Aidlux中,都可以看到推理的显示结果。