使用预训练的 ONNX 格式的目标检测模型(基于 YOLOv8n-pose)姿态监测

12 篇文章 0 订阅

具体步骤如下:

  1. 加载图像:

    • 从指定路径读取一张图像(这里假设图像名为bus.jpg)。
    • 将图像从 BGR 颜色空间转换为 RGB 颜色空间。
  2. 图像预处理:

    • 计算图像的高度、宽度,并确定其中的最大值作为新图像的边长。
    • 创建一个全零的新图像,大小为最长边的正方形,将原始图像复制到新图像中,确保图像的边长是最长边的长度。
    • 将新图像调整为640x640的大小,并进行转置和像素值归一化处理,最后添加一个维度以满足模型输入要求。
  3. 模型推理:

    • 使用onnxruntime加载预训练的 ONNX 模型。
    • 将预处理后的图像输入模型进行推理,得到模型的输出结果。
  4. 结果筛选:

    • 对模型输出结果进行转置操作,然后根据置信度阈值筛选出置信度大于 0.8 的检测结果。
    • 提取这些结果中的边界框信息和置信度,并使用 OpenCV 的非极大值抑制算法去除重叠的边界框,得到最终的检测结果。
  5. 绘制结果:

    • 遍历最终的检测结果,对于每个检测结果,计算边界框的四个顶点坐标,并在原始图像上绘制矩形框。
    • 将检测结果中除边界框和置信度之外的部分按照每三个元素一组进行分割,得到关键点信息。对于每个关键点,根据缩放比例计算其在原始图像中的坐标,并在图像上绘制一个红色的小圆点。
  6. 显示图像:

    • 显示绘制了检测结果的图像。
    • 等待用户按下任意键退出程序,并关闭所有窗口。
import cv2
import numpy as np
import onnxruntime as ort
from ultralytics import YOLO

导入了所需的库,包括 OpenCV(cv2)用于图像处理,numpy用于数值计算,onnxruntime用于加载和运行 ONNX 模型,以及ultralyticsYOLO类用于对象检测。

# model = YOLO('yolov8n-pose.pt')
# model.export(format='onnx')
img_path = ''
frame = cv2.imread('bus.jpg')

使用ultralyticsYOLO模型加载一个名为yolov8n-pose.pt的预训练模型,并将其导出为 ONNX 格式。随后,使用cv2.imread读取一张名为bus.jpg的图像,并将其存储在frame变量中。

rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
h, w, c = rgb_frame.shape
max_slide = max(h, w)
bg_img = np.zeros((max_slide, max_slide, 3), dtype=np.float32)
bg_img[:h, :w] = rgb_frame

将读取的图像从 BGR 颜色空间转换为 RGB 颜色空间。然后计算图像的高度、宽度和通道数,并找到高度和宽度中的较大值作为max_slide。创建一个全零的图像,大小为max_slide x max_slide x 3,然后将原始图像复制到这个新图像中,确保新图像的边长是图像最长边的长度。

image = cv2.resize(bg_img, dsize=(640, 640))
image = np.transpose(image, (2, 0, 1)) / 255
image = np.expand_dims(image, 0)
scale = max_slide / 640

将处理后的图像调整为640x640的大小。接着,对图像进行转置操作并将像素值归一化到[0, 1]范围,然后使用np.expand_dims在第一个维度上添加一个维度,以便符合模型输入的要求。同时,计算图像缩放比例。

session = ort.InferenceSession('yolov8n-pose.onnx', providers=['CPUExecutionProvider'])
input_name = session._inputs_meta[0].name
session_out = session.run(None, {input_name: image})[0][0]

使用onnxruntime加载名为yolov8n-pose.onnx的模型,并指定使用 CPU 进行推理。获取模型输入的名称,并将预处理后的图像输入模型进行推理,得到输出结果。

result = np.transpose(session_out, (1, 0))
result = result[result[:, 4] > 0.8]
bboxes = result[:, 0:4]
confs = result[:, 4]

对模型输出结果进行转置操作,然后筛选出置信度大于 0.8 的检测结果。提取出这些结果中的边界框信息和置信度。

idx = cv2.dnn.NMSBoxes(bboxes, confs, score_threshold=0.5, nms_threshold=0.3)
res = result[idx]

使用 OpenCV 的非极大值抑制(NMS)算法,对检测结果进行筛选,去除重叠的边界框。得到最终的检测结果。

for re in res:
    cx, cy, w, h = re[:4]
    x1 = (cx - w // 2) * scale
    y1 = (cy - h // 2) * scale
    x2 = (cx + w // 2) * scale
    y2 = (cy + h // 2) * scale
    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
    conf = re[5]
    cv2.rectangle(frame, (x1, y1),( x2, y2), color=(0, 255, 1), thickness=3, lineType=cv2.LINE_AA)
    kpoints = np.split(re[5:], re[5:].shape[0] // 3)
    for kpoint in kpoints:
        x_, y_, visibility = kpoint
        x_ = int(x_ * scale)
        y_ = int(y_ * scale)
        cv2.circle(frame, center=(x_, y_), radius=2, color=(0, 0, 255), thickness=2)

遍历最终的检测结果,对于每个检测结果,计算边界框的四个顶点坐标,并在原始图像上绘制矩形框。然后,将检测结果中除边界框和置信度之外的部分按照每三个元素一组进行分割,得到关键点信息。对于每个关键点,根据缩放比例计算其在原始图像中的坐标,并在图像上绘制一个红色的小圆点。

cv2.imshow('win', frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

最后,显示处理后的图像,并等待用户按下任意键退出程序,关闭所有窗口。

完整代码如下:

import cv2
import numpy as np
import onnxruntime as ort

from ultralytics import YOLO

# model=YOLO('yolov8n-pose.pt')
# model.export(format='onnx')
img_path=''
frame=cv2.imread('bus.jpg')
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
#
h, w, c = rgb_frame.shape
max_slide = max(h, w)
bg_img = np.zeros((max_slide, max_slide, 3), dtype=np.float32)
bg_img[:h, :w] = rgb_frame
# 640*640*3
image = cv2.resize(bg_img, dsize=(640, 640))
image = np.transpose(image, (2, 0, 1)) / 255
image = np.expand_dims(image, 0)
scale=max_slide/640
session=ort.InferenceSession('yolov8n-pose.onnx',providers=['CPUExecutionProvider'])
input_name=session._inputs_meta[0].name
session_out = session.run(None, {input_name:image})[0][0]
result=np.transpose(session_out,(1,0))
result=result[result[:,4]>0.8]
bboxes=result[:,0:4]
confs=result[:,4]
idx = cv2.dnn.NMSBoxes(bboxes, confs, score_threshold=0.5, nms_threshold=0.3)
res=result[idx]
for re in res:
    cx, cy, w, h = re[:4]
    x1 = (cx - w // 2) * scale
    y1 = (cy - h // 2) * scale
    x2 = (cx + w // 2) * scale
    y2 = (cy + h // 2) * scale
    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
    conf=re[5]#置信度
    cv2.rectangle(frame,(x1,y1),(x2,y2),color=(0,255,1),thickness=3,lineType=cv2.LINE_AA)
    kpoints=np.split(re[5:],re[5:].shape[0]//3)
    for kpoint in kpoints:
        x_,y_,visibility=kpoint
        x_=int(x_*scale)
        y_=int(y_*scale)
        cv2.circle(frame, center=(x_, y_), radius=2, color=(0, 0, 255), thickness=2)
cv2.imshow('win',frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

 结果如下:

  • 10
    点赞
  • 16
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

西柚与蓝莓

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值