目录一、关键点预测二、图像预处理二、推理三、后处理与可视化3.1、后处理3.2、特征点可视化四、完整pytorch代码yolov8-pose tensorrt一、关键点预测注本篇只是阐述推理流程tensorrt实现后续跟进。yolov8-pose的tensorrt部署代码稍后更新还是在仓库GitHub - FeiYull/TensorRT-Alpha: TensorRT-Alpha supports YOLOv8、YOLOv7、YOLOv6、YOLOv5、YOLOv4、v3、YOLOX、YOLOR...CUDA IS ALL YOU NEED.It also supports end2end CUDA C acceleration and multi-batch inference.也可以关注TensorRT系列教程-CSDN博客以下是官方预测代码from ultralytics import YOLO model YOLO(modelyolov8n-pose.pt) model.predict(sourced:/Data/1.jpg, saveTrue)推理过程无非是图像预处理 - 推理 - 后处理 可视化这三个关键步骤在文件大概247行D:\CodePython\ultralytics\ultralytics\engine\predictor.py代码如下# Preprocess with profilers[0]: im self.preprocess(im0s) # 图像预处理 # Inference with profilers[1]: preds self.inference(im, *args, **kwargs) # 推理 # Postprocess with profilers[2]: self.results self.postprocess(preds, im, im0s) # 后处理二、图像预处理通过debug进入上述self.preprocess函数看到代码实现如下。处理流程大概是padding满足矩形推理图像通道转换即BGR装RGB检查图像数据是否连续存储顺序有HWC转为CHW然后归一化。需要注意原始pytorch框架图像预处理的时候会将图像缩放padding为HxW的图像其中H、W为32倍数而导出tensorrt的时候为了高效推理H、W 固定为640x640。def preprocess(self, im): Prepares input image before inference. Args: im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list. not_tensor not isinstance(im, torch.Tensor) if not_tensor: im np.stack(self.pre_transform(im)) im im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) im np.ascontiguousarray(im) # contiguous im torch.from_numpy(im) img im.to(self.device) img img.half() if self.model.fp16 else img.float() # uint8 to fp16/32 if not_tensor: img / 255 # 0 - 255 to 0.0 - 1.0 return img二、推理图像预处理之后直接推理就行了这里是基于pytorch推理。def inference(self, im, *args, **kwargs): visualize increment_path(self.save_dir / Path(self.batch[0][0]).stem, mkdirTrue) if self.args.visualize and (not self.source_type.tensor) else False return self.model(im, augmentself.args.augment, visualizevisualize)三、后处理与可视化3.1、后处理网络推理输出特征图维度为56x8400其中8400表示候选目标数量56 xywhc points * 17points的长度为3分别为xyc即特征点的坐标和置信度尽管推理输出特征图中每一行既有bbox还有keypoints但是NMS的时候依然只作用于bbox下面代码作了NMS之后将筛选之后的目标中bbox、keypoints进行坐标值缩放缩放到原图尺寸坐标系。def postprocess(self, preds, img, orig_imgs): Return detection results for a given input image or list of images. preds ops.non_max_suppression(preds, self.args.conf, self.args.iou, agnosticself.args.agnostic_nms, max_detself.args.max_det, classesself.args.classes, nclen(self.model.names)) results [] for i, pred in enumerate(preds): orig_img orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs shape orig_img.shape pred[:, :4] ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round() pred_kpts pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:] pred_kpts ops.scale_coords(img.shape[2:], pred_kpts, shape) path self.batch[0] img_path path[i] if isinstance(path, list) else path results.append( Results(orig_imgorig_img, pathimg_path, namesself.model.names, boxespred[:, :6], keypointspred_kpts)) return results3.2、特征点可视化bbox可视化没什么好说的说下17个特征点的可视化在文件D:\CodePython\ultralytics_fire_smoke\ultralytics\utils\plotting.py171行绘制特征点需要注意需要按照预定义的顺序绘制其中特征点置信度需要足够大。def kpts(self, kpts, shape(640, 640), radius5, kpt_lineTrue): Plot keypoints on the image. Args: kpts (tensor): Predicted keypoints with shape [17, 3]. Each keypoint has (x, y, confidence). shape (tuple): Image shape as a tuple (h, w), where h is the height and w is the width. radius (int, optional): Radius of the drawn keypoints. Default is 5. kpt_line (bool, optional): If True, the function will draw lines connecting keypoints for human pose. Default is True. Note: kpt_lineTrue currently only supports human pose plotting. if self.pil: # Convert to numpy first self.im np.asarray(self.im).copy() nkpt, ndim kpts.shape is_pose nkpt 17 and ndim 3 kpt_line is_pose # kpt_lineTrue for now only supports human pose plotting # 绘制特征点 for i, k in enumerate(kpts): color_k [int(x) for x in self.kpt_color[i]] if is_pose else colors(i) x_coord, y_coord k[0], k[1] if x_coord % shape[1] ! 0 and y_coord % shape[0] ! 0: if len(k) 3: conf k[2] if conf 0.5: continue cv2.circle(self.im, (int(x_coord), int(y_coord)), radius, color_k, -1, lineTypecv2.LINE_AA) # 绘制线段 if kpt_line: ndim kpts.shape[-1] for i, sk in enumerate(self.skeleton): pos1 (int(kpts[(sk[0] - 1), 0]), int(kpts[(sk[0] - 1), 1])) pos2 (int(kpts[(sk[1] - 1), 0]), int(kpts[(sk[1] - 1), 1])) if ndim 3: conf1 kpts[(sk[0] - 1), 2] conf2 kpts[(sk[1] - 1), 2] if conf1 0.5 or conf2 0.5: continue if pos1[0] % shape[1] 0 or pos1[1] % shape[0] 0 or pos1[0] 0 or pos1[1] 0: continue if pos2[0] % shape[1] 0 or pos2[1] % shape[0] 0 or pos2[0] 0 or pos2[1] 0: continue cv2.line(self.im, pos1, pos2, [int(x) for x in self.limb_color[i]], thickness2, lineTypecv2.LINE_AA) if self.pil: # Convert im back to PIL and update draw self.fromarray(self.im)这里给一张特征点顺序图四、完整pytorch代码将以上流程合并起来并加以修改完整代码如下import torch import cv2 as cv import numpy as np from ultralytics.data.augment import LetterBox from ultralytics.utils import ops from ultralytics.engine.results import Results import copy # path d:/Data/1.jpg path d:/Data/6406402.jpg device cuda:0 conf 0.25 iou 0.7 # preprocess im cv.imread(path) # letterbox im [im] orig_imgs copy.deepcopy(im) im [LetterBox([640, 640], autoTrue, stride32)(imagex) for x in im] im im[0][None] # im np.stack(im) im im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) im np.ascontiguousarray(im) # contiguous im torch.from_numpy(im) img im.to(device) img img.float() img / 255 # load model pt ckpt torch.load(yolov8n-pose.pt, map_locationcpu) model ckpt[model].to(device).float() # FP32 model model.eval() # inference preds model(img) prediction ops.non_max_suppression(preds, conf, iou, agnosticFalse, max_det300, classesNone, nclen(model.names)) results [] for i, pred in enumerate(prediction): orig_img orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs shape orig_img.shape pred[:, :4] ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round() pred_kpts pred[:, 6:].view(len(pred), *model.kpt_shape) if len(pred) else pred[:, 6:] pred_kpts ops.scale_coords(img.shape[2:], pred_kpts, shape) img_path path results.append( Results(orig_imgorig_img, pathimg_path, namesmodel.names, boxespred[:, :6], keypointspred_kpts)) # show plot_args {line_width: None,boxes: True,conf: True, labels: True} plot_args[im_gpu] img[0] result results[0] plotted_img result.plot(**plot_args) cv.imshow(plotted_img, plotted_img) cv.waitKey(0) cv.destroyAllWindows() print()