devbox
繁依Fanyi0
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

YOLOv8-pose之tensorrt版本_yolov8分类模型 tensorrt

作者：繁依Fanyi0 | 2024-08-21 20:52:40

赞

踩

yolov8分类模型 tensorrt

目录

一、环境安装

1、创建一个虚拟环境

2、下载TensorRT

3.TensorRT安装

4、其他库安装

二、模型转换及代码

（一）模型转换

导出ONNX模型

导出engine模型

（二）代码

完整的tensorrt推理代码：

三、遇到的问题

（一）模型转换遇到的问题

（二）代码使用注意

1、模型修改地方

2、输入参数修改

3、关键点的颜色和大小修改

4、骨架连接顺序和线宽修改

5、代码使用，看注释

6、项目目录结构

参考文献：

一、环境安装

1、创建一个虚拟环境
conda create -n 环境名 python=X.X
2、下载TensorRT

进入官方网站：Log in | NVIDIA Developer
寻找自己对应的版本，我这里选择为：

下载得到 zip 压缩包，解压。

3.TensorRT安装

 1、复制TensorRT-8.4.3.1\bin中内容到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.3\bin

 2、复制TensorRT的include文件夹到CUDA的include文件夹

 3、复制TensorRT-8.4.3.1\lib文件夹中的lib文件到CUDA的lib文件夹，dll文件到CUDA的bin文件夹（注意lib文件和dll文件要分开）

4、使用pip install xxx.whl安装TensorRT-8.4.3.1文件夹中框出来的（里面有whl文件，直接在创建的虚拟环境中pip安装）

使用python检查是否安装成功。

  1. import tensorrt as trt
  2.  
  3. trt.__version__

4、其他库安装
pip install pycuda

  pip install opencv-contrib-python

二、模型转换及代码

（一）模型转换

首先需要将pth模型转换为ONNX模型，再将ONNX模型转换为engine模型。

导出ONNX模型

使用如下脚本：（需要安装ultralytics，直接pip安装即可）

  1. from ultralytics import YOLO
  2.  
  3. # Load a model
  4. model = YOLO("yolov8s-pose.pt")  # load a pretrained model (recommended for training)
  5. success = model.export(format="onnx", opset=11, simplify=True)  # export the model to onnx format
  6. assert success
导出engine模型

使用如下命令：

trtexec.exe --onnx='你的onnx模型'.onnx  --saveEngine='保存的名称'.engine --fp16
（二）代码
完整的tensorrt推理代码：
  1. '''
  2. Author: [egrt]
  3. Date: 2023-03-26 09:39:21
  4. LastEditors: Egrt
  5. LastEditTime: 2023-07-15 22:10:25
  6. Description:
  7. '''
  8. import numpy as np
  9. import time
  10. import tensorrt as trt
  11. import pycuda.driver as cuda
  12. import pycuda.autoinit
  13. import cv2
  14. from numpy import array
  15.  
  16.  
  17. def resize_image(image, size, letterbox_image):
  18.     ih, iw = image.shape[:2]
  19.     h, w = size
  20.     if letterbox_image:
  21.         scale = min(w / iw, h / ih)
  22.         nw = int(iw * scale)
  23.         nh = int(ih * scale)
  24.  
  25.         image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_CUBIC)
  26.         new_image = 128 * np.ones((h, w, 3), dtype=np.uint8)
  27.         new_image[(h - nh) // 2:(h - nh) // 2 + nh, (w - nw) // 2:(w - nw) // 2 + nw, :] = image
  28.     else:
  29.         new_image = cv2.resize(image, (w, h), interpolation=cv2.INTER_CUBIC)
  30.     scale = [iw / w, ih / h]
  31.     return new_image, scale
  32.  
  33.  
  34. def preprocess_input(image):
  35.     image /= 255.0
  36.     return image
  37.  
  38.  
  39. def xywh2xyxy(x):
  40.     """
  41.     Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is the
  42.     top-left corner and (x2, y2) is the bottom-right corner.
  44.     Args:
  45.         x (np.ndarray | torch.Tensor): The input bounding box coordinates in (x, y, width, height) format.
  46.     Returns:
  47.         y (np.ndarray | torch.Tensor): The bounding box coordinates in (x1, y1, x2, y2) format.
  48.     """
  49.     y = np.copy(x)
  50.     y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
  51.     y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
  52.     y[..., 2] = x[..., 0] + x[..., 2] / 2  # bottom right x
  53.     y[..., 3] = x[..., 1] + x[..., 3] / 2  # bottom right y
  54.     return y
  55.  
  56.  
  57. def box_area(boxes: array):
  58.     """
  59.     :param boxes: [N, 4]
  60.     :return: [N]
  61.     """
  62.     return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
  63.  
  64.  
  65. def box_iou(box1: array, box2: array):
  66.     """
  67.     :param box1: [N, 4]
  68.     :param box2: [M, 4]
  69.     :return: [N, M]
  70.     """
  71.     area1 = box_area(box1)  # N
  72.     area2 = box_area(box2)  # M
  73.     # broadcasting, 两个数组各维度大小 从后往前对比一致， 或者 有一维度值为1；
  74.     lt = np.maximum(box1[:, np.newaxis, :2], box2[:, :2])
  75.     rb = np.minimum(box1[:, np.newaxis, 2:], box2[:, 2:])
  76.     wh = rb - lt  # 右下角 - 左上角；
  77.     wh = np.maximum(0, wh)  # [N, M, 2]
  78.     inter = wh[:, :, 0] * wh[:, :, 1]
  79.     iou = inter / (area1[:, np.newaxis] + area2 - inter)
  80.     return iou  # NxM
  81.  
  82.  
  83. def numpy_nms(boxes: array, scores: array, iou_threshold: float):
  84.     idxs = scores.argsort()  # 按分数 降序排列的索引 [N]
  85.     keep = []
  86.     while idxs.size > 0:  # 统计数组中元素的个数
  87.         max_score_index = idxs[-1]
  88.         max_score_box = boxes[max_score_index][None, :]
  89.         keep.append(max_score_index)
  90.  
  91.         if idxs.size == 1:
  92.             break
  93.         idxs = idxs[:-1]  # 将得分最大框 从索引中删除； 剩余索引对应的框 和 得分最大框 计算IoU；
  94.         other_boxes = boxes[idxs]  # [?, 4]
  95.         ious = box_iou(max_score_box, other_boxes)  # 一个框和其余框比较 1XM
  96.         idxs = idxs[ious[0] <= iou_threshold]
  97.  
  98.     keep = np.array(keep)  # Tensor
  99.     return keep
  100.  
  101.  
  102. def non_max_suppression(
  103.         prediction,
  104.         conf_thres=0.25,
  105.         iou_thres=0.45,
  106.         classes=None,
  107.         agnostic=False,
  108.         multi_label=False,
  109.         labels=(),
  110.         max_det=300,
  111.         nc=0,  # number of classes (optional)
  112.         max_time_img=0.05,
  113.         max_nms=30000,
  114.         max_wh=7680,
  115. ):
  116.     """
  117.     Perform non-maximum suppression (NMS) on a set of boxes, with support for masks and multiple labels per box.
  119.     Arguments:
  120.         prediction (np.ndarray): An array of shape (batch_size, num_classes + 4 + num_masks, num_boxes)
  121.             containing the predicted boxes, classes, and masks. The array should be in the format
  122.             output by a model, such as YOLO.
  123.         conf_thres (float): The confidence threshold below which boxes will be filtered out.
  124.             Valid values are between 0.0 and 1.0.
  125.         iou_thres (float): The IoU threshold below which boxes will be filtered out during NMS.
  126.             Valid values are between 0.0 and 1.0.
  127.         classes (List[int]): A list of class indices to consider. If None, all classes will be considered.
  128.         agnostic (bool): If True, the model is agnostic to the number of classes, and all
  129.             classes will be considered as one.
  130.         multi_label (bool): If True, each box may have multiple labels.
  131.         labels (List[List[Union[int, float, np.ndarray]]]): A list of lists, where each inner
  132.             list contains the apriori labels for a given image. The list should be in the format
  133.             output by a dataloader, with each label being a tuple of (class_index, x1, y1, x2, y2).
  134.         max_det (int): The maximum number of boxes to keep after NMS.
  135.         nc (int, optional): The number of classes output by the model. Any indices after this will be considered masks.
  136.         max_time_img (float): The maximum time (seconds) for processing one image.
  137.         max_nms (int): The maximum number of boxes into torchvision.ops.nms().
  138.         max_wh (int): The maximum box width and height in pixels
  140.     Returns:
  141.         (List[np.ndarray]): A list of length batch_size, where each element is an array of
  142.             shape (num_boxes, 6 + num_masks) containing the kept boxes, with columns
  143.             (x1, y1, x2, y2, confidence, class, mask1, mask2, ...).
  144.     """
  145.  
  146.     # Checks
  147.     assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
  148.     assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
  149.     if isinstance(prediction, (list, tuple)):  # YOLOv8 model in validation model, output = (inference_out, loss_out)
  150.         prediction = prediction[0]  # select only inference output
  151.  
  152.     bs = prediction.shape[0]  # batch size
  153.     nc = nc or (prediction.shape[1] - 4)  # number of classes
  154.     nm = prediction.shape[1] - nc - 4
  155.     mi = 4 + nc  # mask start index
  156.     xc = prediction[:, 4:mi].max(axis=1) > conf_thres  # candidates
  157.  
  158.     # Settings
  159.     # min_wh = 2  # (pixels) minimum box width and height
  160.     time_limit = 0.5 + max_time_img * bs  # seconds to quit after
  161.     redundant = True  # require redundant detections
  162.     multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
  163.     merge = False  # use merge-NMS
  164.  
  165.     prediction = np.transpose(prediction, (0, 2, 1))  # shape(1,84,6300) to shape(1,6300,84)
  166.     prediction[..., :4] = xywh2xyxy(prediction[..., :4])  # xywh to xyxy
  167.  
  168.     t = time.time()
  169.     output = [np.zeros((0, 6 + nm)) for _ in range(bs)]
  170.     for xi, x in enumerate(prediction):  # image index, image inference
  171.         # Apply constraints
  172.         # x[((x[:, 2:4] < min_wh) | (x[:, 2:4] > max_wh)).any(1), 4] = 0  # width-height
  173.         x = x[xc[xi]]  # confidence
  174.  
  175.         # Cat apriori labels if autolabelling
  176.         if labels and len(labels[xi]):
  177.             lb = labels[xi]
  178.             v = np.zeros((len(lb), nc + nm + 5))
  179.             v[:, :4] = lb[:, 1:5]  # box
  180.             v[np.arange(len(lb)), lb[:, 0].astype(int) + 4] = 1.0  # cls
  181.             x = np.concatenate((x, v), axis=0)
  182.  
  183.         # If none remain process next image
  184.         if not x.shape[0]:
  185.             continue
  186.  
  187.         # Detections matrix nx6 (xyxy, conf, cls)
  188.         box, cls, mask = np.split(x, (4, 4 + nc), axis=1)
  189.  
  190.         if multi_label:
  191.             i, j = np.where(cls > conf_thres)
  192.             x = np.concatenate((box[i], x[i, 4 + j, None], j[:, None].astype(float), mask[i]), axis=1)
  193.         else:  # best class only
  194.             conf = np.max(cls, axis=1, keepdims=True)
  195.             j = np.argmax(cls, axis=1)
  196.             j = np.expand_dims(j, axis=1)
  197.             x = np.concatenate((box, conf, j.astype(float), mask), axis=1)[conf.reshape(-1) > conf_thres]
  198.  
  199.         # Filter by class
  200.         if classes is not None:
  201.             class_indices = np.array(classes)
  202.             mask = np.any(x[:, 5:6] == class_indices, axis=1)
  203.             x = x[mask]
  204.  
  205.         # Apply finite constraint
  206.         # if not np.isfinite(x).all():
  207.         #     x = x[np.isfinite(x).all(axis=1)]
  208.  
  209.         # Check shape
  210.         n = x.shape[0]  # number of boxes
  211.         if not n:  # no boxes
  212.             continue
  213.         if n > max_nms:  # excess boxes
  214.             sorted_indices = np.argsort(x[:, 4])[::-1]
  215.             x = x[sorted_indices[:max_nms]]  # sort by confidence and remove excess boxes
  216.  
  217.         # Batched NMS
  218.         c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
  219.         boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
  220.         i = numpy_nms(boxes, scores, iou_thres)  # NMS
  221.         i = i[:max_det]  # limit detections
  222.         if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
  223.             # Update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
  224.             iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
  225.             weights = iou * scores[None]  # box weights
  226.             x[i, :4] = np.dot(weights, x[:, :4]).astype(float) / weights.sum(1, keepdims=True)  # merged boxes
  227.             if redundant:
  228.                 i = i[np.sum(iou, axis=1) > 1]  # require redundancy
  229.  
  230.         output[xi] = x[i]
  231.         if (time.time() - t) > time_limit:
  232.             break  # time limit exceeded
  233.  
  234.     return output
  235.  
  236.  
  237. class YOLO(object):
  238.     _defaults = {
  239.         # ---------------------------------------------------------------------#
  240.         #   模型文件存放的路径
  241.         # ---------------------------------------------------------------------#
  242.         "model_path": 'yolov8-pose.engine',
  243.         # ---------------------------------------------------------------------#
  244.         #   输入图像的分辨率大小
  245.         # ---------------------------------------------------------------------#
  246.         "input_shape": [640, 640],
  247.         # ---------------------------------------------------------------------#
  248.         #   只有得分大于置信度的预测框会被保留下来
  249.         # ---------------------------------------------------------------------#
  250.         "confidence": 0.5,
  251.         # ---------------------------------------------------------------------#
  252.         #   非极大抑制所用到的nms_iou大小
  253.         # ---------------------------------------------------------------------#
  254.         "nms_iou": 0.3,
  255.     }
  256.  
  257.     @classmethod
  258.     def get_defaults(cls, n):
  259.         if n in cls._defaults:
  260.             return cls._defaults[n]
  261.         else:
  262.             return "Unrecognized attribute name '" + n + "'"
  263.  
  264.     # ---------------------------------------------------#
  265.     #   初始化YOLO
  266.     # ---------------------------------------------------#
  267.     def __init__(self, **kwargs):
  268.         self.__dict__.update(self._defaults)
  269.         for name, value in kwargs.items():
  270.             setattr(self, name, value)
  271.             self._defaults[name] = value
  272.  
  273.         # ---------------------------------------------------#
  274.         #   获得种类和先验框的数量
  275.         # ---------------------------------------------------#
  276.         self.class_names = ['person']
  277.         self.num_classes = len(self.class_names)
  278.         self.kpts_shape = [17, 3]
  279.         self.bbox_color = (150, 0, 0)
  280.         self.bbox_thickness = 6
  281.         # 框类别文字
  282.         self.bbox_labelstr = {
  283.             'font_size': 1,  # 字体大小
  284.             'font_thickness': 2,  # 字体粗细
  285.             'offset_x': 0,  # X 方向，文字偏移距离，向右为正
  286.             'offset_y': -10,  # Y 方向，文字偏移距离，向下为正
  287.         }
  288.         # 关键点 BGR 配色
  289.         self.kpt_color_map = {
  290.             0: {'color': [255, 128, 0], 'radius': 3},  # radius  半径
  291.             1: {'color': [255, 153, 51], 'radius': 3},  #
  292.             2: {'color': [255, 178, 102], 'radius': 3},  #
  293.             3: {'color': [230, 230, 0], 'radius': 3},
  294.             4: {'color': [255, 153, 255], 'radius': 3},
  295.             5: {'color': [153, 204, 255], 'radius': 3},
  296.             6: {'color': [255, 102, 255], 'radius': 3},
  297.             7: {'color': [255, 51, 255], 'radius': 3},
  298.             8: {'color': [102, 178, 255], 'radius': 3},
  299.             9: {'color': [51, 153, 255], 'radius': 3},
  300.             10: {'color': [255, 153, 153], 'radius': 3},
  301.             11: {'color': [255, 102, 102], 'radius': 3},
  302.             12: {'color': [255, 51, 51], 'radius': 3},
  303.             13: {'color': [153, 255, 153], 'radius': 3},
  304.             14: {'color': [102, 255, 102], 'radius': 3},
  305.             15: {'color': [51, 255, 51], 'radius': 3},
  306.             16: {'color': [0, 255, 0], 'radius': 3},
  307.         }
  308.  
  309.         # 点类别文字
  310.         # self.kpt_labelstr = {
  311.         #     'font_size': 1.5,  # 字体大小
  312.         #     'font_thickness': 3,  # 字体粗细
  313.         #     'offset_x': 10,  # X 方向，文字偏移距离，向右为正
  314.         #     'offset_y': 0,  # Y 方向，文字偏移距离，向下为正
  315.         # }
  316.  
  317.         # 骨架连接 BGR 配色
  318.         self.skeleton_map = [
  319.             {'srt_kpt_id': 0, 'dst_kpt_id': 1, 'color': [196, 75, 255], 'thickness': 2},  # thickness 线宽
  320.             {'srt_kpt_id': 0, 'dst_kpt_id': 2, 'color': [180, 187, 28], 'thickness': 2},  #
  321.             {'srt_kpt_id': 1, 'dst_kpt_id': 3, 'color': [47, 255, 173], 'thickness': 2},  #
  322.             {'srt_kpt_id': 2, 'dst_kpt_id': 4, 'color': [47, 255, 173], 'thickness': 2},
  323.             {'srt_kpt_id': 3, 'dst_kpt_id': 5, 'color': [47, 255, 173], 'thickness': 2},
  324.             {'srt_kpt_id': 5, 'dst_kpt_id': 7, 'color': [47, 255, 173], 'thickness': 2},
  325.             {'srt_kpt_id': 7, 'dst_kpt_id': 9, 'color': [47, 255, 173], 'thickness': 2},
  326.             {'srt_kpt_id': 4, 'dst_kpt_id': 6, 'color': [47, 255, 173], 'thickness': 2},
  327.             {'srt_kpt_id': 6, 'dst_kpt_id': 8, 'color': [47, 255, 173], 'thickness': 2},
  328.             {'srt_kpt_id': 8, 'dst_kpt_id': 10, 'color': [47, 255, 173], 'thickness': 2},
  329.             {'srt_kpt_id': 5, 'dst_kpt_id': 6, 'color': [47, 255, 173], 'thickness': 2},
  330.             {'srt_kpt_id': 5, 'dst_kpt_id': 11, 'color': [47, 255, 173], 'thickness': 2},
  331.             {'srt_kpt_id': 11, 'dst_kpt_id': 13, 'color': [47, 255, 173], 'thickness': 2},
  332.             {'srt_kpt_id': 13, 'dst_kpt_id': 15, 'color': [47, 255, 173], 'thickness': 2},
  333.             {'srt_kpt_id': 6, 'dst_kpt_id': 12, 'color': [47, 255, 173], 'thickness': 2},
  334.             {'srt_kpt_id': 12, 'dst_kpt_id': 14, 'color': [47, 255, 173], 'thickness': 2},
  335.             {'srt_kpt_id': 14, 'dst_kpt_id': 16, 'color': [47, 255, 173], 'thickness': 2},
  336.         ]
  337.         self.generate()
  338.  
  339.     # ---------------------------------------------------#
  340.     #   生成模型
  341.     # ---------------------------------------------------#
  342.     def generate(self):
  343.         # ---------------------------------------------------#
  344.         #   建立yolo模型，载入yolo模型的权重
  345.         # ---------------------------------------------------#
  346.         engine = self.load_engine(self.model_path)
  347.         self.context = engine.create_execution_context()
  348.         self.inputs, self.outputs, self.bindings = [], [], []
  349.         self.stream = cuda.Stream()
  350.         for binding in engine:
  351.             size = engine.get_binding_shape(binding)
  352.             dtype = trt.nptype(engine.get_binding_dtype(binding))
  353.             host_mem = np.empty(size, dtype=dtype)
  354.             host_mem = np.ascontiguousarray(host_mem)
  355.             device_mem = cuda.mem_alloc(host_mem.nbytes)
  356.             self.bindings.append(int(device_mem))
  357.             if engine.binding_is_input(binding):
  358.                 self.inputs.append({'host': host_mem, 'device': device_mem})
  359.             else:
  360.                 self.outputs.append({'host': host_mem, 'device': device_mem})
  361.  
  362.     def load_engine(self, engine_path):
  363.         TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
  364.         with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
  365.             return runtime.deserialize_cuda_engine(f.read())
  366.  
  367.     def forward(self, img):
  368.         self.inputs[0]['host'] = np.ravel(img)
  369.         # transfer data to the gpu
  370.         for inp in self.inputs:
  371.             cuda.memcpy_htod_async(inp['device'], inp['host'], self.stream)
  372.         # run inference
  373.         self.context.execute_async_v2(
  374.             bindings=self.bindings,
  375.             stream_handle=self.stream.handle)
  376.         # fetch outputs from gpu
  377.         for out in self.outputs:
  378.             cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
  379.         # synchronize stream
  380.         self.stream.synchronize()
  381.         return [out['host'] for out in self.outputs]
  382.  
  383.     # ---------------------------------------------------#
  384.     #   检测图片
  385.     # ---------------------------------------------------#
  386.     def detect_image(self, image):
  387.         # ---------------------------------------------------------#
  388.         #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
  389.         #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
  390.         # ---------------------------------------------------------#
  391.         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  392.         image_data, scale = resize_image(image, (self.input_shape[1], self.input_shape[0]), False)
  393.         # ---------------------------------------------------------#
  394.         #   添加上batch_size维度
  395.         #   h, w, 3 => 3, h, w => 1, 3, h, w
  396.         # ---------------------------------------------------------#
  397.         image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
  398.         # ---------------------------------------------------------#
  399.         #   将图像输入网络当中进行预测！
  400.         # ---------------------------------------------------------#
  401.         outputs = self.forward(image_data)[::-1]
  402.         # ---------------------------------------------------------#
  403.         #   将预测框进行堆叠，然后进行非极大抑制
  404.         # ---------------------------------------------------------#
  405.         results = non_max_suppression(outputs, conf_thres=self.confidence, iou_thres=self.nms_iou, nc=1)[0]
  406.         if results is None:
  407.             return image
  408.  
  409.         top_label = np.array(results[:, 5], dtype='int32')
  410.         top_conf = results[:, 4]
  411.         top_boxes = results[:, :4]
  412.         top_kpts = results[:, 6:].reshape(len(results), self.kpts_shape[0], self.kpts_shape[1])
  413.  
  414.         # ---------------------------------------------------------#
  415.         #   图像绘制
  416.         # ---------------------------------------------------------#
  417.         for i, c in list(enumerate(top_label)):
  418.             predicted_class = self.class_names[int(c)]
  419.             box = top_boxes[i]
  420.             score = top_conf[i]
  421.             left, top, right, bottom = box.astype('int32')
  422.             left = int(left * scale[0])
  423.             top = int(top * scale[1])
  424.             right = int(right * scale[0])
  425.             bottom = int(bottom * scale[1])
  426.             image = cv2.rectangle(image, (left, top), (right, bottom), self.bbox_color, self.bbox_thickness)
  427.             label = '{} {:.2f}'.format(predicted_class, score)
  428.             # 写框类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
  429.             image = cv2.putText(image, label,
  430.                                 (left + self.bbox_labelstr['offset_x'], top + self.bbox_labelstr['offset_y']),
  431.                                 cv2.FONT_HERSHEY_SIMPLEX, self.bbox_labelstr['font_size'], self.bbox_color,
  432.                                 self.bbox_labelstr['font_thickness'])
  433.  
  434.             bbox_keypoints = top_kpts[i]  # 该框所有关键点坐标和置信度
  435.  
  436.             # 画该框的骨架连接
  437.             for skeleton in self.skeleton_map:
  438.                 # 获取起始点坐标
  439.                 srt_kpt_id = skeleton['srt_kpt_id']
  440.                 srt_kpt_x = int(bbox_keypoints[srt_kpt_id][0] * scale[0])
  441.                 srt_kpt_y = int(bbox_keypoints[srt_kpt_id][1] * scale[1])
  442.                 # 获取终止点坐标
  443.                 dst_kpt_id = skeleton['dst_kpt_id']
  444.                 dst_kpt_x = int(bbox_keypoints[dst_kpt_id][0] * scale[0])
  445.                 dst_kpt_y = int(bbox_keypoints[dst_kpt_id][1] * scale[1])
  446.                 # 获取骨架连接颜色
  447.                 skeleton_color = skeleton['color']
  448.                 # 获取骨架连接线宽
  449.                 skeleton_thickness = skeleton['thickness']
  450.                 # 画骨架连接
  451.                 image = cv2.line(image, (srt_kpt_x, srt_kpt_y), (dst_kpt_x, dst_kpt_y), color=skeleton_color,
  452.                                  thickness=skeleton_thickness)
  453.  
  454.             # 画该框的关键点
  455.             for kpt_id in self.kpt_color_map:
  456.                 # 获取该关键点的颜色、半径、XY坐标
  457.                 kpt_color = self.kpt_color_map[kpt_id]['color']
  458.                 kpt_radius = self.kpt_color_map[kpt_id]['radius']
  459.                 kpt_x = int(bbox_keypoints[kpt_id][0] * scale[0])
  460.                 kpt_y = int(bbox_keypoints[kpt_id][1] * scale[1])
  461.                 # 画圆：图片、XY坐标、半径、颜色、线宽（-1为填充）
  462.                 image = cv2.circle(image, (kpt_x, kpt_y), kpt_radius, kpt_color, -1)
  463.                 # 写关键点类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
  464.                 # kpt_label = str(self.kpt_color_map[kpt_id]['name'])
  465.                 # image = cv2.putText(image)
  466.  
  467.         image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
  468.         return image
  469.  
  470.  
  471. if __name__ == '__main__':
  472.     yolo = YOLO()
  473.     # ----------------------------------------------------------------------------------------------------------#
  474.     #   mode用于指定测试的模式：
  475.     #   'predict'           表示单张图片预测，如果想对预测过程进行修改，如保存图片，截取对象等，可以先看下方详细的注释
  476.     #   'video'             表示视频检测，可调用摄像头或者视频进行检测，详情查看下方注释。
  477.     #   'fps'               表示测试fps，使用的图片是img里面的street.jpg，详情查看下方注释。
  478.     #   'dir_predict'       表示遍历文件夹进行检测并保存。默认遍历img文件夹，保存img_out文件夹，详情查看下方注释。
  479.     # ----------------------------------------------------------------------------------------------------------#
  480.     mode = "video"
  481.     # ----------------------------------------------------------------------------------------------------------#
  482.     #   video_path          用于指定视频的路径，当video_path=0时表示检测摄像头
  483.     #                       想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
  484.     #   video_save_path     表示视频保存的路径，当video_save_path=""时表示不保存
  485.     #                       想要保存视频，则设置如video_save_path = "yyy.mp4"即可，代表保存为根目录下的yyy.mp4文件。
  486.     #   video_fps           用于保存的视频的fps
  487.     #
  488.     #   video_path、video_save_path和video_fps仅在mode='video'时有效
  489.     #   保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。
  490.     # ----------------------------------------------------------------------------------------------------------#
  491.     video_path = 'two.mp4'
  492.     video_save_path = "one_out.mp4"
  493.     video_fps = 25.0
  494.     # ----------------------------------------------------------------------------------------------------------#
  495.     #   test_interval       用于指定测量fps的时候，图片检测的次数。理论上test_interval越大，fps越准确。
  496.     #   fps_image_path      用于指定测试的fps图片
  497.     #
  498.     #   test_interval和fps_image_path仅在mode='fps'有效
  499.     # ----------------------------------------------------------------------------------------------------------#
  500.     test_interval = 100
  501.     fps_image_path = "img/test.jpg"
  502.     # -------------------------------------------------------------------------#
  503.     #   dir_origin_path     指定了用于检测的图片的文件夹路径
  504.     #   dir_save_path       指定了检测完图片的保存路径
  505.     #
  506.     #   dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
  507.     # -------------------------------------------------------------------------#
  508.     dir_origin_path = "img/"
  509.     dir_save_path = "img_out/"
  510.  
  511.     if mode == "predict":
  512.         '''
  513.         1、如果想要进行检测完的图片的保存，利用r_image.save("img.jpg")即可保存，直接在predict.py里进行修改即可。 
  514.         2、如果想要获得预测框的坐标，可以进入yolo.detect_image函数，在绘图部分读取top，left，bottom，right这四个值。
  515.         3、如果想要利用预测框截取下目标，可以进入yolo.detect_image函数，在绘图部分利用获取到的top，left，bottom，right这四个值
  516.         在原图上利用矩阵的方式进行截取。
  517.         4、如果想要在预测图上写额外的字，比如检测到的特定目标的数量，可以进入yolo.detect_image函数，在绘图部分对predicted_class进行判断，
  518.         比如判断if predicted_class == 'car': 即可判断当前目标是否为车，然后记录数量即可。利用draw.text即可写字。
  519.         '''
  520.         while True:
  521.             img = input('Input image filename:')
  522.             try:
  523.                 image = cv2.imread(img)
  524.             except:
  525.                 print('Open Error! Try again!')
  526.                 continue
  527.             else:
  528.                 r_image = yolo.detect_image(image)
  529.                 cv2.imshow('result', r_image)
  530.                 c = cv2.waitKey(0)
  531.  
  532.     elif mode == "video":
  533.         capture = cv2.VideoCapture(video_path)
  534.         if video_save_path != "":
  535.             fourcc = cv2.VideoWriter_fourcc(*'XVID')
  536.             size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
  537.             out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
  538.  
  539.         ref, frame = capture.read()
  540.         if not ref:
  541.             raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
  542.  
  543.         fps = 0.0
  544.         while (True):
  545.             t1 = time.time()
  546.             # 读取某一帧
  547.             ref, frame = capture.read()
  548.             if not ref:
  549.                 break
  550.             # 进行检测
  551.             frame = yolo.detect_image(frame)
  552.  
  553.             fps = (fps + (1. / (time.time() - t1))) / 2
  554.             print("fps= %.2f" % (fps))
  555.             frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
  556.  
  557.             cv2.imshow("video", frame)
  558.             c = cv2.waitKey(1) & 0xff
  559.             if video_save_path != "":
  560.                 out.write(frame)
  561.  
  562.             if c == 27:
  563.                 capture.release()
  564.                 break
  565.  
  566.         print("Video Detection Done!")
  567.         capture.release()
  568.         if video_save_path != "":
  569.             print("Save processed video to the path :" + video_save_path)
  570.             out.release()
  571.         cv2.destroyAllWindows()
  572.  
  573.     elif mode == "fps":
  574.         img = cv2.imread(fps_image_path)
  575.         tact_time = yolo.get_FPS(img, test_interval)
  576.         print(str(tact_time) + ' seconds, ' + str(1 / tact_time) + 'FPS, @batch_size 1')
  577.  
  578.     elif mode == "dir_predict":
  579.         import os
  580.  
  581.         from tqdm import tqdm
  582.  
  583.         img_names = os.listdir(dir_origin_path)
  584.         for img_name in tqdm(img_names):
  585.             if img_name.lower().endswith(
  586.                     ('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
  587.                 image_path = os.path.join(dir_origin_path, img_name)
  588.                 image = cv2.imread(image_path)
  589.                 r_image = yolo.detect_image(image)
  590.                 if not os.path.exists(dir_save_path):
  591.                     os.makedirs(dir_save_path)
  592.                 r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
  593.  
  594.     else:
  595.         raise AssertionError("Please specify the correct mode: 'predict', 'video', 'fps', 'dir_predict'.")
  596.  
  597.

三、遇到的问题

（一）模型转换遇到的问题

在使用命令转engine模型时，会出现如下图的问题：

这个不影响，只是警告，等待几分钟就转换完成了

（二）代码使用注意
1、模型修改地方

2、输入参数修改

3、关键点的颜色和大小修改

4、骨架连接顺序和线宽修改

5、代码使用，看注释

6、项目目录结构

参考文献：

睿智的目标检测——YOLOv8-Pose的TensorRT推理__白鹭先生_的博客-CSDN博客

YOLOV8模型训练+部署（实战）-CSDN博客

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

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

  

闽ICP备14008679号