DeepStream5.0系列之yolov5使用

0 背景
本文旨在对 deepstream 中使用 yolov5 的方法做一介绍

测试环境: Ubuntu 18.04, CUDA 10.2, T4, jetpack4.4

软件版本:yolov5:3.0,3.1

项目地址:https://github.com/DanaHan/Yolov5-in-Deepstream-5.0

1 安装流程

1.1 创建 conda 环境
Tips:如果是在 jetson 设备上,可以不用创建 conda 虚拟环境

conda create -n yolov5 python=3.6

1.2 下载源码

mkdir deepstream_yolov5
cd deepstream_yolov5
git clone https://github.com/DanaHan/Yolov5-in-Deepstream-5.0.git
git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git

1.3 安装环境
注意:这块实测 3.0 版本和 3.1 版本的没有问题,最新的 4.0 版本可能会有检测框问题

cd yolov5
git checkout v3.1 #切换到3.1版本
pip install scikit-build
pip install -r requirements.txt

2 模型准备

2.1 下载预训练模型
在 yolov5 界面下载预训练模型(3.1版本模型),这里我们以 YOLOv5x 为例介绍
地址:https://github.com/ultralytics/yolov5
下载好模型后放到 yolov5/weights 目录下

2.2 生成 engine

cd yolov5
cp ../tensorrtx/yolov5/gen_wts.py ./

修改 gen_wts.py 中的模型名称,需要根据自己用的模型做对应修改,主要是第8行和第11行


import torch
import struct
from utils.torch_utils import select_device
 
# Initialize
device = select_device('cpu')
# Load model
model = torch.load('weights/yolov5x.pt', map_location=device)['model'].float()  # load to FP32
model.to(device).eval()
 
f = open('yolov5x.wts', 'w')
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
    vr = v.reshape(-1).cpu().numpy()
    f.write('{} {} '.format(k, len(vr)))
    for vv in vr:
        f.write(' ')
        f.write(struct.pack('>f',float(vv)).hex())
    f.write('\n')


运行后会在本地生成 yolov5x.wts 文件,然后拷贝到 Yolov5-in-Deepstream-5.0 文件夹中

cp yolov5x.wts ../Yolov5-in-Deepstream-5.0
cd ../Yolov5-in-Deepstream-5.0

修改 yolov5.cpp 文件,将 NET 宏改成自己对应的模型

#define NET x // s m l x
然后编译

mkdir build
cd build
cmake ..
make
sudo ./yolov5 -s

运行后会生成 yolov5x.engine 和 libmyplugin.so 文件,首先检查下engine运行结果是否正确

创建测试图片

mkdir …/samples

网上下载两张 coco 数据集图片

sudo ./yolov5 -d  ../samples

输出如下

[11/02/2020-15:08:27] [W] [TRT] Current optimization profile is: 0. Please ensure there are no enqueued operations pending in this context prior to switching profiles
9ms
8ms
并且可以在build目录下生成检测结果图片,格式是_name.jpg,打开图片可以看到检测框正常

验证完毕后,将生成的模型和动态库拷贝到 deepstream 目录中

cp yolov5x.engine ../Deepstream\ 5.0/
cp libmyplugins.so ../Deepstream\ 5.0/

3 deepstream部署
在 Deepstream 5.0/nvdsinfer_custom_impl_Yolo 目录中运行 make 编译,生成 libnvdsinfer_custom_impl_Yolo.so 文件

返回上一级,修改 config_infer_primary_yoloV5.txt 文件,总共有以下几处需要修改

model-engine-file=yolov5s.engine --> model-engine-file=yolov5x.engine
custom-lib-path=objectDetector_Yolo_V5/nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so --> custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so
原工程中缺少 labels.txt 文件,把 coco 数据集的 80 类 labels 拷贝过来即可

然后根据自己的需求,修改 deepstream_app_config_yoloV5.txt 即可

然后载入定制化的动态库,并运行即可

LD_PRELOAD=./libmyplugins.so deepstream-app -c deepstream_app_config_yoloV5.txt

会重新生成 engine 文件,并可以看到运行结果

注意,engine 文件要在 Yolov5-in-Deepstream-5.0 工程中生成,在 tensorrtx 工程中生成 engine 再导过来运行会出现一大堆框

————————————————
版权声明:本文为CSDN博主「ZONG_XP」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/zong596568821xp/article/details/109444343

附相关资源文件(包含labels.txt、yolov5s.wts、libmyplugins.so、yolov5s.engine):
https://download.csdn.net/download/qq_42393859/16061529

相关推荐
©️2020 CSDN 皮肤主题: 程序猿惹谁了 设计师:白松林 返回首页