YOLOv8 Pose使用RKNN进行推理

关注微信公众号:朱sir的小站,发送202411081即可免费获取源代码下载链接

一、简单介绍

        YOLOv8-Pose 是一种基于 YOLOv8 架构的姿态估计模型,能够识别图像中的关键点位置,这些关键点通常表示人体的关节、特征点或其他显著位置。该模型在 COCO 关键点数据集上训练,适合多种姿势估计任务。

二、ONNX推理

1.首先需要先将Pytorch模型转换为Onnx模型,下载pt模型

这里给出官方的权重下载地址:点击此处,选择yolov8m-pose.pt进行下载​​​​​​​

2.进行ONNX转换

首先安装ultralytics

pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

再进行转换

from ultralytics import YOLO

# Load a model
model = YOLO("models/yolov8m-pose.pt")  # load an official model

# Export the model
model.export(format="onnx")

3.进行ONNX推理

在整个推理中分为两个文件,把它们放在同一目录底下:

onnx_inference.py-主要执行文件

import numpy as np
import cv2
import matplotlib.cm as cm
import onnxruntime as ort

import yolo_nms

session = ort.InferenceSession('models/yolov8m-pose.onnx',providers=['CPUExecutionProvider'])

input_name = session.get_inputs()[0].name


def model_inference(input=None):
    output = session.run([], {input_name: input})
    return output[0]


IMG_SZ = (736, 480)

sk = [15, 13, 13, 11, 16, 14, 14, 12, 11, 12,
      5, 11, 6, 12, 5, 6, 5, 7, 6, 8, 7, 9, 8, 10,
      1, 2, 0, 1, 0, 2, 1, 3, 2, 4, 3, 5, 4, 6]


def preprocess_img(frame):
    frame = cv2.resize(frame, (640, 640))
    img = frame[:, :, ::-1]
    img = img / 255.00
    img = np.asarray(img, dtype=np.float32)
    img = np.expand_dims(img, 0)
    img = img.transpose(0, 3, 1, 2)
    return img


def single_non_max_suppression(prediction):
    argmax = np.argmax(prediction[4, :])
    x = (prediction.T)[argmax]

    box = x[:4]  # Cx,Cy,w,h
    conf = x[4]
    keypts = x[5:]

    return box, conf, keypts


def post_process_multi(img, output, score_threshold=10):
    boxes, conf_scores, keypt_vectors = yolo_nms.non_max_suppression(output, score_threshold)

    for keypts, conf in zip(keypt_vectors, conf_scores):
        plot_keypoints(img, keypts, score_threshold)
    return img


def post_process_single(img, output, score_threshold=10):
    box, conf, keypts = single_non_max_suppression(output)
    keypts = smooth_pred(keypts)
    plot_keypoints(img, keypts, score_threshold)
    return img


def plot_keypoints(img, keypoints, threshold=10):
    for i in range(0, len(sk) // 2):
        pos1 = (int(keypoints[3 * sk[2 * i]]), int(keypoints[3 * sk[2 * i] + 1]))
        pos2 = (int(keypoints[3 * sk[2 * i + 1]]), int(keypoints[3 * sk[2 * i + 1] + 1]))
        conf1 = keypoints[3 * sk[2 * i] + 2]
        conf2 = keypoints[3 * sk[2 * i + 1] + 2]

        color = (cm.jet(i / (len(sk) // 2))[:3])
        color = [int(c * 255) for c in color[::-1]]
        if conf1 > threshold and conf2 > threshold:  # For a limb, both the keypoint confidence must be greater than 0.5
            cv2.line(img, pos1, pos2, color, thickness=8)

    for i in range(0, len(keypoints) // 3):
        x = int(keypoints[3 * i])
        y = int(keypoints[3 * i + 1])
        conf = keypoints[3 * i + 2]
        if conf > threshold:  # Only draw the circle if confidence is above some threshold
            cv2.circle(img, (x, y), 3, (0, 0, 0), -1)


keypoints_old = None


def smooth_pred(keypoints):
    global keypoints_old
    if keypoints_old is None:
        keypoints_old = keypoints.copy()
        return keypoints

    smoothed_keypoints = []
    for i in range(0, len(keypoints), 3):
        x_keypoint = keypoints[i]
        y_keypoint = keypoints[i + 1]
        conf = keypoints[i + 2]
        x_keypoint_old = keypoints_old[i]
        y_keypoint_old = keypoints_old[i + 1]
        conf_old = keypoints_old[i + 2]
        x_smoothed = (conf * x_keypoint + conf_old * x_keypoint_old) / (conf + conf_old)
        y_smoothed = (conf * y_keypoint + conf_old * y_keypoint_old) / (conf + conf_old)
        smoothed_keypoints.extend([x_smoothed, y_smoothed, (conf + conf_old) / 2])
    keypoints_old = smoothed_keypoints
    return smoothed_keypoints


if __name__ == "__main__":
    image_path = 'image/test.png'
    frame = cv2.imread(image_path)
    original_height, original_width = frame.shape[:2]
    input_img = preprocess_img(frame)
    output = model_inference(input_img)
    frame = cv2.resize(frame, (640, 640))
    print("Prediction shape:", output[0].shape)
    frame = post_process_multi(frame, output[0], score_threshold=0.2)
    frame_final = cv2.resize(frame, (original_width, original_height))
    cv2.imwrite('result.png',frame_final)

yolo_nms.py-后处理文件

import numpy as np


def nms(boxes, scores, threshold=0.5):
    # Sort the bounding boxes by their confidence scores in descending order
    indices = np.argsort(scores)[::-1]
    boxes = boxes[indices]

    # Initialize a list of selected bounding box indices
    selected_indices = []

    # Loop over the sorted list of bounding boxes
    while len(boxes) > 0:
        # Select the box or keypoint with the highest confidence score
        selected_index = indices[0]
        selected_indices.append(selected_index)

        # Compute the overlap between the selected box and all other boxes
        ious = compute_iou(boxes[0], boxes[1:])

        # Remove all boxes that overlap with the selected box by more than the threshold
        indices = indices[1:][ious <= threshold]
        boxes = boxes[1:][ious <= threshold]

    # Convert the list of selected box indices to an array and return it
    return np.array(selected_indices)


def compute_iou(box, boxes):
    """Compute the intersection-over-union (IoU) between a bounding box and a set of other bounding boxes."""
    x1 = np.maximum(box[0], boxes[:, 0])
    y1 = np.maximum(box[1], boxes[:, 1])
    x2 = np.minimum(box[2], boxes[:, 2])
    y2 = np.minimum(box[3], boxes[:, 3])
    intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0)
    area1 = (box[2] - box[0]) * (box[3] - box[1])
    area2 = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    union = area1 + area2 - intersection
    return intersection / union


# slower, outputs all boxes
def non_max_suppression(prediction, conf_thres=0.05):
    confidences = prediction[4, :]
    conf_mask = confidences > conf_thres
    x = (prediction.T)[conf_mask]
    if x.shape[0] == 0:
        return [], [], []

    box, conf, mask = np.split(x, (4, 5), axis=1)
    boxxyxy = np.copy(box)
    boxxyxy[..., 0] = box[..., 0] - box[..., 2] / 2  # top left x
    boxxyxy[..., 1] = box[..., 1] - box[..., 3] / 2  # top left y
    boxxyxy[..., 2] = box[..., 0] + box[..., 2] / 2  # bottom right x
    boxxyxy[..., 3] = box[..., 1] + box[..., 3] / 2  # bottom right y
    box = boxxyxy

    x = np.concatenate((box, conf, mask), axis=1)

    x = x[x[:, 4].argsort()[::-1]]
    boxes, scores = x[:, :4], x[:, 4]

    indices = nms(boxes, scores)
    x = x[indices]

    return x[:, :4], x[:, 4], x[:, 5:]

安装完所需的包之后,选择完需推理图片和模型位置之后,运行onnx_inference.py即可进行推理。

三、RKNN推理

        需要安装RKNNtoolkit2,根据不同的Ubuntu版本和python版本进行选择,Ubuntu20.04建议python3.8,Ubuntu22.04建议python3.10。注意,这里也需要上面的yolo_nms.py。

rknn_inference.py

import cv2
import numpy as np
from rknn.api import *

import yolo_nms
from onnx_inference import plot_keypoints

ONNX_MODEL='yolov8m-pose.onnx'
RKNN_MODEL='yolov8m-pose.rknn'
IMG_PATH='test.png'

def preprocess_img(frame):
    frame = cv2.resize(frame, (640, 640))
    img = frame[:, :, ::-1]
    img = img / 255.00
    img = np.asarray(img, dtype=np.float32)
    img = np.expand_dims(img, 0)
    img = img.transpose(0, 3, 1, 2)
    return img

image_input = cv2.imread(IMG_PATH)
original_height, original_width = image_input.shape[:2]
image = preprocess_img(image_input)

rknn = RKNN(verbose=True)
rknn.config(mean_values=[0,0,0], std_values=[1,1,1], target_platform='rk3568')
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
    print('Load model failed!')
print('done')
ret = rknn.build(do_quantization=False)
if ret != 0:
    print('Load model failed!')
print('done')
ret = rknn.init_runtime()
# ret = rknn.init_runtime('rk3566')
if ret != 0:
    print('Init runtime environment failed!')
# print('done')
print('--> Export rknn model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    print('Export rknn model failed!')
    # exit(ret)
print('done')
outputs_rknn = rknn.inference(inputs=[image], data_format=['nchw'])

result = outputs_rknn[0]
def post_process_multi(img, output, score_threshold=10):
    boxes, conf_scores, keypt_vectors = yolo_nms.non_max_suppression(output, score_threshold)

    for keypts, conf in zip(keypt_vectors, conf_scores):
        plot_keypoints(img, keypts, score_threshold)
    return img
frame = cv2.resize(image_input, (640, 640))
frame = post_process_multi(frame,result[0],score_threshold=0.2)
frame_final = cv2.resize(frame, (original_width, original_height))
cv2.imwrite('rknn_result.png',frame_final)

        安装完成依赖后,直接进行运行即可。

关注微信公众号:朱sir的小站,发送202411081即可免费获取源代码下载链接

你可能感兴趣的:(实用项目部署,YOLO,人工智能,python,linux,pip)