devbox
程序自动化专家
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

部署Yolov5模型到jetson nano上_build cuda_10.2_r440.tc440_70.29663091_0

作者：程序自动化专家 | 2024-01-28 23:21:03

赞

踩

build cuda_10.2_r440.tc440_70.29663091_0

目录

一、检查是否安装cuda

二、安装好pip3，系统已经自带python3.6.9

三、检测是否安装gpu版本的tensorflow

四、安装pycuda

五、下载tensorrtx源码

六、模型测试

一、检查是否安装cuda

nvcc -V

  1. ljx@ljx-desktop:~/pycuda2/tensorrtx-yolov5-v5.0/yolov5$ nvcc -V
  2. nvcc: NVIDIA (R) Cuda compiler driver
  3. Copyright (c) 2005-2021 NVIDIA Corporation
  4. Built on Sun_Feb_28_22:34:44_PST_2021
  5. Cuda compilation tools, release 10.2, V10.2.300
  6. Build cuda_10.2_r440.TC440_70.29663091_0
  7. ljx@ljx-desktop:~/pycuda2/tensorrtx-yolov5-v5.0/yolov5$
  1. cd /usr/src/cudnn_samples_v8/mnistCUDNN
  2. sudo make
  3. sudo chmod a+x mnistCUDNN
  4. ./mnistCUDNN

二、安装好pip3，系统已经自带python3.6.9

sudo apt-get install python3-pip python3-dev

 三、检测是否安装gpu版本的tensorflow

1.安装方法之前的文章有这里举个例子

sudo apt-get install libhdf5-serial-dev hdf5-tools
pip3 install --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v46 tensorflow-gpu==2.6.0+nv19.3 --user

 2.检测方法举两个例子

  1. ljx@ljx-desktop:~/pycuda2/pycuda-2021.1$ python3
  2. Python 3.6.9 (default, Dec  8 2021, 21:08:43)
  3. [GCC 8.4.0] on linux
  4. Type "help", "copyright", "credits" or "license" for more information.
  5. >>> import torchvision
  6. >>> print(torchvision.__version__)
  7. 0.11.1
  8. >>> import tensorflow as tf
  9. >>> a = tf.constant(1.)
  10. 2022-02-21 21:25:38.178350: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  11. 2022-02-21 21:25:38.179671: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  12. 2022-02-21 21:25:38.180036: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  13. 2022-02-21 21:25:38.194555: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  14. 2022-02-21 21:25:38.196004: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  15. 2022-02-21 21:25:38.197011: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  16. 2022-02-21 21:26:58.812460: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  17. 2022-02-21 21:26:58.873885: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  18. 2022-02-21 21:26:58.909564: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  19. 2022-02-21 21:26:59.039953: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 41 MB memory:  -> device: 0, name: NVIDIA Tegra X1, pci                bus id: 0000:00:00.0, compute capability: 5.3
  20. >>> import os
  21. >>> os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
  22. >>> a = tf.constant(1.)
  23. >>> b = tf.constant(2.)
  24. >>> print(a+b)
  25. tf.Tensor(3.0, shape=(), dtype=float32)
  26. >>> print('GPU:', tf.test.is_gpu_available())
  27. WARNING:tensorflow:From <stdin>:1: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
  28. Instructions for updating:
  29. Use `tf.config.list_physical_devices('GPU')` instead.
  30. 2022-02-21 21:32:15.515633: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  31. 2022-02-21 21:32:15.517432: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  32. 2022-02-21 21:32:15.518313: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  33. 2022-02-21 21:32:15.527565: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  34. 2022-02-21 21:32:15.528595: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  35. 2022-02-21 21:32:15.529327: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /device:GPU:0 with 41 MB memory:  -> device: 0, name: NVIDIA Tegra X1, pci bus id: 0000:00:00.0, compute capability: 5.3
  36. GPU: True
  37. >>>
  1. ljx@ljx-desktop:~/pycuda2$ cat demo3.py
  2. import tensorflow as tf
  3.  
  4. tf.compat.v1.disable_eager_execution()
  5. with tf.device('/cpu:0'):
  6.     a = tf.constant([1.0,2.0,3.0],shape=[3],name='a')
  7.     b = tf.constant([1.0,2.0,3.0],shape=[3],name='b')
  8. with tf.device('/gpu:1'):
  9.     c = a+b
  10.  
  11. sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True,log_device_placement=True))
  12. sess.run(tf.compat.v1.global_variables_initializer())
  13. print(sess.run(c))
  14. ljx@ljx-desktop:~/pycuda2$ python3 demo3.py
  15. 2022-02-24 13:36:43.842123: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  16. 2022-02-24 13:36:45.249622: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  17. 2022-02-24 13:36:45.251626: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  18. 2022-02-24 13:38:19.897324: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  19. 2022-02-24 13:38:20.908341: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  20. 2022-02-24 13:38:20.941767: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1017] ARM64 does not support NUMA - returning NUMA node zero
  21. 2022-02-24 13:38:21.589736: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 39 MB memory:  -> device: 0, name: NVIDIA Tegra X1, pci bus id: 0000:00:00.0, compute capability: 5.3
  22. 2022-02-24 13:38:22.835843: I tensorflow/core/common_runtime/direct_session.cc:361] Device mapping:
  23. /job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: NVIDIA Tegra X1, pci bus id: 0000:00:00.0, compute capability: 5.3
  24.  
  25. add: (AddV2): /job:localhost/replica:0/task:0/device:GPU:0
  26. 2022-02-24 13:38:28.950709: I tensorflow/core/common_runtime/placer.cc:114] add: (AddV2): /job:localhost/replica:0/task:0/device:GPU:0
  27. init: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
  28. 2022-02-24 13:38:28.988627: I tensorflow/core/common_runtime/placer.cc:114] init: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
  29. a: (Const): /job:localhost/replica:0/task:0/device:CPU:0
  30. 2022-02-24 13:38:28.988931: I tensorflow/core/common_runtime/placer.cc:114] a: (Const): /job:localhost/replica:0/task:0/device:CPU:0
  31. b: (Const): /job:localhost/replica:0/task:0/device:CPU:0
  32. 2022-02-24 13:38:28.989207: I tensorflow/core/common_runtime/placer.cc:114] b: (Const): /job:localhost/replica:0/task:0/device:CPU:0
  33. [2. 4. 6.]

四、安装pycuda

官方解决方案【链接】
不想去看的话，直接下载这个链接的源码，同下步骤进行安装即可

 pycuda · PyPI

  1. tar zxvf pycuda-2021.1.tar.gz    
  2. cd pycuda-2021.1/  
  3. python3 configure.py --cuda-root=/usr/local/cuda-10.2
  4. sudo python3 setup.py install

demo测试

  1. ljx@ljx-desktop:~/pycuda2$ python3 demo2.py
  2. [[ 19.436962    39.908886    20.68723    ...  -8.1019335  -15.546103
  3.   -17.154585  ]
  4.  [-19.714169    -0.6291714    9.462954   ... -15.174974    -4.1439514
  5.    18.460089  ]
  6.  [-17.491064   -34.86578    -12.999788   ... -17.18811     10.867537
  7.     0.05436563]
  8.  ...
  9.  [ 45.716812   -32.27492     -0.5752983  ... -31.032787    -4.8378153
  10.     7.907672  ]
  11.  [  6.989045   -13.123575    -2.8372145  ...  21.856476     5.0534296
  12.   -15.905795  ]
  13.  [ 17.042442     0.354123    -7.9831614  ... -11.882836    20.23512
  14.   -19.761951  ]]
  15. [[ 19.436964    39.908894    20.687223   ...  -8.101934   -15.54609
  16.   -17.154581  ]
  17.  [-19.71417     -0.62916106   9.46296    ... -15.174983    -4.1439533
  18.    18.460089  ]
  19.  [-17.491072   -34.86579    -12.999789   ... -17.188126    10.867537
  20.     0.05437115]
  21.  ...
  22.  [ 45.716824   -32.27491     -0.57529545 ... -31.03278     -4.8378134
  23.     7.907671  ]
  24.  [  6.989043   -13.123584    -2.8372157  ...  21.856468     5.053428
  25.   -15.905798  ]
  26.  [ 17.042446     0.35412684  -7.98316    ... -11.882843    20.23511
  27.   -19.761948  ]]

  1. ljx@ljx-desktop:~/pycuda2$ cat demo2.py
  2. import numpy as np
  3. import pycuda.autoinit
  4. import pycuda.driver as cuda
  5. from pycuda.compiler import SourceModule
  6.  
  7.  
  8. mod = SourceModule("""
  9. #define BLOCK_SIZE 16
  11. typedef struct {
  12.     int width;
  13.     int height;
  14.     int stride;
  15.     int __padding;    //为了和64位的elements指针对齐
  16.     float* elements;
  17. } Matrix;
  19. // 读取矩阵元素
  20. __device__ float GetElement(const Matrix A, int row, int col)
  21. {
  22.     return A.elements[row * A.stride + col];
  23. }
  25. // 赋值矩阵元素
  26. __device__ void SetElement(Matrix A, int row, int col, float value)
  27. {
  28.     A.elements[row * A.stride + col] = value;
  29. }
  31. // 获取 16x16 的子矩阵
  32.  __device__ Matrix GetSubMatrix(Matrix A, int row, int col)
  33. {
  34.     Matrix Asub;
  35.     Asub.width    = BLOCK_SIZE;
  36.     Asub.height   = BLOCK_SIZE;
  37.     Asub.stride   = A.stride;
  38.     Asub.elements = &A.elements[A.stride * BLOCK_SIZE * row + BLOCK_SIZE * col];
  39.     return Asub;
  40. }
  42. __global__ void matrix_mul(Matrix *A, Matrix *B, Matrix *C)
  43. {
  44.     int blockRow = blockIdx.y;
  45.     int blockCol = blockIdx.x;
  46.     int row = threadIdx.y;
  47.     int col = threadIdx.x;
  49.     Matrix Csub = GetSubMatrix(*C, blockRow, blockCol);
  51.     // 每个线程通过累加Cvalue计算Csub的一个值
  52.     float Cvalue = 0;
  54.     // 为了计算Csub遍历所有需要的Asub和Bsub
  55.     for (int m = 0; m < (A->width / BLOCK_SIZE); ++m)
  56.     {
  57.         Matrix Asub = GetSubMatrix(*A, blockRow, m);
  58.         Matrix Bsub = GetSubMatrix(*B, m, blockCol);
  60.         __shared__ float As[BLOCK_SIZE][BLOCK_SIZE];
  61.         __shared__ float Bs[BLOCK_SIZE][BLOCK_SIZE];
  63.         As[row][col] = GetElement(Asub, row, col);
  64.         Bs[row][col] = GetElement(Bsub, row, col);
  66.         __syncthreads();
  68.         for (int e = 0; e < BLOCK_SIZE; ++e)
  69.             Cvalue += As[row][e] * Bs[e][col];
  71.         __syncthreads();
  72.     }
  74.     SetElement(Csub, row, col, Cvalue);
  75. }
  76. """)
  77.  
  78.  
  79. class MatrixStruct(object):
  80.     def __init__(self, array):
  81.         self._cptr = None
  82.  
  83.         self.shape, self.dtype = array.shape, array.dtype
  84.         self.width = np.int32(self.shape[1])
  85.         self.height = np.int32(self.shape[0])
  86.         self.stride = self.width
  87.         self.elements = cuda.to_device(array)                      # 分配内存并拷贝数组数据至device，返回其地址
  88.  
  89.     def send_to_gpu(self):
  90.         self._cptr = cuda.mem_alloc(self.nbytes())                 # 分配一个C结构体所占的内存
  91.         cuda.memcpy_htod(int(self._cptr), self.width.tobytes())    # 拷贝数据至device，下同
  92.         cuda.memcpy_htod(int(self._cptr)+4, self.height.tobytes())
  93.         cuda.memcpy_htod(int(self._cptr)+8, self.stride.tobytes())
  94.         cuda.memcpy_htod(int(self._cptr)+16, np.intp(int(self.elements)).tobytes())
  95.  
  96.     def get_from_gpu(self):
  97.         return cuda.from_device(self.elements, self.shape, self.dtype)  # 从device取回数组数据
  98.  
  99.     def nbytes(self):
  100.         return self.width.nbytes * 4 + np.intp(0).nbytes
  101.  
  102.  
  103. a = np.random.randn(400,400).astype(np.float32)
  104. b = np.random.randn(400,400).astype(np.float32)
  105. c = np.zeros_like(a)
  106.  
  107. A = MatrixStruct(a)
  108. B = MatrixStruct(b)
  109. C = MatrixStruct(c)
  110. A.send_to_gpu()
  111. B.send_to_gpu()
  112. C.send_to_gpu()
  113.  
  114. matrix_mul = mod.get_function("matrix_mul")
  115. matrix_mul(A._cptr, B._cptr, C._cptr, block=(16,16,1), grid=(25,25))
  116. result = C.get_from_gpu()
  117. print(np.dot(a,b))
  118. print(result)

五、下载tensorrtx源码

进入tensorrtx的官网，下载你训练时对应的yolov5的版本，点击左上角的master-->tags-->yolov5

下载完成后，来到下载目录下，输入以下命令解压,我这里是v5.0版本

unzip tensorrtx-yolov5-v5.0.zip

 把之前训练的模型生成的wts权重文件放到tensorrtx的yolov5文件夹中

没有wts文件只是想体验强大的jetson nano的同学可以先下载一下五类垃圾分类权重文件https://blog.csdn.net/xiaoyuan2157

链接: https://pan.baidu.com/s/1nciB7Xn1vXj9ZfBAoj39Bw 提取码: r74h 

来到tensorrtx的yolov5文件夹，打开yololayer.h的代码，修改CLASS_NUM

创建进入文件夹build并cmake ..

  1. mkdir build
  2. cd build
  3. cmake ..
  4. make

 生成引擎文件

sudo ./yolov5 -s ../best.wts best.engine s

这一段模型引擎生成的命令解释如下

sudo ./yolov5 -s/ [.wts文件路径] [.engine文件名称] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw]

稍作等待后，出现Build engine successfully！表示生成完成，这时build文件夹里面会多出一个best.engine文件

六、模型测试

根据官方的yolov5_trt改的代码来测试一下 

  1. ljx@ljx-desktop:~/pycuda2/tensorrtx-yolov5-v5.0/yolov5$ cat yolov5_trt2.py
  2. """
  3. # Yolov5 基于pytorch，修改起来更加方便快捷；
  4. # yolov5自带anchor生成器，自动为你的数据集生成最优化的anchor；
  5. # yolov5的整体AP比yolov4更高。
  6. """
  7. import ctypes
  8. import os
  9. import random
  10. import sys
  11. import threading
  12. import time
  13. # 安装串口函数库 sudo pip3 install pyserial
  14. import serial
  15. import serial as se  # 导入串口库,这里是用于串口通信的库，需要在命令行输入
  16. #pip3 install pyserial
  17. import cv2
  18. import numpy as np  # 构造ndarray对象
  19. import pycuda.autoinit
  20. import pycuda.driver as cuda
  21. import tensorrt as trt
  22. from time import sleep
  23.  
  24.  
  25. # from jetcam.csi_camera import CSICamera
  26. # import torch
  27. # import torchvision#在nano上安装这两个库是有些麻烦的特别是torchvision。
  28.  
  29. INPUT_W = 640
  30. INPUT_H = 640
  31. CONF_THRESH = 0.8  # 概率阈值
  32. IOU_THRESHOLD = 0.1
  33.  
  34.  
  35. # 定义画框函数
  36. def plot_one_box(x, img, color=None, label=None, line_thickness=None):
  37.     '''
  38.     description: Plots one bounding box on image img,
  39.                  this function comes from YoLov5 project.
  40.     param:
  41.         x:      a box likes [x1,y1,x2,y2]
  42.         img:    a opencv image object
  44.         label:  str
  45.         line_thickness: int
  46.     return:
  47.         no return
  48.     '''
  49.  
  50.     # img, result_boxes, result_scores, result_classid = yolov5_wrapper.infer(img)
  51.     # img = draw_boxes(img, result_boxes, result_scores, result_classid)
  52.     tl = (
  53.             line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
  54.     )  # line/font thickness
  55.     color = color or [random.randint(0, 255) for _ in range(3)]
  56.     c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
  57.     cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
  58.     # print("left:(" + str(c1[0]) + "," + str(c1[1]) +")","right:(" + str(c2[0]) + "," + str(c2[1])+ ")")
  59.     a = int(c1[0])
  60.     b = int(c2[0])
  61.     c = int(c1[1])
  62.     d = int(c2[1])
  63.     x1 = (b + a) / 2
  64.  
  65.     x = int(x1)
  66.     y1 = (d + c) / 2
  67.     y = int(y1)
  68.     r = label[2:6] #rate
  69.     sleep(0.0009)
  70.     c =str(label[0]) #class
  71.  
  72.     if label:
  73.         tf = max(tl - 1, 1)  # font thickness
  74.         t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
  75.         c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
  76.         cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
  77.         cv2.putText(
  78.             img,
  79.             label,
  80.             (c1[0], c1[1] - 2),
  81.             0,
  82.             tl / 3,
  83.             [225, 255, 255],
  84.             thickness=tf,
  85.             lineType=cv2.LINE_AA,
  86.         )
  87.     return x, y
  88.  
  89.  
  90. # 画框函数
  91. def draw_boxes(image_raw, result_boxes, result_scores, result_classid):
  92.     max_scores = -1
  93.     max_index = -1
  94.     max_x,max_y = -1,-1
  95.     for i in range(len(result_boxes)):
  96.         box = result_boxes[i]
  97.         x, y = plot_one_box(
  98.             box,
  99.             image_raw,
  100.             label="{}:{:.2f}".format(
  101.                 categories[int(result_classid[i])], result_scores[i]
  102.             )
  103.         )
  104.         # print(result_classid[i])
  105.         # se.write((str(x) + ',' + str(y) + ',' + str(result_classid[i]) + '\r\n').encode())
  106.         # global max_score
  107.         if result_boxes.all() > max_scores:
  108.             max_scores = result_scores[i]
  109.             max_index = i
  110.             max_x, max_y = x, y
  111.  
  112.     if max_scores != -1:
  113.         c = int(result_classid[max_index])
  114.         output_str = ('[' + str(x) + ',' + str(y) + ',' +str(c) + ']'+'\r\n')
  115.         print(output_str)
  116.         se.write(output_str.encode())
  117.         sleep(0.0009)
  118.  
  119.     return image_raw
  120.  
  121.  
  122. # yolov5模型转到TensorRT中推理
  123. # 定义yolov5转trt的类 start
  124. class YoLov5TRT(object):
  125.     """
  126.     description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops.
  127.     """
  128.  
  129.     def __init__(self, engine_file_path):
  130.         # Create a Context on this device,
  131.         self.ctx = cuda.Device(0).make_context()
  132.         stream = cuda.Stream()
  133.         TRT_LOGGER = trt.Logger(trt.Logger.INFO)
  134.         runtime = trt.Runtime(TRT_LOGGER)
  135.  
  136.         # Deserialize the engine from file
  137.         with open(engine_file_path, "rb") as f:
  138.             engine = runtime.deserialize_cuda_engine(f.read())
  139.         context = engine.create_execution_context()
  140.  
  141.         host_inputs = []
  142.         cuda_inputs = []
  143.         host_outputs = []
  144.         cuda_outputs = []
  145.         bindings = []
  146.  
  147.         for binding in engine:
  148.             size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
  149.             dtype = trt.nptype(engine.get_binding_dtype(binding))
  150.             # Allocate host and device buffers
  151.             host_mem = cuda.pagelocked_empty(size, dtype)
  152.             cuda_mem = cuda.mem_alloc(host_mem.nbytes)
  153.             # Append the device buffer to device bindings.
  154.             bindings.append(int(cuda_mem))
  155.             # Append to the appropriate list.
  156.             if engine.binding_is_input(binding):
  157.                 host_inputs.append(host_mem)
  158.                 cuda_inputs.append(cuda_mem)
  159.             else:
  160.                 host_outputs.append(host_mem)
  161.                 cuda_outputs.append(cuda_mem)
  162.  
  163.         # Store
  164.         self.stream = stream
  165.         self.context = context
  166.         self.engine = engine
  167.         self.host_inputs = host_inputs
  168.         self.cuda_inputs = cuda_inputs
  169.         self.host_outputs = host_outputs
  170.         self.cuda_outputs = cuda_outputs
  171.         self.bindings = bindings
  172.  
  173.     # 释放引擎，释放GPU显存，释放CUDA流
  174.     def __del__(self):
  175.         print("delete object to release memory")
  176.  
  177.     def infer(self, input_image_path):
  178.         threading.Thread.__init__(self)
  179.         # Make self the active context, pushing it on top of the context stack.
  180.         self.ctx.push()
  181.         # Restore
  182.         stream = self.stream
  183.         context = self.context
  184.         engine = self.engine
  185.         host_inputs = self.host_inputs
  186.         cuda_inputs = self.cuda_inputs
  187.         host_outputs = self.host_outputs
  188.         cuda_outputs = self.cuda_outputs
  189.         bindings = self.bindings
  190.         # Do image preprocess
  191.         input_image, image_raw, origin_h, origin_w = self.preprocess_image(
  192.             input_image_path
  193.         )
  194.         # Copy input image to host buffer
  195.         np.copyto(host_inputs[0], input_image.ravel())
  196.         start = time.time()
  197.         # Transfer input data  to the GPU.
  198.         cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
  199.         # Run inference.
  200.         context.execute_async(bindings=bindings, stream_handle=stream.handle)
  201.         # Transfer predictions back from the GPU.
  202.         cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
  203.         # Synchronize the stream
  204.         stream.synchronize()
  205.         end = time.time()
  206.         # Remove any context from the top of the context stack, deactivating it.
  207.         self.ctx.pop()
  208.         # Here we use the first row of output in that batch_size = 1
  209.         output = host_outputs[0]
  210.         # Do postprocess
  211.         result_boxes, result_scores, result_classid = self.post_process(
  212.             output, origin_h, origin_w
  213.         )
  214.  
  215.  
  216.         return image_raw, result_boxes, result_scores, result_classid
  217.  
  218.     def destroy(self):
  219.         # Remove any context from the top of the context stack, deactivating it.
  220.         self.ctx.pop()
  221.  
  222.     def preprocess_image(self, image_raw):
  223.         """
  224.         description: Read an image from image path, convert it to RGB,
  225.                      resize and pad it to target size, normalize to [0,1],
  226.                      transform to NCHW format.
  227.         param:
  228.             input_image_path: str, image path
  229.         return:
  230.             image:  the processed image
  231.             image_raw: the original image
  232.             h: original height
  233.             w: original width
  234.         """
  235.         h, w, c = image_raw.shape
  236.         image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)
  237.         # Calculate widht and height and paddings
  238.         r_w = INPUT_W / w
  239.         r_h = INPUT_H / h
  240.         if r_h > r_w:
  241.             tw = INPUT_W
  242.             th = int(r_w * h)
  243.             tx1 = tx2 = 0
  244.             ty1 = int((INPUT_H - th) / 2)
  245.             ty2 = INPUT_H - th - ty1
  246.         else:
  247.             tw = int(r_h * w)
  248.             th = INPUT_H
  249.             tx1 = int((INPUT_W - tw) / 2)
  250.             tx2 = INPUT_W - tw - tx1
  251.             ty1 = ty2 = 0
  252.         # Resize the image with long side while maintaining ratio
  253.         image = cv2.resize(image, (tw, th))
  254.         # Pad the short side with (128,128,128)
  255.         image = cv2.copyMakeBorder(
  256.             image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128)
  257.         )
  258.         image = image.astype(np.float32)
  259.         # Normalize to [0,1]
  260.         image /= 255.0
  261.         # HWC to CHW format:
  262.         image = np.transpose(image, [2, 0, 1])
  263.         # CHW to NCHW format
  264.         image = np.expand_dims(image, axis=0)
  265.         # Convert the image to row-major order, also known as "C order":
  266.         image = np.ascontiguousarray(image)
  267.         return image, image_raw, h, w
  268.  
  269.     def xywh2xyxy(self, origin_h, origin_w, x):
  270.         """
  271.         description:    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
  272.         param:
  273.             origin_h:   height of original image
  274.             origin_w:   width of original image
  275.             x:          A boxes tensor, each row is a box [center_x, center_y, w, h]
  276.         return:
  277.             y:          A boxes tensor, each row is a box [x1, y1, x2, y2]
  278.         """
  279.         y = np.zeros_like(x)
  280.         # y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
  281.         r_w = INPUT_W / origin_w
  282.         r_h = INPUT_H / origin_h
  283.         if r_h > r_w:
  284.             y[:, 0] = x[:, 0] - x[:, 2] / 2
  285.             y[:, 2] = x[:, 0] + x[:, 2] / 2
  286.             y[:, 1] = x[:, 1] - x[:, 3] / 2 - (INPUT_H - r_w * origin_h) / 2
  287.             y[:, 3] = x[:, 1] + x[:, 3] / 2 - (INPUT_H - r_w * origin_h) / 2
  288.             y /= r_w
  289.         else:
  290.             y[:, 0] = x[:, 0] - x[:, 2] / 2 - (INPUT_W - r_h * origin_w) / 2
  291.             y[:, 2] = x[:, 0] + x[:, 2] / 2 - (INPUT_W - r_h * origin_w) / 2
  292.             y[:, 1] = x[:, 1] - x[:, 3] / 2
  293.             y[:, 3] = x[:, 1] + x[:, 3] / 2
  294.             y /= r_h
  295.  
  296.         return y
  297.  
  298.     # 往YoLov5TRT这个类中加入一个方法，此处是用numpy的方式实现nms
  299.     def nms(self, boxes, scores, iou_threshold=IOU_THRESHOLD):  # 非极大值抑制，是目标检测框架中的后处理模块
  300.         # 空间距离结合并交比(IOU)完成聚类划分
  301.         x1 = boxes[:, 0]
  302.         y1 = boxes[:, 1]
  303.         x2 = boxes[:, 2]
  304.         y2 = boxes[:, 3]
  305.         areas = (y2 - y1 + 1) * (x2 - x1 + 1)
  306.         scores = scores
  307.         keep = []
  308.         index = scores.argsort()[::-1]
  309.         while index.size > 0:
  310.             i = index[0]  # every time the first is the biggst, and add it directly
  311.             keep.append(i)
  312.  
  313.             x11 = np.maximum(x1[i], x1[index[1:]])  # calculate the points of overlap
  314.             y11 = np.maximum(y1[i], y1[index[1:]])
  315.             x22 = np.minimum(x2[i], x2[index[1:]])
  316.             y22 = np.minimum(y2[i], y2[index[1:]])
  317.  
  318.             w = np.maximum(0, x22 - x11 + 1)  # the weights of overlap
  319.             h = np.maximum(0, y22 - y11 + 1)  # the height of overlap
  320.  
  321.             overlaps = w * h
  322.             ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
  323.  
  324.             idx = np.where(ious <= iou_threshold)[0]
  325.             index = index[idx + 1]  # because index start from 1
  326.             # print(overlaps)
  327.             # print(x1)
  328.             # sleep(1)
  329.  
  330.         return keep
  331.  
  332.     # 把nms的结果赋值给indices变量，改写post_process函数
  333.     def post_process(self, output, origin_h, origin_w):
  334.         """
  335.         description: postprocess the prediction
  336.         param:
  337.             output:     A tensor likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...]
  338.             origin_h:   height of original image
  339.             origin_w:   width of original image
  340.         return:
  341.             result_boxes: finally boxes, a boxes tensor, each row is a box [x1, y1, x2, y2]
  342.             result_scores: finally scores, a tensor, each element is the score correspoing to box
  343.             result_classid: finally classid, a tensor, each element is the classid correspoing to box
  344.         """
  345.         # Get the num of boxes detected
  346.         num = int(output[0])
  347.         # Reshape to a two dimentional ndarray
  348.         pred = np.reshape(output[1:], (-1, 6))[:num, :]
  349.         # to a torch Tensor
  350.         # pred = torch.Tensor(pred).cuda()#去掉这行，用torchvision库中的nms方法来完成非极大值抑制。
  351.         # Get the boxes
  352.         boxes = pred[:, :4]
  353.         # Get the scores
  354.         scores = pred[:, 4]
  355.         # Get the classid
  356.         classid = pred[:, 5]
  357.         # Choose those boxes that score > CONF_THRESH
  358.         si = scores > CONF_THRESH
  359.         boxes = boxes[si, :]
  360.         scores = scores[si]
  361.         classid = classid[si]
  362.         # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]
  363.         boxes = self.xywh2xyxy(origin_h, origin_w, boxes)
  364.         # Do nms
  365.         # 去掉cpu方法，因为ndarray没有这个方法
  366.         # indices = torchvision.ops.nms(boxes, scores, iou_threshold=IOU_THRESHOLD).cpu()
  367.         # result_boxes = boxes[indices, :].cpu()
  368.         # result_scores = scores[indices].cpu()
  369.         # result_classid = classid[indices].cpu()
  370.  
  371.         indices = self.nms(boxes, scores, IOU_THRESHOLD)
  372.         result_boxes = boxes[indices, :]
  373.         result_scores = scores[indices]
  374.         result_classid = classid[indices]
  375.         # print(result_boxes)
  376.         # print(result_classid)
  377.  
  378.         return result_boxes, result_scores, result_classid
  379.  
  380.  
  381. class myThread(threading.Thread):
  382.     def __init__(self, func, args):
  383.         threading.Thread.__init__(self)
  384.         self.func = func
  385.         self.args = args
  386.  
  387.     def run(self):
  388.         self.func(*self.args)
  389.  
  390.  
  391. # 摄像头检测
  392. def detect_camera(camera, yolov5_wrapper):
  393.     # def detect_camera(x,camera, yolov5_wrapper):
  394.     count = 0
  395.  
  396.     # 开始循环检测
  397.     while True:
  398.         # img = camera.read()#CSI摄像头
  399.         ret, img = camera.read()  # usb摄像头用这个
  400.         img, result_boxes, result_scores, result_classid = yolov5_wrapper.infer(img)
  401.         img = draw_boxes(img, result_boxes, result_scores, result_classid)
  402.  
  403.         count = count + 1
  404.  
  405.         cv2.imshow("result", img)  # 显示结果
  406.         if cv2.waitKey(1) == ord('q'):
  407.             break
  408.  
  409.  
  410. # 定义摄像头函数
  411. def main_camera():
  412.     camera = cv2.VideoCapture(0)  # usb摄像头用这个
  413.     # camera = CSICamera(capture_device=0, width=640, height=480)
  414.     # load custom plugins
  415.     camera.set(3, 640)
  416.     camera.set(4, 480)
  417.     PLUGIN_LIBRARY = "build/libmyplugins.so"
  418.     ctypes.CDLL(PLUGIN_LIBRARY)
  419.     engine_file_path = "build/yolov5s.engine"
  420.  
  421.     # YoLov5TRT instance
  422.     yolov5_wrapper = YoLov5TRT(engine_file_path)
  423.     print("start detection!")
  424.     detect_camera(camera, yolov5_wrapper)
  425.     # camera.release() #  使用cv方法打开摄像头才需要这句
  426.     cv2.destroyAllWindows()
  427.     print("\nfinish!")
  428.  
  429.  
  430. if __name__ == "__main__":
  431.     # load custom plugins      修改成你build出来的引擎的相对路径
  432.     PLUGIN_LIBRARY = "build/libmyplugins.so"
  433.     ctypes.CDLL(PLUGIN_LIBRARY)
  434.     engine_file_path = "build/yolov5s.engine"
  435.     se = serial.Serial('/dev/ttyTHS1', 115200, timeout=0.5)  # 设置使用的引脚、波特率和超时时间 8接R,10接T
  436.     # load coco labels
  437.  
  438.     # categories = ['battery', 'orange', 'bottle', 'paper_cup', 'spitball']  # 垃圾种类
  439.     categories = ['0', '1', '2', '3', '4']  # 垃圾种类
  440.  
  441.  
  442.     main_camera()

都是按照大佬们的博客复制学习的，真尴尬哈哈

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/article/detail/42164?site
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号