三天从YOLOV8关键点检测入门到实战（第三天）——用onnx部署yolov8

1引用

[1] 同济子豪兄的github项目
[2] 小破站关键点检测视频

onnx部署的优点就是快，比原来模型要快3倍。这里大佬的部署可谓是无敌，读完之后膜拜了。在此向大佬的开源精神致敬，这一系列代码让我少走路很多弯路。

2代码以及解析

2.1代码解析

导入一下常用的库，因为我的老师会在部署的时候，将torch这个比较大的框架去掉，但这边作者是保留的。实际项目中可以考虑去掉torch框架，能够减少软件的体量。

import cv2
import numpy as np
from PIL import Image

import onnxruntime

import torch
# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

import matplotlib.pyplot as plt
%matplotlib inline
1
2
3
4
5
6
7
8
9
10
11
12

这里需要指定关键点的shape，因为是3角板，所以是3个点，每个点有两个坐标和一个置信度信息，所以shape是[3, 3]

kpts_shape = [3, 3] # 关键点 shape
1

创建onnx的推理引擎

ort_session = onnxruntime.InferenceSession('checkpoint/Triangle_215_yolov8l_pretrain.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

1
2

获取输入层的信息，这是模型的本身信息。

获取模型的输出层信息。对于onnx部署，我们需要提前了解的就是输入输出层的信息。

用opencv读取图像

图像的预处理工作。从之前推理引擎中可以看到输入图片的类型必须是[1,3,640,640]。

求图像的缩放比例，通常按照缩放大的比例作为标准

对于图像需要进行归一化，然后变成tensor

将图像输入到推理引擎。之后将输出的转成tensor，因为yolov8处理结果的时候，预测结果的类型就是tensor的类型。yolov有很多成熟的解析结果的函数，比如非极大值抑制等等都是基于tensor的预测结果的。

# ONNX Runtime 推理预测
ort_output = ort_session.run(output_name, {input_name[0]: input_tensor.cpu().numpy()})[0]

# 转 Tensor
preds = torch.Tensor(ort_output)
1
2
3
4
5

这里使用了yolov8自带的非极大值抑制的函数。pred = preds[0]输出结果为torch.Size([3, 15])，表示了三个图片，每个图片上有15个参数。

from ultralytics.utils import ops
preds = ops.non_max_suppression(preds, conf_thres=0.25, iou_thres=0.7, nc=1)
pred = preds[0]
1
2
3

len(pred_det)表示长度，这里等于pred的行数。因为一行记录了一个预测框的信息。

pred_det = pred[:, :6].cpu().numpy()
num_bbox = len(pred_det)
print('预测出 {} 个框'.format(num_bbox))
1
2
3

这里可以看到pred的每个部分的参数代表的含义。第一次觉得这么具体的。

我们知道的输入图像尺寸要求的就是640*640，因此对于预测结果来说，预测框坐标值被缩放了，因此我们需要将xy重新映射会原图像上。此外画图需要保证坐标点为整数uint32。

# 目标检测框 XYXY 坐标
# 还原为缩放之前原图上的坐标
pred_det[:, 0] = pred_det[:, 0] * x_ratio
pred_det[:, 1] = pred_det[:, 1] * y_ratio
pred_det[:, 2] = pred_det[:, 2] * x_ratio
pred_det[:, 3] = pred_det[:, 3] * y_ratio
bboxes_xyxy = pred_det[:, :4].astype('uint32')
1
2
3
4
5
6
7

我们再来解析关键点的数据。关键点数据为pred的第6列到最后，一共3个框*每个框9个值=27个值。

我们将关键点坐标映射回原图

# 还原为缩放之前原图上的坐标
bboxes_keypoints[:,:,0] = bboxes_keypoints[:,:,0] * x_ratio
bboxes_keypoints[:,:,1] = bboxes_keypoints[:,:,1] * y_ratio
bboxes_keypoints = bboxes_keypoints.astype('uint32')
1
2
3
4

下面这部分和前两天的一个模版。

2.2opencv可视化模版

# 框（rectangle）可视化配置
bbox_color = (150, 0, 0)             # 框的 BGR 颜色
bbox_thickness = 6                   # 框的线宽

# 框类别文字
bbox_labelstr = {
    'font_size':4,         # 字体大小
    'font_thickness':10,   # 字体粗细
    'offset_x':0,          # X 方向，文字偏移距离，向右为正
    'offset_y':-80,        # Y 方向，文字偏移距离，向下为正
}

# 关键点 BGR 配色
kpt_color_map = {
    0:{'name':'angle_30', 'color':[255, 0, 0], 'radius':40},      # 30度角点
    1:{'name':'angle_60', 'color':[0, 255, 0], 'radius':40},      # 60度角点
    2:{'name':'angle_90', 'color':[0, 0, 255], 'radius':40},      # 90度角点
}

# 点类别文字
kpt_labelstr = {
    'font_size':4,             # 字体大小
    'font_thickness':10,       # 字体粗细
    'offset_x':30,             # X 方向，文字偏移距离，向右为正
    'offset_y':120,            # Y 方向，文字偏移距离，向下为正
}

# 骨架连接 BGR 配色
skeleton_map = [
    {'srt_kpt_id':0, 'dst_kpt_id':1, 'color':[196, 75, 255], 'thickness':3},        # 30度角点-60度角点
    {'srt_kpt_id':0, 'dst_kpt_id':2, 'color':[180, 187, 28], 'thickness':3},        # 30度角点-90度角点
    {'srt_kpt_id':1, 'dst_kpt_id':2, 'color':[47,255, 173], 'thickness':3},         # 60度角点-90度角点
]

for idx in range(num_bbox): # 遍历每个框
    
    # 获取该框坐标
    bbox_xyxy = bboxes_xyxy[idx] 
    
    # 获取框的预测类别（对于关键点检测，只有一个类别）
    bbox_label = 'sjb_rect'
    
    # 画框
    img_bgr = cv2.rectangle(img_bgr, (bbox_xyxy[0], bbox_xyxy[1]), (bbox_xyxy[2], bbox_xyxy[3]), bbox_color, bbox_thickness)
    
    # 写框类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
    img_bgr = cv2.putText(img_bgr, bbox_label, (bbox_xyxy[0]+bbox_labelstr['offset_x'], bbox_xyxy[1]+bbox_labelstr['offset_y']), cv2.FONT_HERSHEY_SIMPLEX, bbox_labelstr['font_size'], bbox_color, bbox_labelstr['font_thickness'])
    
    bbox_keypoints = bboxes_keypoints[idx] # 该框所有关键点坐标和置信度
    
    # 画该框的骨架连接
    for skeleton in skeleton_map:
        
        # 获取起始点坐标
        srt_kpt_id = skeleton['srt_kpt_id']
        srt_kpt_x = bbox_keypoints[srt_kpt_id][0]
        srt_kpt_y = bbox_keypoints[srt_kpt_id][1]
        
        # 获取终止点坐标
        dst_kpt_id = skeleton['dst_kpt_id']
        dst_kpt_x = bbox_keypoints[dst_kpt_id][0]
        dst_kpt_y = bbox_keypoints[dst_kpt_id][1]
        
        # 获取骨架连接颜色
        skeleton_color = skeleton['color']
        
        # 获取骨架连接线宽
        skeleton_thickness = skeleton['thickness']
        
        # 画骨架连接
        img_bgr = cv2.line(img_bgr, (srt_kpt_x, srt_kpt_y),(dst_kpt_x, dst_kpt_y),color=skeleton_color,thickness=skeleton_thickness)
        
    # 画该框的关键点
    for kpt_id in kpt_color_map:
        
        # 获取该关键点的颜色、半径、XY坐标
        kpt_color = kpt_color_map[kpt_id]['color']
        kpt_radius = kpt_color_map[kpt_id]['radius']
        kpt_x = bbox_keypoints[kpt_id][0]
        kpt_y = bbox_keypoints[kpt_id][1]
        
        # 画圆：图片、XY坐标、半径、颜色、线宽（-1为填充）
        img_bgr = cv2.circle(img_bgr, (kpt_x, kpt_y), kpt_radius, kpt_color, -1)
        
        # 写关键点类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
        # kpt_label = str(kpt_id) # 写关键点类别 ID
        kpt_label = str(kpt_color_map[kpt_id]['name']) # 写关键点类别名称
        img_bgr = cv2.putText(img_bgr, kpt_label, (kpt_x+kpt_labelstr['offset_x'], kpt_y+kpt_labelstr['offset_y']), cv2.FONT_HERSHEY_SIMPLEX, kpt_labelstr['font_size'], kpt_color, kpt_labelstr['font_thickness'])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

最后得出的图像是这样子的

plt.imshow(img_bgr[:,:,::-1])
plt.show()
1
2

3总结

其实我们可以看到应用onnx来预测确实是准确性有所降低，这里也是用的最小的模型yolov8n-pose的模型，可能用大的模型会更好。但重点是一些套路，比如onnx推理引擎的使用，opencv可视化的模版，这些套路对以后开发是很有帮助的。

4相关文章

[1]三天从YOLOV8关键点检测入门到实战（第一天）——初识YOLOV8
[2] 三天从YOLOV8关键点检测入门到实战（第二天）——用python调用YOLOV8预测图片并解析结果
[3]三天从YOLOV8关键点检测入门到实战（第二天）——用python调用YOLOV8预测视频并解析结果
[4]三天从YOLOV8关键点检测入门到实战（第三天）——用onnx部署yolov8