14、TX2上运行YOLO5

一、下载yolov5、tensorrtx，yolov5s.pt

预备基础环境：
python3
torch1.7或者更高

1、生成权重文件（可以在pc上或者是在TX2上生成都可以）

1、在TX2的终端直接输入下面的命令

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git

2、下载yolov5s.pt

https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt

3、文件搬移
①

cp tensorrtx/yolov5/gen_wts.py yolov5

②
将yolov5s.pt（预训练模型）放到yolov5的weights文件夹下，当然也可以是其他位置，具体调整可以参考下面的python脚本

4、生成权重文件

python gen_wts.py

执行这个脚本的时候可能会出现如下问题，大概如下（可能不全，但是方法是一样的）
①其实是需要一些依赖的环境和包，
这是环境中缺少了tqdm进度条的安装包，需要使用conda或者pip命令进行安装

pip install tqdm

缺少其他的包就按照同样的方法即可。

②python库中urllib3 (1.22) or chardet (2.2.1) 的版本不兼容
解决方法如下：

pip uninstall urllib3
pip uninstall chardet
pip install requests

2、生成部署引擎（必须在部署的硬件平台上编译，此处在TX2上）

将yolov5s.wts文件放到tensorrtx/yolov5文件夹中

mv yolov5s.wts  ~/tensorrtx/yolov5

在yolov5.cpp文件中还可以修改fp16还是fp32(nano不支持int8)、device(选择哪一个GPU设备)、nms_thresh(nms的阈值)、conf_thresh(conf的置信度)、batch_size(批次大小),此处可以使用默认的参数。

打开yololayer.h文件，修改他的num总数，根据你训练模型的类个数来(如果是你自己针对特定的数据集训练的模型的话)。

除此之外，我们还可以修改输入图片的尺寸，但必须是32的倍数。缩小输入尺寸可以一定程度上加快推理速度。

然后反序列化引擎对图像做处理

sudo ./yolov5 -d yolov5s.engine ../samples

二、视频检测

代码如下：

import time
import cv2
import pycuda.autoinit  # This is needed for initializing CUDA driver
import numpy as np
import ctypes
import tensorrt as trt
import pycuda.driver as cuda

import threading
import random

INPUT_W = 608
INPUT_H = 608
CONF_THRESH = 0.2
IOU_THRESHOLD = 0.4

categories = ['vehicle','bicyle','pedestrain','road_sign']


def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param:
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return

    """
    tl = (
            line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )


def draw_boxes(image_raw, result_boxes, result_scores, result_classid):
    for i in range(len(result_boxes)):
        box = result_boxes[i]
        plot_one_box(
            box,
            image_raw