1,337
社区成员
发帖
与我相关
我的任务
分享我的NVIDIA开发者之旅——Jetson平台上基于YOLOV5-TRT-DeepSort目标检测跟踪应用
前言
车流量统计、人流量统计等在实际场景中实用性很强，我们通过判断某区域内的车流人流量，来做一些决策。一般而言，这类任务通常需要一定的实时性，需要边缘计算而不是将视频流通过高延时网络发送到服务器端，然后进行计算。边缘端负责将预警事件上报或者结合后续处理自动进行控制。这样我们就需要在边缘端实现目标检测和跟踪，其中跟踪主要用于计数。
在Jetson平台上，我们可以通过DeepStream来进行检测跟踪计数，NVIDIA已经针对DeepStream做了大量优化，因此使用DeepStream来部署此类任务通常较为合适。但是考虑到DeepStream并不能满足所有的场景需求，并且对我来说，DeepStream也太厚重了，因此我决定用TensorRT+DeepSort实现这一功能。
深度学习模型选择
对于目标检测而言，受限于jetson平台算力，我们选择一些较为轻量级的模型，同时受限于监控镜头的视角、画质、焦段，又要尽可能降低误检、漏检率，因此我们选择使用YOLOV5作为目标检测网络。加之较为熟悉YOLOV5的性能，在Jetson Nano上YOLOV5S模型使用TensorRT加速后，推理大概耗时90ms，在Jetson NX Xavier上推理大概耗时15ms。选用YOLOV5模型，实时性可以得到保证，同时模型精度较高，误检漏检在实际测试中都很低。
数据集制作
根据公路上的实际情况，将待检测的目标分为：行人、自行车、电动自行车、小汽车、公交车、卡车6个类别，分别准备200张图像，并使用lableImg软件对图像进行标注，完成数据集制作。
模型训练
从YoloV5项目的github主页拉取YoloV5项目.
git clone https://github.com/ultralytics/yolov5
修改/data/coco128.yaml文件中数据路径与nc与names数据，并注释掉最下面的download。数据路径要保证与数据集存放路径对应.
path: ../datasets/mydata # dataset path
train: images/train # train images
val: images/val # val images
test: # test images
nc: 6
names:[ 'person', 'bicycle', 'elecBicycle','car', 'bus', 'truck']
也可以自己新建yaml文件，训练时通过参数指定即可。我是比较懒，直接改了默认的yaml文件。
修改/models/yolov5s.yaml文件中的nc，改成对应的6类。
新建目录nvidia，并将mydata目录与yolov5目录放入到nvidia下。
mkdir nvidia && mv /mydata /nvidia && mv /yolov5 /nvidia
这里路径要对应自己的真实路径。
从NGC官网下载pytorch镜像：
进入NGC主页，登录NGC账号，在NGC页面点击containers，搜索pytorch，点击进入，注意L4T那个是给jetson设备用的，不适用于X86架构。
选择版本，我用的是nvcr.io/nvidia/pytorch:21.11-py3，复制指令，在终端执行：
docker pull nvcr.io/nvidia/pytorch:21.11-py3
启动容器：
docker run -it --rm -gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v /nvidia:/workspace/nvidia nvcr.io/nvidia/pytorch:21.11-py3 /bin/bash
其中--rm指定容器推出后自动销毁。
进入容器后，执行:
python3 train.py
进入网络训练，训练完成后得到pt文件。
TRT加速
YOLOv5使用TensorRT加速一般有两种方式，第一种为先将pt文件导出为onnx文件，然后利用TensorRT的onnx parser或者trtexec解析onnx文件，序列化为engine，之后加载、反序列化。在pytorch容器内执行：
python3 export.py
导出onnx文件，我们需要新建TensorRT容器，执行：
docker run -it --rm -gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v /nvidia:/workspace/nvidia nvcr.io/nvidia/tensorrt:22.03-py3 /bin/bash
进入TensorRT容器，执行：
trtexec --onnx=yolov5m.pt --saveEngine=yolov5m.engine --fp32
由于我们设计好了，pytorch容器和TensorRT容器挂载了同一目录，因此我们可以用上述命令导出yolov5的engine文件。
第二种为使用Github上开源的tensorrtx项目里的yolov5，此项目实现了诸多网络的TensorRT版本，包括yolov5。执行
git clone https://github.com/wang-xinyu/tensorrtx.git
将项目下载到本地yolov5目录下。
此方法需要对代码进行编译，得到so文件。
执行：
cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5
将gen_wts.py文件放到yolov5目录下，此文件的作用是将训练好的pt文件转换为wts文件。
在pytorch容器中执行：
python3 gen_wts.py -w yolov5m.pt -o yolov5m.wts
获得wts文件。
进入TensorRT容器中，执行：
cd {tensorrtx}/yolov5/
vim yololayer.h
修改代码中的类别数：
static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000;
static constexpr int CLASS_NUM = 6; //将coco128的80类改为6
static constexpr int INPUT_H = 640;
static constexpr int INPUT_W = 640; //网络输入宽高
继续执行:
mkdir build && \
cd build && \
cp {ultralytics}/yolov5/yolov5m.wts {tensorrtx}/yolov5/build && \
cmake .. && \
make
生成yolov5可执行文件和libmyplugins.so库文件，我们用python进行TensorRT推理时，只需要这个libmyplugins.so库文件。
接下来，我们将wts文件转换为序列化的engine。执行：
./yolov5 -s [.wts] [.engine] [n/s/m/l/x/n6/s6/m6/l6/x6 or c/c6 gd gw]
如:
./yolov5 -s yolov5m.wts yolov5m.engine m
完成后，我们可以测试下转换后的推理效果，执行：
sudo ./yolov5 -d [.engine] [image folder]
下面我们解释下两种转换方法的不同之处：第一种方法是NVIDIA推荐的方法，先导出一个onnx文件，然后用onnx parser解析它，重新构建TensorRT网络，如果有不能支持的层，我们用onnxGraphSurgeon修改onnx网络，直到TensorRT能够完整构建出整个网络，同时也能使用Polygraphy对比逐层输出，从而灵活调整各层转换为TensorRT层后的精度。这种方式针对YoloV5网络来说的不足指出是，相较于第二种，后期需要处理大量nms的运算，在Jetson NX上，需要占用10~15ms的时间。YoloV5 v6.1中集成的TensorRT加速就是采用的这种方法。
第二种方式其实是wang-xinyu等人使用TensorRT API重新一层层构建了YoloV5网络，重新设计了nms的后续处理，节省了大量nms时间。不足之处是使用起来比较麻烦。
因为在Jetson上性能相对台式机显卡而言比较弱，因此我们采用第二种方式进行模型转换，使用TensorRT。
实际部署中，tensorrtx/yolov5给了一个使用python进行推理的模板：yolov5_trt.py，我们基于此修改我们自己的TensorRT推理代码。
我们首先修改主函数部分：
if __name__ == "__main__":
# load custom plugin and engine
PLUGIN_LIBRARY = "build/libmyplugins.so"
engine_file_path = "build/yolov5s.engine"
ctypes.CDLL(PLUGIN_LIBRARY)
这几行代码用来载入plugin和engine，此时engine仅仅只是一个文件名，实际并未加载，也未被反序列化。
然后是修改类别名，由于我们并不用检测部分代码画框，所以这部分其实可以不做，但如果只是用来进行检测推理，labels这里需要根据实际情况进行相应修改。
yolov5_wrapper = YoLov5TRT(engine_file_path)
模型初始化
这行代码对YoloV5TRT进行初始化，我们进入这个初始化过程。
class YoLov5TRT(object):
"""
description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops.
"""
def __init__(self, engine_file_path):
# Create a Context on this device,
self.ctx = cuda.Device(0).make_context()
stream = cuda.Stream()
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(TRT_LOGGER)
# Deserialize the engine from file
with open(engine_file_path, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
host_inputs = []
cuda_inputs = []
host_outputs = []
cuda_outputs = []
bindings = []
for binding in engine:
print('bingding:', binding, engine.get_binding_shape(binding))
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
cuda_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(cuda_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
self.input_w = engine.get_binding_shape(binding)[-1]
self.input_h = engine.get_binding_shape(binding)[-2]
host_inputs.append(host_mem)
cuda_inputs.append(cuda_mem)
else:
host_outputs.append(host_mem)
cuda_outputs.append(cuda_mem)
# Store
self.stream = stream
self.context = context
self.engine = engine
self.host_inputs = host_inputs
self.cuda_inputs = cuda_inputs
self.host_outputs = host_outputs
self.cuda_outputs = cuda_outputs
self.bindings = bindings
self.batch_size = engine.max_batch_size
以上的初始化代码，我们可以看到使用TensorRT推理的一般步骤，首先调用
TRT_LOGGER = trt.Logger(trt.Logger.INFO)
产生一个logger，然后用这个logger，构建runtime,
runtime = trt.Runtime(TRT_LOGGER)
再使用runtime，进行反序列化：
with open(engine_file_path, "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
engine文件反序列化完成之后，我们构建一个context，这个context不是cuda里的context，它除了包含了CUDA context以外还包含了trt logger、trt engine指针等内容。
context = engine.create_execution_context()
之后我们执行for循环，获取engine的输入shape、输出shape、开辟内存、显存空间，从而构建一个bindings，这个bindings其实就是一段显存空间，包含了网络需要的输入和输出，它通过host_inputs、cuda_inputs、host_outputs、cuda_outputs来构建输入图片和输出推理结果与TensorRT的关联。
由于在Jetson上，CPU和GPU共用一块Memory，因此我们可以把host_outputs、cuda_outputs设置成同一块空间，这样我们可以避免一次Memory Copy，节省一点时间。由于时间关系，这个我们先放一边，改天再写这个教程。
执行推理
成功构建好bindings后，我们执行推理，推理主要调用infer函数：
def infer(self, raw_image_generator):
threading.Thread.__init__(self)
# Make self the active context, pushing it on top of the context stack.
self.ctx.push()
# Restore
stream = self.stream
context = self.context
engine = self.engine
host_inputs = self.host_inputs
cuda_inputs = self.cuda_inputs
host_outputs = self.host_outputs
cuda_outputs = self.cuda_outputs
bindings = self.bindings
# Do image preprocess
batch_image_raw = []
batch_origin_h = []
batch_origin_w = []
batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w])
for i, image_raw in enumerate(raw_image_generator):
input_image, image_raw, origin_h, origin_w = self.preprocess_image(image_raw)
batch_image_raw.append(image_raw)
batch_origin_h.append(origin_h)
batch_origin_w.append(origin_w)
np.copyto(batch_input_image[i], input_image)
batch_input_image = np.ascontiguousarray(batch_input_image)
# Copy input image to host buffer
np.copyto(host_inputs[0], batch_input_image.ravel())
start = time.time()
# Transfer input data to the GPU.
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
# Run inference.
context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
# Synchronize the stream
stream.synchronize()
end = time.time()
# Remove any context from the top of the context stack, deactivating it.
self.ctx.pop()
# Here we use the first row of output in that batch_size = 1
output = host_outputs[0]
因为我们只需要一张图片，因此不需要使用它的get_img_path_batches，这个是将几张图合并成一个generator，提高推理的吞吐量，我们只有一路视频，使用batches反而会增加不少额外处理，因此我们不使用generator，我们直接传图像进去。
def infer(self, img):
threading.Thread.__init__(self)
# Make self the active context, pushing it on top of the context stack.
self.ctx.push()
# Restore
stream = self.stream
context = self.context
engine = self.engine
host_inputs = self.host_inputs
cuda_inputs = self.cuda_inputs
host_outputs = self.host_outputs
cuda_outputs = self.cuda_outputs
bindings = self.bindings
input_image1 = np.empty(shape=[1, 3, self.input_h, self.input_w])
input_image, image_raw, origin_h, origin_w = self.preprocess_image(input_image1)
input_image = np.ascontiguousarray(input_image)
# Copy input image to host buffer
np.copyto(host_inputs[0], input_image.ravel())
start = time.time()
# Transfer input data to the GPU.
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
# Run inference.
context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
# Synchronize the stream
stream.synchronize()
end = time.time()
# Remove any context from the top of the context stack, deactivating it.
self.ctx.pop()
# Here we use the first row of output in that batch_size = 1
output = host_outputs[0]
传入图像以后，
batch_image_raw = []
batch_origin_h = []
batch_origin_w = []
这三行不再需要.
batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w])
中，self.batch_size改为1：
input_image1 = np.empty(shape=[1, 3, self.input_h, self.input_w])
for i, image_raw in enumerate(raw_image_generator):
不再需要执行，只需要对图像执行process_image()即可：
input_image, image_raw, origin_h, origin_w = self.preprocess_image(image)
preprocess_image()的作用是把图像长边缩放到网络输入的长边，然后对短边进行加黑边到网络输入的短边，对我们的网络来说，是缩放至640*640。
下一步将输入图像Copy到host buffer里：
np.copyto(host_inputs[0], image.ravel())
接着将host buffer Copy至GPU
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
此函数是个异步函数，只是发布任务。
执行推理：
context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
此函数也是异步函数，bindings即为初始化时创建的bindings，memcpy_htod_async函数已经将图像数据copy到bindings里了，推理的结果也会写入bindings的output空间。
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
执行memcpy_dtoh_async()将输出的结果，Copy回CPU内存空间，我们的CPU代码就可以拿到结果，进行nms之类的后续操作了。
因为前面三个函数都是异步函数，因此我们需要加同步函数，确保CPU拿到想要的数据，调用：
stream.synchronize()
进行同步。
我们拿到output之后，调用post_process进行后处理，tensorrtx版的yolov5舍弃了大量候选框，使得后续处理耗时很小，节省了不少时间。
我们就可以拿到result_boxes, result_scores, result_classid，拿到目标坐标、置信度和类别，这样推理的主体部分就算完成了。
结合DeepSort进行检测跟踪计数
首先从github上将unbox_yolov5_deepsort_counting项目拉到Jetson NX上，然后将{tensorrtx}/yolov5也复制到Jetson NX上，按第二种方法，重新生成engine和so库，因为在训练机器上生成的so库和engine文件不能在jetson上用，因此这里必须重新生成一次，生成过程跟上述过程完全一致。
我们将新生成的libmyplugins.so、yolov5m.engine、yolov5_trt.py放入unbox_yolov5_deepsort_counting目录下。
第一步我们先把依赖的包安装好，保证非TRT版的能够正常跑起来，执行:
python3 main.py
程序正常跑起来以后，我们分析一下main.py里的代码部分。程序第四行导入Detector：
from detector import Detector
第63行初始化detector：
detector = Detector()
第79行调用detector.detect()对图像执行真正的推理：
bboxes = detector.detect(im)
我们进入这个detector.detect()函数看一下，返回值具体是什么：
for *x, conf, cls_id in det:
lbl = self.names[int(cls_id)]
if lbl not in ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck']:
continue
pass
x1, y1 = int(x[0]), int(x[1])
x2, y2 = int(x[2]), int(x[3])
boxes.append(
(x1, y1, x2, y2, lbl, conf))
return boxes
可以看出这里box里含有4个坐标，lbl是类别名字，conf应该就算置信度，我们在yolov5_trt.py里推理后拿到的结果是框的4个坐标，类别的id和conf置信度，因此我们需要修改yolov5_trt.py，将类别id变成类别名字，同时将返回值封装成boxes。
在使用TensorRT进行YoloV5推理的时候，我们拿到了拿到result_boxes, result_scores, result_classid，我们需要再加入一个后处理，将其封装为boxes：
boxes = []
for i in range(len(result_boxes)):
box = result_boxes[i]
if box is not None and len(box):
for k in box:
x1 = int(box[0])
y1 = int(box[1])
x2 = int(box[2])
y2 = int(box[3])
label = categories[int(result_classid[i])]
conf = result_scores[i]
boxes.append((x1,y1,x2,y2,label,conf))
其中categories是name列表：
categories = ["person", "bicycle", "elecBicycle","car", "bus", "truck"]
其余部分，包括图像预处理、推理后的nms等都不用遍，我们将接口封装为与main.py文件里调用的接口一样，修改main.py文件detector部分的导入，初始化和推理的调用，即可完成对原来的检测器的替换，再执行：
python3 main.py
即为使用TensorRT版的检测器。这样我们就完成了一个使用TensorRT + DeepSort的检测跟踪计数程序。
...全文
请发表友善的回复…
发表回复
TracyGC 2024-03-28
- 打赏
- 举报
为什么infer里面的def infer(self, img)的imgg参数没有用到
m0_47941105 2023-06-25
- 打赏
- 举报
服务器厂商 欢迎翻牌
陈浥尘 2023-06-15
- 打赏
- 举报
“第二种方式其实是wang-xinyu等人使用TensorRT API重新一层层构建了YoloV5网络，重新设计了nms的后续处理，节省了大量nms时间。不足之处是使用起来比较麻烦。”在wang-xinyu的版本中，python代码里面有nms的过程，感觉区别不大，在agx orin上实测速度也差距不大
turndown 2023-04-23
- 打赏
- 举报
写的好详细，赞！！！
陈小春学安全 2023-01-03
- 打赏
- 举报
很不错的内容，干货满满，已支持师傅，期望师傅能输出更多干货，并强烈给师傅五星好评
另外，如果可以的话，期待师傅能给正在参加年度博客之星评选的我一个五星好评，您的五星好评都是对我的支持与鼓励（帖子中有大额红包惊喜哟，不要忘记评了五星后领红包哟）
⭐ ⭐ ⭐ ⭐ ⭐ 博主信息⭐ ⭐ ⭐ ⭐ ⭐
博主：橙留香Park
本人原力等级：5
链接直达：https://bbs.csdn.net/topics/611387568
微信直达：Blue_Team_Park
⭐ ⭐ ⭐ ⭐ ⭐ 五星必回！！！⭐ ⭐ ⭐ ⭐ ⭐
点赞五星好评回馈小福利：抽奖赠书 | 总价值200元，书由君自行挑选（从此页面参与抽奖的同学，只需五星好评后，参与抽奖）
YT 2022-07-01
- 打赏
- 举报
给力给力呢！
陈小春学安全 2023-01-03
- 举报
很不错的内容，干货满满，已支持师傅，期望师傅能输出更多干货，并强烈给师傅五星好评
另外，如果可以的话，期待师傅能给正在参加年度博客之星评选的我一个五星好评，您的五星好评都是对我的支持与鼓励（帖子中有大额红包惊喜哟，不要忘记评了五星后领红包哟）
⭐ ⭐ ⭐ ⭐ ⭐ 博主信息⭐ ⭐ ⭐ ⭐ ⭐
博主：橙留香Park
本人原力等级：5
链接直达：https://bbs.csdn.net/topics/611387568
微信直达：Blue_Team_Park
⭐ ⭐ ⭐ ⭐ ⭐ 五星必回！！！⭐ ⭐ ⭐ ⭐ ⭐
点赞五星好评回馈小福利：抽奖赠书 | 总价值200元，书由君自行挑选（从此页面参与抽奖的同学，只需五星好评后，参与抽奖）
xiaomaku 2022-06-30
- 打赏
- 举报
爱了爱了
陈小春学安全 2023-01-03
- 举报
很不错的内容，干货满满，已支持师傅，期望师傅能输出更多干货，并强烈给师傅五星好评
另外，如果可以的话，期待师傅能给正在参加年度博客之星评选的我一个五星好评，您的五星好评都是对我的支持与鼓励（帖子中有大额红包惊喜哟，不要忘记评了五星后领红包哟）
⭐ ⭐ ⭐ ⭐ ⭐ 博主信息⭐ ⭐ ⭐ ⭐ ⭐
博主：橙留香Park
本人原力等级：5
链接直达：https://bbs.csdn.net/topics/611387568
微信直达：Blue_Team_Park
⭐ ⭐ ⭐ ⭐ ⭐ 五星必回！！！⭐ ⭐ ⭐ ⭐ ⭐
点赞五星好评回馈小福利：抽奖赠书 | 总价值200元，书由君自行挑选（从此页面参与抽奖的同学，只需五星好评后，参与抽奖）
连798 2022-06-30
- 打赏
- 举报
👍👍
陈小春学安全 2023-01-03
- 举报
很不错的内容，干货满满，已支持师傅，期望师傅能输出更多干货，并强烈给师傅五星好评
另外，如果可以的话，期待师傅能给正在参加年度博客之星评选的我一个五星好评，您的五星好评都是对我的支持与鼓励（帖子中有大额红包惊喜哟，不要忘记评了五星后领红包哟）
⭐ ⭐ ⭐ ⭐ ⭐ 博主信息⭐ ⭐ ⭐ ⭐ ⭐
博主：橙留香Park
本人原力等级：5
链接直达：https://bbs.csdn.net/topics/611387568
微信直达：Blue_Team_Park
⭐ ⭐ ⭐ ⭐ ⭐ 五星必回！！！⭐ ⭐ ⭐ ⭐ ⭐
点赞五星好评回馈小福利：抽奖赠书 | 总价值200元，书由君自行挑选（从此页面参与抽奖的同学，只需五星好评后，参与抽奖）
扫地的小何尚 2022-06-30
- 打赏
- 举报
手动点赞
