具体步骤如下:
-
加载图像:
- 从指定路径读取一张图像(这里假设图像名为
bus.jpg
)。 - 将图像从 BGR 颜色空间转换为 RGB 颜色空间。
- 从指定路径读取一张图像(这里假设图像名为
-
图像预处理:
- 计算图像的高度、宽度,并确定其中的最大值作为新图像的边长。
- 创建一个全零的新图像,大小为最长边的正方形,将原始图像复制到新图像中,确保图像的边长是最长边的长度。
- 将新图像调整为
640x640
的大小,并进行转置和像素值归一化处理,最后添加一个维度以满足模型输入要求。
-
模型推理:
- 使用
onnxruntime
加载预训练的 ONNX 模型。 - 将预处理后的图像输入模型进行推理,得到模型的输出结果。
- 使用
-
结果筛选:
- 对模型输出结果进行转置操作,然后根据置信度阈值筛选出置信度大于 0.8 的检测结果。
- 提取这些结果中的边界框信息和置信度,并使用 OpenCV 的非极大值抑制算法去除重叠的边界框,得到最终的检测结果。
-
绘制结果:
- 遍历最终的检测结果,对于每个检测结果,计算边界框的四个顶点坐标,并在原始图像上绘制矩形框。
- 将检测结果中除边界框和置信度之外的部分按照每三个元素一组进行分割,得到关键点信息。对于每个关键点,根据缩放比例计算其在原始图像中的坐标,并在图像上绘制一个红色的小圆点。
-
显示图像:
- 显示绘制了检测结果的图像。
- 等待用户按下任意键退出程序,并关闭所有窗口。
import cv2
import numpy as np
import onnxruntime as ort
from ultralytics import YOLO
导入了所需的库,包括 OpenCV(cv2
)用于图像处理,numpy
用于数值计算,onnxruntime
用于加载和运行 ONNX 模型,以及ultralytics
的YOLO
类用于对象检测。
# model = YOLO('yolov8n-pose.pt')
# model.export(format='onnx')
img_path = ''
frame = cv2.imread('bus.jpg')
使用ultralytics
的YOLO
模型加载一个名为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()
结果如下: