Jetson Nano部署YOLOv5与Tensorrtx加速——（自己走一遍全过程记录）

说在前面

搞了一下Jetson nano和YOLOv5，网上的资料大多重复也有许多的坑，在配置过程中摸爬滚打了好几天，出坑后决定写下这份教程供自己备忘。

事先声明，这篇文章的许多内容本身并不是原创，而是将配置过程中的文献进行了搜集整理，但是所有步骤都1:1复刻我的配置过程，包括其中的出错和解决途径，但是每个人的设备和网络上的包都是不断更新的，不能保证写下这篇文章之后的版本在兼容性上没有问题，总之提前祝自己好运！

一、烧录镜像

1、镜像选择

这里我选择的是亚博智能，它已经将镜像大部分给配置好了。

获取链接：（提取码：o6a4）
镜像的下载地址

里面已经安装好了如下的东西：
CUDA10.2，CUDNNv8，tensorRT，opencv4.1.1，python2，python3,tensorflow2.3，jetpack4.4.1，yolov4-tiny和yolov4，jetson-inference包（含资料中的训练模型），jetson-gpio库，安装pytorch1.6和torchvesion0.7，安装node v15.0.1，npm7.0.3，jupterlab，jetcham，已开启VNC服务。

2、镜像烧录方法

烧录方法参考这一篇文章，很简单的。
镜像烧录方法

3、Jetson nano 系统初始化设置

插卡！开机！最好连接上屏幕，不差这几个钱了。之后的很多命令需要用到root权限，我们需要开启root用户。

sudo passwd root

之后设置密码即可

开发板需要插上网线或者插上免驱动的无线网卡联网！！！

①做个小备份

sudo cp /etc/apt/sources.list /etc/apt/sources.list.bak
sudo gedit /etc/apt/sources.list

②删除所有，替换为如下的东西

deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic main multiverse restricted universe
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-security main multiverse restricted universe
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-updates main multiverse restricted universe
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-backports main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-security main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-updates main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-backports main multiverse restricted universe

题外话：如何更换源呢？
Jetson Nano 烧录的镜像是国外的源，安装软件和升级软件包的速度非常慢，甚至还会常常出现网络错误，更换源的步骤如下：
①先备份原本的source.list文件。

sudo cp /etc/apt/sources.list /etc/apt/sources.list.bak

②编辑source.list，并更换国内源。

sudo gedit /etc/apt/sources.list

③按 “i” 开始输入，删除所有内容，复制并更换源。（这里选清华源或中科大源其中一个，然后保存）

# 清华源
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic main multiverse restricted universe
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-security main multiverse restricted universe
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-updates main multiverse restricted universe
deb http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-backports main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-security main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-updates main multiverse restricted universe
deb-src http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/ bionic-backports main multiverse restricted universe# 中科大源
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-updates main restricted
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic universe
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-updates universe
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic multiverse
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-updates multiverse
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-backports main restricted universe multiverse
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-security main restricted
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-security universe
deb http://mirrors.ustc.edu.cn/ubuntu-ports/ bionic-security multiverse

④更新软件

# 更新软件
sudo apt-get update
sudo apt-get upgrade

二、开始配置所需的环境，安装各种支持包

1、配置CUDA

Jetson nano内置好了CUDA,但需要配置环境变量才能使用，打开命令行添加环境变量即可，我这里是CUDA10.2如果不是使用我的镜像就需要根据自己的CUDA版本去填写路径了。

#打开终端，输入命令
vi .bashrc

拉到最后，在最后添加这些

export PATH=/usr/local/cuda-10.2/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export CUDA_ROOT=/usr/local/cuda

应用当前配置（刷新一下）

source ~/.bashrc

查看是否配置成功

nvcc -V

**出来这个就算成功了**

2、安装pip3

sudo apt-get update
sudo apt-get install python3-pip python3-dev -y

3、安装jtop

安装jtop库这个可以监控自己的设备CPU、GPU工作状态

sudo -H pip3 install jetson-stats
sudo jtop		#运行jtop（第一次可能不行，第二次就好了）  按【q】退出

4、配置可能需要用到的库

sudo apt-get install build-essential make cmake cmake-curses-gui -y
sudo apt-get install git g++ pkg-config curl -y
sudo apt-get install libatlas-base-dev gfortran libcanberra-gtk-module libcanberra-gtk3-module -y
sudo apt-get install libhdf5-serial-dev hdf5-tools -y
sudo apt-get install nano locate screen -y

5、安装所需要的依赖环境

sudo apt-get install libfreetype6-dev -y
sudo apt-get install protobuf-compiler libprotobuf-dev openssl -y
sudo apt-get install libssl-dev libcurl4-openssl-dev -y
sudo apt-get install cython3 -y

6、安装opencv的系统级依赖,一些编解码的库

sudo apt-get install build-essential -y
sudo apt-get install cmake git libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev -y
sudo apt-get install python-dev python-numpy libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff5-dev libdc1394-22-dev -y
sudo apt-get install libavcodec-dev libavformat-dev libswscale-dev libv4l-dev liblapacke-dev -y
sudo apt-get install libxvidcore-dev libx264-dev -y
sudo apt-get install libatlas-base-dev gfortran -y
sudo apt-get install ffmpeg -y

7、更新CMake

这一步是必须的，因为ARM架构的很多东西都要从源码编译

wget http://www.cmake.org/files/v3.13/cmake-3.13.0.tar.gz
tar xpvf cmake-3.13.0.tar.gz cmake-3.13.0/  #解压
cd cmake-3.13.0/
./bootstrap --system-curl	# 漫长的等待,做一套眼保健操...
make -j4 #编译  同样是漫长的等待...
echo 'export PATH=~/cmake-3.13.0/bin/:$PATH' >> ~/.bashrc
source ~/.bashrc #更新.bashrc

8、U盘兼容

之后的步骤可能需要使用U盘把大文件拷入开发板，但是对于大容量设备可能会出现无法挂载，一条安装命令解决。

sudo apt-get install exfat-utils

三、安装pytorch

Jetson nano上的Linux其实不是x86架构,而是类似手机的ARM架构，这也就导致它的很多包和普通的Linux上的不是通用的。也是踩过的坑之一，pytorch官网下载的包，在实际使用时无法调用开发板的显卡（这是个大问题，失去显卡的开发板算力暴跌！）。这里的PyTorch以及接下来的torchvision等包都需要安装Nvidia官网给出的版本。

1.下载PyTorch1.8
我已经下载好了，现成的安装包下载链接奉上：（提取码：yvex）
安装包

2.安装PyTorch1.8
把下载的东西用U盘拷到Jetson nano开发板上，建议放桌面上，好找。
sudo pip3 install …# 直接把.whl拖到命令窗口中，让它自动填充文件位置
安装需要略漫长的等待。
在这里插入图片描述

四、安装torchvision 0.9.0版本

PyTorch和torchvision版本是需要对应的，上一步下载的那个正好是对应的。

1.提前安装好我们需要的依赖

sudo apt-get install libopenmpi2
sudo apt-get install libopenblas-dev
sudo apt-get install libjpeg-dev zlib1g-dev

2.安装torchvision 0.9.0
同样需要特殊的匹配Jetson nano的版本，步骤三中个人链接里包含了这个torchvision。把下载的包拷到开发板上，同样建议放桌面上。

cd torchvision	# 进入到这个包的目录下
export BUILD_VERSION=0.9.0
sudo python3 setup.py install		# 安装（估计要20、30分钟不止吧）

3.检验一下是否成功安装

python3
import torch
import torchvision
print(torch.cuda.is_available())	# 这一步如果输出True那么就成功了！
quit()	# 最后退出python编译

在这里插入图片描述

五、下载YOLOv5-5.0源代码

在自己的电脑上或服务器上训练好。这里如何训练，不做过多解释，可以去B站找一些视频学习一下。我的项目是检测电梯按键。需要数据集和训练权重以及各种yolov5改进代码的同学可以滴滴私信联系我。

六、安装使YOLOv5成功运行需依赖的包

注意：下载过程如果因为网络原因失败的话可以在命令后加上 -i https://pypi.tuna.tsinghua.edu.cn/simple 来使用清华镜像源

1、

sudo pip3 install matplotlib==3.2.2
sudo pip3 install --upgrade Cython	#更新一下这个包

2、numpy有些特殊，已经自带了，但是是apt-get安装的，所以先卸掉原来的，也方便之后包的管理

sudo apt-get remove python-numpy
sudo pip3 install numpy==1.19.4
sudo pip3 install scipy==1.4.1.	# 这个包安装巨慢，耐心等待

3、这之后的一些包我在安装时都没有指定版本,这里的指令是根据之后pip3 list补上的

sudo pip3 install tqdm==4.61.2
sudo pip3 install seaborn==0.11.1
sudo pip3 install scikit-build==0.11.1	# 安装opencv需要这个包
sudo pip3 install opencv-python==4.5.3.56	# 不出意外也是一个相当漫长的过程
sudo pip3 install tensorboard==2.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
sudo pip3 install --upgrade PyYAML	# 我升级到了5.4.1 也可以sudo pip3 install PyYAML==5.4.1
sudo pip3 install thop
sudo pip3 install pycocotools

4、根据YOLOv5官方给的所需的安装包清单，仔细对照，查漏补缺的给安装好。

安装命令输入格式：sudo pip3 install  .................
# base ----------------------------------------
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
Pillow
PyYAML>=5.3.1
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.41.0# logging -------------------------------------
tensorboard>=2.4.1
wandb# plotting ------------------------------------
seaborn>=0.11.0
pandas# export --------------------------------------
coremltools>=4.1
onnx>=1.8.1
scikit-learn==0.19.2  # for coreml quantization# extras --------------------------------------
thop  # FLOPS computation
pycocotools>=2.0  # COCO mAP

5、运行检测脚本
在源码的detect.py同目录下，打开终端，运行下面的命令。
效果还可以，启动模型要很久，预测效果还可以。之后就可以在自己的inference中的output中看到自己预测的图片了。
接着打开detecy.py检测脚本，修改一下检测资源参数，改为调用摄像头进行实时视频预测，大概10fps,应该说不算差，但是是有提升办法的。

python3 detect.py --source /path/to/xxx.jpg --weights /path/to/best.pt --conf-thres 0.7
或者是：
python3 detect.py

七、来一波TensorRT加速？

1、安装pycuda-2019

① （网络好的时候用这个方法）

在线安装pycuda

pip3 install pycuda

②（你的网络不好的时候用下面这个方法）

提取码：t94b 下载链接
下载完之后解压。
进入解压出来的文件。

tar zxvf pycuda-2019.1.2.tar.gz    
cd pycuda-2019.1.2/  
python3 configure.py --cuda-root=/usr/local/cuda-10.2
sudo python3 setup.py install

出现这个就说明正在编译文件安装，等待一段时间后即可安装完成。
在这里插入图片描述
安装完出现：

就表明安装成功了。

但是使用的时候还得配置一下一些必要的东西不然会报错：

FileNotFoundError: [Errno 2] No such file or directory: ‘nvcc’

将nvcc的完整路径硬编码到Pycuda的compiler.py文件中的compile_plain()
中，大约在第 73 行的位置中加入下面段代码！

nvcc = '/usr/local/cuda/bin/'+nvcc

在这里插入图片描述

2、TensorRT加速

这时我们要用到一个大佬的开源，GitHub地址如下：

https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

大佬是真的牛批，好好看一下吧。不仅有yolov5的，还有好多算法的，大佬都给做了相关的加速，大佬给他的项目起名叫TensorRTx，比原版的TensorRT加速更好用。
需要下载两个东西：
第一是：YOLOv5原版的开源程序（选择v5.0版本）
第二是：将大佬开源的项目tensorrtx，下载到自己的windows电脑上
然后，把tensorrtx文件夹整体，复制粘贴到yolov5-5.0原版程序的文件夹中。
我为了自己理解方便，和之后的操作，稍微改了一下文件夹名称：
（当然我都把东西准备好了，下载就行：提取码：私信聊）
下载
YOLOv5原版程序文件夹改名为yolov5（Tensorrtx）如下图所示：
在这里插入图片描述
把tenserrtx文件夹整体改名为：tensorrtx-yolov5-v5.0，复制粘贴到yolov5（Tensorrtx）的文件夹中，如下图所示：

下面开始真正的操作了：
①生成.wts文件（在windows电脑上操作即可）
1.将训练得到的.pt权重文件改名为yolov5s.pt（必须改成这个名，没有为什么），把它放到yolov5（Tenserrtx）文件夹中。
2.将这个文件 yolov5-5.0（Tensorrtx）\tensorrtx-yolov5-v5.0\yolov5\gen_wts.py
复制粘贴到yolov5（Tensorrtx）文件夹中。

注意： 此时yolov5（Tensorrtx）文件夹中有了 yolov5s.pt和gen_wts.py这两个文件。

然后，在yolov5（Tensorrtx）文件夹中右击鼠标，打开终端，激活在anaconda中自己创建的虚拟环境
比如：conda activate torch1.10 。
然后输入命令：

python gen_wts.py -w yolov5s.pt -o yolov5s.wts

(问题：在anaconda中自己创建虚拟环境不会？那你就去B站找视频自己学一下。YOLOv5的权重都训练好了，这个不可能不会的。)
文件内会生成一个文件：yolov5s.wts

② build（在Jetson nano上弄）（这一步是生成引擎文件）
1.将上述生成的.wts文件用U盘复制到Jetson nano里的yolov5-5.0（Tensorrtx）\tensorrtx-yolov5-v5.0\yolov5文件夹中。
2.打开上述文件夹里的yololayer.h文件，修改CLASS_NUM的数量（根据自己训练模型的类的个数来设，我的是55）。
3.此时上述文件夹里有（.wts 是在windows电脑上生成的）（yolov5.cpp 未进行过改动）（yololayer.h 已经改为自己训练的类数了）这三个。
4.在上述文件夹中打开终端，依次运行指令：

mkdir build
cd build
cmake ..
make
sudo ./yolov5 -s ../yolov5s.wts yolov5s.engine s

稍微等待之后，在build文件夹中便通过tensorrtx生成了基于C++的engine引擎部署文件了。但是我C++水平不怎么样，对它有种心理上的抵触，把他搞成python的吧。

③USB摄像头实时检测加速
由于本人C++语言很一般，所以只能硬着头皮修改了下yolov5-5.0（Tensorrtx）\tensorrtx-yolov5-v5.0\yolov5文件夹中的yolov5_trt.py脚本，脚本的代码格式较差，但是能够实现加速，有需要的可以作为一个参考。在文件夹下新建一个yolo_trt_test.py文件。复制下面 v4.0或者v5.0的代码到yolo_trt_test.py。
需要自行更改的地方：yolov5s.engine的路径要改成自己的、检测物体的类别名称要改为自己的。

①v5.0代码

"""
An example that uses TensorRT's Python api to make inferences.
"""
import ctypes
import os
import shutil
import random
import sys
import threading
import time
import cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import torchvision
import argparseCONF_THRESH = 0.5
IOU_THRESHOLD = 0.4def get_img_path_batches(batch_size, img_dir):ret = []batch = []for root, dirs, files in os.walk(img_dir):for name in files:if len(batch) == batch_size:ret.append(batch)batch = []batch.append(os.path.join(root, name))if len(batch) > 0:ret.append(batch)return retdef 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 objectcolor:  color to draw rectangle, such as (0,255,0)label:  strline_thickness: intreturn:no return"""tl = (line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1)  # line/font thicknesscolor = 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 thicknesst_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filledcv2.putText(img,label,(c1[0], c1[1] - 2),0,tl / 3,[225, 255, 255],thickness=tf,lineType=cv2.LINE_AA,)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 filewith 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_sizedtype = trt.nptype(engine.get_binding_dtype(binding))# Allocate host and device buffershost_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)# Storeself.stream = streamself.context = contextself.engine = engineself.host_inputs = host_inputsself.cuda_inputs = cuda_inputsself.host_outputs = host_outputsself.cuda_outputs = cuda_outputsself.bindings = bindingsself.batch_size = engine.max_batch_sizedef infer(self, input_image_path):threading.Thread.__init__(self)# Make self the active context, pushing it on top of the context stack.self.ctx.push()self.input_image_path = input_image_path# Restorestream = self.streamcontext = self.contextengine = self.enginehost_inputs = self.host_inputscuda_inputs = self.cuda_inputshost_outputs = self.host_outputscuda_outputs = self.cuda_outputsbindings = self.bindings# Do image preprocessbatch_image_raw = []batch_origin_h = []batch_origin_w = []batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w])input_image, image_raw, origin_h, origin_w = self.preprocess_image(input_image_path)batch_origin_h.append(origin_h)batch_origin_w.append(origin_w)np.copyto(batch_input_image, input_image)batch_input_image = np.ascontiguousarray(batch_input_image)# Copy input image to host buffernp.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 streamstream.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 = 1output = host_outputs[0]# Do postprocessresult_boxes, result_scores, result_classid = self.post_process(output, origin_h, origin_w)# Draw rectangles and labels on the original imagefor j in range(len(result_boxes)):box = result_boxes[j]plot_one_box(box,image_raw,label="{}:{:.2f}".format(categories[int(result_classid[j])], result_scores[j]),)return image_raw, end - startdef destroy(self):# Remove any context from the top of the context stack, deactivating it.self.ctx.pop()def get_raw_image(self, image_path_batch):"""description: Read an image from image path"""for img_path in image_path_batch:yield cv2.imread(img_path)def get_raw_image_zeros(self, image_path_batch=None):"""description: Ready data for warmup"""for _ in range(self.batch_size):yield np.zeros([self.input_h, self.input_w, 3], dtype=np.uint8)def preprocess_image(self, input_image_path):"""description: Convert BGR image to RGB,resize and pad it to target size, normalize to [0,1],transform to NCHW format.param:input_image_path: str, image pathreturn:image:  the processed imageimage_raw: the original imageh: original heightw: original width"""image_raw = input_image_pathh, w, c = image_raw.shapeimage = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)# Calculate widht and height and paddingsr_w = self.input_w / wr_h = self.input_h / hif r_h > r_w:tw = self.input_wth = int(r_w * h)tx1 = tx2 = 0ty1 = int((self.input_h - th) / 2)ty2 = self.input_h - th - ty1else:tw = int(r_h * w)th = self.input_htx1 = int((self.input_w - tw) / 2)tx2 = self.input_w - tw - tx1ty1 = ty2 = 0# Resize the image with long side while maintaining ratioimage = cv2.resize(image, (tw, th))# Pad the short side with (128,128,128)image = cv2.copyMakeBorder(image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128))image = image.astype(np.float32)# Normalize to [0,1]image /= 255.0# HWC to CHW format:image = np.transpose(image, [2, 0, 1])# CHW to NCHW formatimage = np.expand_dims(image, axis=0)# Convert the image to row-major order, also known as "C order":image = np.ascontiguousarray(image)return image, image_raw, h, wdef xywh2xyxy(self, origin_h, origin_w, x):"""description:    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-rightparam:origin_h:   height of original imageorigin_w:   width of original imagex:          A boxes tensor, each row is a box [center_x, center_y, w, h]return:y:          A boxes tensor, each row is a box [x1, y1, x2, y2]"""y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)r_w = self.input_w / origin_wr_h = self.input_h / origin_hif r_h > r_w:y[:, 0] = x[:, 0] - x[:, 2] / 2y[:, 2] = x[:, 0] + x[:, 2] / 2y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2y /= r_welse:y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2y[:, 1] = x[:, 1] - x[:, 3] / 2y[:, 3] = x[:, 1] + x[:, 3] / 2y /= r_hreturn ydef post_process(self, output, origin_h, origin_w):"""description: postprocess the predictionparam:output:     A tensor likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...] origin_h:   height of original imageorigin_w:   width of original imagereturn:result_boxes: finally boxes, a boxes tensor, each row is a box [x1, y1, x2, y2]result_scores: finally scores, a tensor, each element is the score correspoing to boxresult_classid: finally classid, a tensor, each element is the classid correspoing to box"""# Get the num of boxes detectednum = int(output[0])# Reshape to a two dimentional ndarraypred = np.reshape(output[1:], (-1, 6))[:num, :]# to a torch Tensorpred = torch.Tensor(pred).cuda()# Get the boxesboxes = pred[:, :4]# Get the scoresscores = pred[:, 4]# Get the classidclassid = pred[:, 5]# Choose those boxes that score > CONF_THRESHsi = scores > CONF_THRESHboxes = boxes[si, :]scores = scores[si]classid = classid[si]# Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]boxes = self.xywh2xyxy(origin_h, origin_w, boxes)# Do nmsindices = torchvision.ops.nms(boxes, scores, iou_threshold=IOU_THRESHOLD).cpu()result_boxes = boxes[indices, :].cpu()result_scores = scores[indices].cpu()result_classid = classid[indices].cpu()return result_boxes, result_scores, result_classidclass inferThread(threading.Thread):def __init__(self, yolov5_wrapper):threading.Thread.__init__(self)self.yolov5_wrapper = yolov5_wrapperdef infer(self , frame):batch_image_raw, use_time = self.yolov5_wrapper.infer(frame)# for i, img_path in enumerate(self.image_path_batch):#     parent, filename = os.path.split(img_path)#     save_name = os.path.join('output', filename)#     # Save image#     cv2.imwrite(save_name, batch_image_raw[i])# print('input->{}, time->{:.2f}ms, saving into output/'.format(self.image_path_batch, use_time * 1000))return batch_image_raw,use_timeclass warmUpThread(threading.Thread):def __init__(self, yolov5_wrapper):threading.Thread.__init__(self)self.yolov5_wrapper = yolov5_wrapperdef run(self):batch_image_raw, use_time = self.yolov5_wrapper.infer(self.yolov5_wrapper.get_raw_image_zeros())print('warm_up->{}, time->{:.2f}ms'.format(batch_image_raw[0].shape, use_time * 1000))if __name__ == "__main__":# load custom pluginsparser = argparse.ArgumentParser()parser.add_argument('--engine', nargs='+', type=str, default="build/yolov5s.engine", help='.engine path(s)') #改为自己的路径parser.add_argument('--save', type=int, default=0, help='save?')opt = parser.parse_args()PLUGIN_LIBRARY = "build/libmyplugins.so"engine_file_path = opt.enginectypes.CDLL(PLUGIN_LIBRARY)# load coco labelscategories = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light","fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow","elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee","skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard","tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple","sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch","potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone","microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear","hair drier", "toothbrush"] #改为自己的检测类别名称# a YoLov5TRT instanceyolov5_wrapper = YoLov5TRT(engine_file_path)cap = cv2.VideoCapture(0)try:thread1 = inferThread(yolov5_wrapper)thread1.start()thread1.join()while 1:_,frame = cap.read()img,t=thread1.infer(frame)cv2.imshow("result", img)if cv2.waitKey(1) & 0XFF == ord('q'):  # 1 millisecondbreakfinally:# destroy the instancecap.release()cv2.destroyAllWindows()yolov5_wrapper.destroy()

②v4.0代码

"""
An example that uses TensorRT's Python api to make inferences.
"""
import ctypes
import os
import random
import sys
import threading
import timeimport cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import torchvisionINPUT_W = 608
INPUT_H = 608
CONF_THRESH = 0.15
IOU_THRESHOLD = 0.45
int_box=[0,0,0,0]
int_box1=[0,0,0,0]
fps1=0.0
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 objectcolor:  color to draw rectangle, such as (0,255,0)label:  strline_thickness: intreturn:no return"""tl = (line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1)  # line/font thicknesscolor = color or [random.randint(0, 255) for _ in range(3)]c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))C2 = c2cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)if label:tf = max(tl - 1, 1)  # font thicknesst_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]c2 = c1[0] + t_size[0], c1[1] + t_size[1] + 8cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filledcv2.putText(img,label,(c1[0], c1[1]+t_size[1] + 5),0,tl / 3,[255,255,255],thickness=tf,lineType=cv2.LINE_AA,)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.cfx = cuda.Device(0).make_context()stream = cuda.Stream()TRT_LOGGER = trt.Logger(trt.Logger.INFO)runtime = trt.Runtime(TRT_LOGGER)# Deserialize the engine from filewith 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:size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_sizedtype = trt.nptype(engine.get_binding_dtype(binding))# Allocate host and device buffershost_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):host_inputs.append(host_mem)cuda_inputs.append(cuda_mem)else:host_outputs.append(host_mem)cuda_outputs.append(cuda_mem)# Storeself.stream = streamself.context = contextself.engine = engineself.host_inputs = host_inputsself.cuda_inputs = cuda_inputsself.host_outputs = host_outputsself.cuda_outputs = cuda_outputsself.bindings = bindingsdef infer(self, input_image_path):global int_box,int_box1,fps1# threading.Thread.__init__(self)# Make self the active context, pushing it on top of the context stack.self.cfx.push()# Restorestream = self.streamcontext = self.contextengine = self.enginehost_inputs = self.host_inputscuda_inputs = self.cuda_inputshost_outputs = self.host_outputscuda_outputs = self.cuda_outputsbindings = self.bindings# Do image preprocessinput_image, image_raw, origin_h, origin_w = self.preprocess_image(input_image_path)# Copy input image to host buffernp.copyto(host_inputs[0], input_image.ravel())# Transfer input data  to the GPU.cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)# Run inference.context.execute_async(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 streamstream.synchronize()# Remove any context from the top of the context stack, deactivating it.self.cfx.pop()# Here we use the first row of output in that batch_size = 1output = host_outputs[0]# Do postprocessresult_boxes, result_scores, result_classid = self.post_process(output, origin_h, origin_w)# Draw rectangles and labels on the original imagefor i in range(len(result_boxes)):box1 = result_boxes[i]plot_one_box(box1,image_raw,label="{}:{:.2f}".format(categories[int(result_classid[i])], result_scores[i]),)return image_raw# parent, filename = os.path.split(input_image_path)# save_name = os.path.join(parent, "output_" + filename)# # 　Save image# cv2.imwrite(save_name, image_raw)def destroy(self):# Remove any context from the top of the context stack, deactivating it.self.cfx.pop()def preprocess_image(self, input_image_path):"""description: Read an image from image path, convert it to RGB,resize and pad it to target size, normalize to [0,1],transform to NCHW format.param:input_image_path: str, image pathreturn:image:  the processed imageimage_raw: the original imageh: original heightw: original width"""image_raw = input_image_pathh, w, c = image_raw.shapeimage = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)# Calculate widht and height and paddingsr_w = INPUT_W / wr_h = INPUT_H / hif r_h > r_w:tw = INPUT_Wth = int(r_w * h)tx1 = tx2 = 0ty1 = int((INPUT_H - th) / 2)ty2 = INPUT_H - th - ty1else:tw = int(r_h * w)th = INPUT_Htx1 = int((INPUT_W - tw) / 2)tx2 = INPUT_W - tw - tx1ty1 = ty2 = 0# Resize the image with long side while maintaining ratioimage = cv2.resize(image, (tw, th))# Pad the short side with (128,128,128)image = cv2.copyMakeBorder(image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128))image = image.astype(np.float32)# Normalize to [0,1]image /= 255.0# HWC to CHW format:image = np.transpose(image, [2, 0, 1])# CHW to NCHW formatimage = np.expand_dims(image, axis=0)# Convert the image to row-major order, also known as "C order":image = np.ascontiguousarray(image)return image, image_raw, h, wdef xywh2xyxy(self, origin_h, origin_w, x):"""description:    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-rightparam:origin_h:   height of original imageorigin_w:   width of original imagex:          A boxes tensor, each row is a box [center_x, center_y, w, h]return:y:          A boxes tensor, each row is a box [x1, y1, x2, y2]"""y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)r_w = INPUT_W / origin_wr_h = INPUT_H / origin_hif r_h > r_w:y[:, 0] = x[:, 0] - x[:, 2] / 2y[:, 2] = x[:, 0] + x[:, 2] / 2y[:, 1] = x[:, 1] - x[:, 3] / 2 - (INPUT_H - r_w * origin_h) / 2y[:, 3] = x[:, 1] + x[:, 3] / 2 - (INPUT_H - r_w * origin_h) / 2y /= r_welse:y[:, 0] = x[:, 0] - x[:, 2] / 2 - (INPUT_W - r_h * origin_w) / 2y[:, 2] = x[:, 0] + x[:, 2] / 2 - (INPUT_W - r_h * origin_w) / 2y[:, 1] = x[:, 1] - x[:, 3] / 2y[:, 3] = x[:, 1] + x[:, 3] / 2y /= r_hreturn ydef post_process(self, output, origin_h, origin_w):"""description: postprocess the predictionparam:output:     A tensor likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...]origin_h:   height of original imageorigin_w:   width of original imagereturn:result_boxes: finally boxes, a boxes tensor, each row is a box [x1, y1, x2, y2]result_scores: finally scores, a tensor, each element is the score correspoing to boxresult_classid: finally classid, a tensor, each element is the classid correspoing to box"""# Get the num of boxes detectednum = int(output[0])# Reshape to a two dimentional ndarraypred = np.reshape(output[1:], (-1, 6))[:num, :]# to a torch Tensorpred = torch.Tensor(pred).cuda()# Get the boxesboxes = pred[:, :4]# Get the scoresscores = pred[:, 4]# Get the classidclassid = pred[:, 5]# Choose those boxes that score > CONF_THRESHsi = scores > CONF_THRESHboxes = boxes[si, :]scores = scores[si]classid = classid[si]# Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]boxes = self.xywh2xyxy(origin_h, origin_w, boxes)# Do nmsindices = torchvision.ops.nms(boxes, scores, iou_threshold=IOU_THRESHOLD).cpu()result_boxes = boxes[indices, :].cpu()result_scores = scores[indices].cpu()result_classid = classid[indices].cpu()return result_boxes, result_scores, result_classidclass myThread(threading.Thread):def __init__(self, func, args):threading.Thread.__init__(self)self.func = funcself.args = argsdef run(self):self.func(*self.args)if __name__ == "__main__":# load custom pluginsPLUGIN_LIBRARY = "build/libmyplugins.so"ctypes.CDLL(PLUGIN_LIBRARY)engine_file_path = "yolov5s.engine"# load coco labelscategories = ["person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light","fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow","elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee","skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard","tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple","sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch","potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone","microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear","hair drier", "toothbrush"]# a  YoLov5TRT instanceyolov5_wrapper = YoLov5TRT(engine_file_path)cap = cv2.VideoCapture(0)while 1:_,image =cap.read()img=yolov5_wrapper.infer(image)cv2.imshow("result", img)if cv2.waitKey(1) & 0XFF == ord('q'):  # 1 millisecondbreakcap.release()cv2.destroyAllWindows()yolov5_wrapper.destroy()