YOLOv8-pose推理代码

AIOT魔法师

已于 2024-03-22 16:15:51 修改

阅读量318

点赞数 1

文章标签： YOLO

于 2024-03-03 22:04:23 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_41967600/article/details/136439490

版权

import cv2
import numpy as np
from ultralytics import YOLO

class PoseDetector:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8],
                         [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
        self.pose_palette = np.array([
            [255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255], [153, 204, 255],
            [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255], [255, 153, 153], [255, 102, 102],
            [255, 51, 51], [153, 255, 153], [102, 255, 102], [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0],
            [255, 255, 255]
        ], dtype=np.uint8)
        self.kpt_color = self.pose_palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]
        self.limb_color = self.pose_palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]

    @staticmethod
    def