-
Notifications
You must be signed in to change notification settings - Fork 421
Prometheus目标检测Yolov5 TensorTR与点击跟踪
注意,运行Yolov5-TensorRT需要安装CUDA、TensorRT,如果没有安装,请参考附录1
先看速度测试:
平台 | Yolov5模型 | 输入分辨率 | 检测帧率 |
---|---|---|---|
Nvidia GTX 1080Ti | yolov5s | 608 | 65.2 |
Nvidia Jetson Xavier NX | yolov5s | 608 | 27.1 |
首先安装2个依赖项:
# 安装Opencv
sudo apt-get install libopencv-dev
# 安装pycuda
pip3 install pycuda
在Prometheus目录下,输入:
cd <path-to-prometheus>
chmod +x ./Scripts/install_detection_yolov5tensorrt.sh
./Scripts/install_detection_yolov5tensorrt.sh
或者输入以下命令:
git clone https://gitee.com/jario-jin/prometheus-yolov5-tensorrt.git Modules/object_detection_yolov5tensorrt
cd Modules/object_detection_yolov5tensorrt
pip3 install -r requirements.txt
下载训练好的权重文件,以官方权重yolov5s.pt
为例,用如下代码生成yolov5s.wts
文件
cd <path-to-prometheus>/Modules/object_detection_yolov5tensorrt
python3 gen_wts.py
编译tensorrtx-yolov5:
cd <path-to-prometheus>/Modules/object_detection_yolov5tensorrt
mkdir build
cd build
cmake ..
make
sudo ./yolov5 -s # 生成yolov5s.engine文件
启动yolov5-tensorrt服务端(负责读取摄像头,检测目标)
cd <path-to-prometheus>
chmod +x ./Scripts/start_yolov5tensorrt_server.sh
./Scripts/start_yolov5tensorrt_server.sh
启动yolov5-tensorrt客户端(负责从服务端读取检测结果,并发布ros消息)
roslaunch prometheus_detection yolov5_nvidia_tensorrt.launch
- 点击鼠标左键开始跟踪目标,点击鼠标右键取消跟踪,恢复到检测模式
- 参数详解见附录2
- 如何训练自己的模型见附录3
以CUDA10.0为例,进入CUDA下载官网https://developer.nvidia.com/cuda-10.0-download-archive
,选择Linux→x86_64→Ubuntu→18.04→deb(local)下载.deb格式安装包,然后运行以下指令:
cd Downloads
sudo dpkg -i cuda-repo-ubuntu1804-10-0-local-10.0.130-410.48_1.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-10-0-local-10.0.130-410.48/7fa2af80.pub
sudo apt-get update
sudo apt-get install cuda
如果遇到依赖错误,可采用aptitude安装:
sudo apt-get install aptitude
sudo aptitude install cuda
以Tensorrt7.0.0为例,进入Tensorrt下载官网https://developer.nvidia.com/nvidia-tensorrt-7x-download
,选择Tensorrt 7.0.0.11 for Ubuntu 1804 and CUDA 10.0 DEB local repo packages
下载.deb格式安装包,然后运行以下指令:
sudo dpkg -i nv-tensorrt-repo-ubuntu1804-cuda10.0-trt7.0.0.11-ga-20191216_1-1_amd64.deb
sudo apt-get update
sudo apt-get install tensorrt
以<path-to-Prometheus>/Modules/object_detection/launch/yolov5_nvidia_tensorrt.launch
为例:
<launch>
<node pkg="prometheus_detection" type="yolov5_tensorrt_client.py" name="yolov5_tensorrt_client" output="screen">
<param name="output_topic" value="/prometheus/object_detection/yolov5_tensorrt_det"/>
<param name="camera_parameters" type="string" value="$(find prometheus_detection)/shell/calib_webcam_640x480.yaml" />
<param name="object_names_txt" value="coco"/>
</node>
</launch>
其中:
-
output_topic:检测结果输出话题(消息类型:
MultiDetectionInfo.msg
) - camera_parameters:相机参数文件(为了估计视线角误差、目标位置)
-
object_names_txt:目标类别描述
txt
问题(具体见:<path-to-Prometheus>/Modules/object_detection/py_nodes/yolov5_tensorrt_client/class_desc/coco.txt
)
默认输出话题,消息类型:
##`MultiDetectionInfo.msg`
Header header
## 检测到的目标数量
int32 num_objs
## 每个目标的检测结果
DetectionInfo[] detection_infos
##`DetectionInfo.msg`
## 是否检测到目标
bool detected
## 目标类别名称
string object_name
## 类别ID
int32 category
## 0表示相机坐标系, 1表示机体坐标系, 2表示惯性坐标系
int32 frame
## 目标位置[相机系下:右方x为正,下方y为正,前方z为正]
float32[3] position
## 目标姿态-欧拉角-(z,y,x)
float32[3] attitude
## 目标姿态-四元数-(qx,qy,qz,qw)
float32[4] attitude_q
## 视线角度[相机系下:右方x角度为正,下方y角度为正]
float32[2] sight_angle
## 偏航角误差
float32 yaw_error
注意:默认情况下每个目标仅有detected
, object_name
, category
, frame
, sight_angle
的输出。
如果想输出position
,需要在类别描述文件(如<path-to-Prometheus>/Modules/object_detection/py_nodes/yolov5_tensorrt_client/class_desc/coco.txt
)中填写目标的高度与宽度(单位:m
)。
如上图,以人(person)为例,设置宽度
0.5m
,高度1.8m
然后,需要修改源代码,这里以高度估计人距离摄像机的距离,并以此估计position
。
源代码位置:<path-to-Prometheus>/Modules/object_detection/py_nodes/yolov5_tensorrt_client/yolov5_tensorrt_client.py
,(112-114
行)
cls==0
用来判断是否为类别—person,cls_hs[cls]
用来读取目标高度(读到的数就是我们写在coco.txt
里的1.8m
),camera_matrix[1][1]
为垂直方向像素焦距(由相机标定参数文件决定),h
物体的像素高度(为实时检测结果,并被归一化到0-1)
下载数据集标注工具,下载地址: Spire Web或者百度网盘 (密码: l9e7) ,数据集管理软件SpireImageTools:gitee地址或者github地址。
-
解压,打开标注软件
SpireImageTools_x.x.x.exe
首先点击Tools->Setting...
,填写一个 save path
(所有的标注文件都会存储在这个文件夹中)
-
将拍摄的视频转为图像 (如果采集的是图像,则跳过这一步骤),点击
Input->Video
, 选择要标注的视频。
然后,点击Tools->Video to Image
点击OK 后,等待完成,结果会存储在:
-
打开需要标注的图像,点击菜单
Input->Image Dir
, 找到需要标注的图像所在文件夹 ,按Ctrl+A
,全选,打开所有图像:
-
点击菜单:
Tools->Annotate Image->Box Label
,开始标注图像
在label
中填写待标注目标名称,然后将对话框拖到一边。
-
开始标注,在主窗口中开始标注,鼠标滚轮放大缩小图像,按住左键移动可视图像区域不断点击左键将目标框包围,使用
Yolo
训练时,点击2个点即可:
标注时,如果点错,按鼠标右键可以取消。标注完成后,如果不满意,可以点击绿色边框(边框会变红,如下图所示),按
Delete
删除
- 继续标注行人类别:
-
全部标注完成后,将标注输出为
Yolo
格式,准备训练——在标注完成之后,按下Ctrl+o
点击OK
即可,需要等待转换。
- 注意,如下两个文件夹是我们训练
Yolov5
需要的
在准备好scaled_images
和Yolo_labels
两个文件夹之后,我们就可以训练Yolov5
了。首先,创建一个car_person.yaml
,将其放到<path-to-Prometheus>/Modules/object_detection_yolov5tensorrt/data/
文件夹下。car_person.yaml
的具体内容如下:
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: data/car_person/images/train/
val: data/car_person/images/train/
# number of classes
nc: 2
# class names
names: ['car', 'person']
注意1:
car_person
是自定义名称,我们这次标注的数据集仅有这2个类别。 注意2:names: ['car', 'person']
这里的类别顺序需要跟Yolo_categories.names
里的类别顺序一致。
- 将训练图像与标注拷贝到对应位置
首先,在<path-to-Prometheus>/Modules/object_detection_yolov5tensorrt/data/
下新建一个文件夹car_person
。然后,在car_person
下再新建2个文件夹images
和labels
。最后,将准备好的scaled_images
拷贝到images
下,并重命名为train
;将准备好的Yolo_labels
拷贝到labels
下,并重命名为train
。
结合
car_person.yaml
里的内容,我想你应该明白上面目录结构的含义啦。
- 开始训练
cd <path-to-Prometheus>/Modules/object_detection_yolov5tensorrt/
python3 train.py --img 640 --batch 16 --epochs 5 --data data/car_person.yaml --weights weights/yolov5s.pt
显示以上内容说明训练成功!可以增加训练期数(--epochs 5
)提升效果。
- 部署训练好的模型
刚刚训练好的模型会保存在<path-to-Prometheus>/Modules/object_detection_yolov5tensorrt/runs/exp?/weights/best.pt
,?
需根据自己的情况而定(最新训练的模型?
为最大的数字),将best.pt
重命名为yolov5s.pt
,拷贝到<path-to-Prometheus>/Modules/object_detection_yolov5tensorrt/
下,然后执行:
cd <path-to-prometheus>/Modules/object_detection_yolov5tensorrt
python3 gen_wts.py
感谢使用Prometheus自主无人机软件平台!