TensorRT部署YOLOv5 06 视频推理

发表于 更新于 分类于 ，

在上一篇文章中，对使用TensorRT引擎进行单幅图像推理进行了介绍，本文仍然使用Python API，在上篇文章的基础上，对视频实时推理的实现进行介绍

workflow

视频实时推理的过程大致如下

  1. 加载TensorRT引擎文件并进行反序列化，创建执行上下文context

  2. 根据引擎申请输入输出Buffers，并进行绑定

  3. 读入一帧图像数据

  4. 图像预处理

  5. 调用推断API进行推断

  6. 输出解码

  7. 画框并显示

  8. 返回步骤3

这里从视频文件读入图像数据以及画框显示使用的都是opencv

实现

大部分代码处理与上一篇文章相同，视频推理的整个程序主体如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
with get_engine(args.engine, TRT_LOGGER) as engine: \
    engine.create_execution_context() as context:
    # alloc buffers and binding
    inputs, outputs, bindings, stream = utils.allocate_buffers(engine)
    # open video file
    cap = cv2.VideoCapture(args.video)
    if cap.isOpened():
        # create window
        cv2.namedWindow('Video display', cv2.WINDOW_AUTOSIZE)
        # read frame
        ret, frame = cap.read()
        if not ret:
            break
        # image preprocess
        im = utils.image_preprocess(frame, args, dtype)
        # image push to model input
        inputs[0].host = im
        # do inference
        inference_outs = do_inference_v2(context, bindings, inputs, 
            outputs, stream)                
        # infernece out convert to preds
        preds = []
        for i, out in enumerate(inference_outs):
            pred = out.reshape(pred_shape[i])
            preds.append(pred)
        preds.reverse()

        # decode boxes
        out_boxes, out_scores, out_classes = decode_box(preds, anchors, 
            anchor_masks, args.num_classes, 
            [args.input_size, args.input_size],
            [frame.shape[1], frame.shape[0]])

        # draw
        if len(out_boxes) > 0:
        frame = draw_output(frame, (out_boxes, out_scores, out_classes, 
            len(out_boxes)))

        # show frame
        cv2.imshow('Video display', frame)

        # wait exit
        key = cv2.waitKey(1) & 0xff
        if key == 27:
            break
    # release capture
    cap.release()
    cv2.destroyAllWindows()

运行该程序，加载自己训练的yolov5m的20分类模型，结果帧率只有1fps+，在TensorRT引擎转换阶段，是可以大致清楚一个模型的推理延迟的，而目前的帧率只有1fps显然过于低了。该程序中使用opencv进行h.264视频的加载和解码，使用的是软解码，实时性能会比较差，Nvidia开发平台内置了硬件h264编解码器，而该编解码加速器适配到了Gstreamer，因此在上层应用中，使用opencv时可以构建Gstreamer的pipeline，将视频传送给硬件编解码器进行解码，然后再设置输出端为appsink，即当前程序接收解码后的数据

以下是构建h264的Gstreamer pipeline

1
2
3
4
5
6
7
8
9
10
11
def video_pipeline(filesrc, width, height):
    return (
        'filesrc location={} !' 
        'qtdemux ! '
        'h264parse ! '
        'nvv4l2decoder ! '
        'nvvidconv ! '
        'video/x-raw,width={},height={},format=BGRx ! '
        'videoconvert ! '
        'appsink'.format(filesrc, width, height)
    )

将opencv创建Capture的代码改写一下

1
2
cap = cv2.VideoCapture(video_pipeline(args.video, args.width, 
    args.height), cv2.CAP_GSTREAMER)

改写后运行代码，帧率提升到了3fps+，后文将继续介绍如何进行帧率提升的优化方法

运行效果