基于YOLO V8的人体骨架提取

发布时间 2023-08-09 22:10:33作者: tlott

题外话:我记得没错的话,最多两年前,当时YOLO系列才出到V5,怎么现在都出到V8了啊。

1、前言

YOLOv8是 ultralytics 公司在 2023 年 1月 10 号开源的 YOLOv5 的下一个重大更新版本,目前支持图像分类、物体检测和实例分割任务,该公司同样也是YOLOV5的提出者,本文将面对YOLOV8-pose的本地部署进行一系列的讲解,并对本人在进行部署时出现的问题进行逐一排查与解决。

2、源码、权重文件的下载以及Yolov8的安装

YOLOv8提供了两种使用方法,第一种便是按部就班的在本地部署源码以及权重文件,另一种则是已封装好的、可直接使用命令行进行调用的方式。以下将对两种方法进行讲解。

2.1 源码及权重文件的下载

YOLOV8源代码地址:https://github.com/ultralytics/ultralytics
其权重文件同样包含在上述网址中,如下图所示,其包含检测、分类、骨架提取等多个任务的权重文件,本文重点聚焦于人体骨架提取任务上,即YOLOv8n-pose等权重文件。
图1 权重文件

2.2 已封装的YOLOv8的安装

首先打开anaconda prompt并进入虚拟环境,输入以下指令即可安装。

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

3、本地部署的两种方式

3.1 命令行部署

使用命令行对图片进行骨架提取:

yolo pose predict model=模型名称 source=图片地址 device=0
yolo pose predict model=yolov8n-pose.pt source=images/yolo.png device=0



同理使用命令行对视频以及调用摄像头进行骨架提取:

yolo pose predict model=模型名称 source=图片地址 device=0
yolo pose predict model=yolov8n-pose.pt source=videos/yolo.mp4 device=0
yolo pose predict model=yolov8n-pose.pt source=0 show

以上便是使用命令行进行骨架提取的方式,在进行图像分割、识别等功能同样可以使用命令行运行,极大的便利了开发,但是缺少了DIY的效果,如果需要进行自己的训练以及使用还是需要回归到源码上。

3.2 源码本地部署

首先下载YOLOv8源码及对应权重文件,并配置所需的环境。

pip install -r requirements.txt

其次编写相应代码进行使用,本文仅提供图片骨架提取的demo。

# 制作者:tlott
from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt
import torch

# 有 GPU 就用 GPU,没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:', device)
# 载入预训练模型
model = YOLO('yolov8n-pose.pt')
model.to(device)
img_path = r'C:\Users\86158\yolo.png'
results = model(img_path)
# 转成整数的 numpy array
bboxes_xyxy = results[0].boxes.xyxy.cpu().numpy().astype('uint32')
# print(results[0].keypoints)
bboxes_keypoints = results[0].keypoints.data.cpu().numpy()
num_bbox = len(results[0].boxes.cls)
img_bgr = cv2.imread(img_path)
plt.imshow(img_bgr[:,:,::-1])
plt.show()
# 框(rectangle)可视化配置
bbox_color = (150, 0, 0)             # 框的 BGR 颜色
bbox_thickness = 6                   # 框的线宽
# 框类别文字
bbox_labelstr = {
    'font_size':6,         # 字体大小
    'font_thickness':14,   # 字体粗细
    'offset_x':0,          # X 方向,文字偏移距离,向右为正
    'offset_y':-80,        # Y 方向,文字偏移距离,向下为正
}

# 关键点 BGR 配色
kpt_color_map = {
    0: {'name': 'Nose', 'color': [0, 0, 255], 'radius': 6},  # 鼻尖
    1: {'name': 'Right Eye', 'color': [255, 0, 0], 'radius': 6},  # 右边眼睛
    2: {'name': 'Left Eye', 'color': [255, 0, 0], 'radius': 6},  # 左边眼睛
    3: {'name': 'Right Ear', 'color': [0, 255, 0], 'radius': 6},  # 右边耳朵
    4: {'name': 'Left Ear', 'color': [0, 255, 0], 'radius': 6},  # 左边耳朵
    5: {'name': 'Right Shoulder', 'color': [193, 182, 255], 'radius': 6},  # 右边肩膀
    6: {'name': 'Left Shoulder', 'color': [193, 182, 255], 'radius': 6},  # 左边肩膀
    7: {'name': 'Right Elbow', 'color': [16, 144, 247], 'radius': 6},  # 右侧胳膊肘
    8: {'name': 'Left Elbow', 'color': [16, 144, 247], 'radius': 6},  # 左侧胳膊肘
    9: {'name': 'Right Wrist', 'color': [1, 240, 255], 'radius': 6},  # 右侧手腕
    10: {'name': 'Left Wrist', 'color': [1, 240, 255], 'radius': 6},  # 左侧手腕
    11: {'name': 'Right Hip', 'color': [140, 47, 240], 'radius': 6},  # 右侧胯
    12: {'name': 'Left Hip', 'color': [140, 47, 240], 'radius': 6},  # 左侧胯
    13: {'name': 'Right Knee', 'color': [223, 155, 60], 'radius': 6},  # 右侧膝盖
    14: {'name': 'Left Knee', 'color': [223, 155, 60], 'radius': 6},  # 左侧膝盖
    15: {'name': 'Right Ankle', 'color': [139, 0, 0], 'radius': 6},  # 右侧脚踝
    16: {'name': 'Left Ankle', 'color': [139, 0, 0], 'radius': 6},  # 左侧脚踝
}

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

# 骨架连接 BGR 配色
skeleton_map = [
    {'srt_kpt_id': 15, 'dst_kpt_id': 13, 'color': [0, 100, 255], 'thickness': 1},  # 右侧脚踝-右侧膝盖
    {'srt_kpt_id': 13, 'dst_kpt_id': 11, 'color': [0, 255, 0], 'thickness': 1},  # 右侧膝盖-右侧胯
    {'srt_kpt_id': 16, 'dst_kpt_id': 14, 'color': [255, 0, 0], 'thickness': 1},  # 左侧脚踝-左侧膝盖
    {'srt_kpt_id': 14, 'dst_kpt_id': 12, 'color': [0, 0, 255], 'thickness': 1},  # 左侧膝盖-左侧胯
    {'srt_kpt_id': 11, 'dst_kpt_id': 12, 'color': [122, 160, 255], 'thickness': 1},  # 右侧胯-左侧胯
    {'srt_kpt_id': 5, 'dst_kpt_id': 11, 'color': [139, 0, 139], 'thickness': 1},  # 右边肩膀-右侧胯
    {'srt_kpt_id': 6, 'dst_kpt_id': 12, 'color': [237, 149, 100], 'thickness': 1},  # 左边肩膀-左侧胯
    {'srt_kpt_id': 5, 'dst_kpt_id': 6, 'color': [152, 251, 152], 'thickness': 1},  # 右边肩膀-左边肩膀
    {'srt_kpt_id': 5, 'dst_kpt_id': 7, 'color': [148, 0, 69], 'thickness': 1},  # 右边肩膀-右侧胳膊肘
    {'srt_kpt_id': 6, 'dst_kpt_id': 8, 'color': [0, 75, 255], 'thickness': 1},  # 左边肩膀-左侧胳膊肘
    {'srt_kpt_id': 7, 'dst_kpt_id': 9, 'color': [56, 230, 25], 'thickness': 1},  # 右侧胳膊肘-右侧手腕
    {'srt_kpt_id': 8, 'dst_kpt_id': 10, 'color': [0, 240, 240], 'thickness': 1},  # 左侧胳膊肘-左侧手腕
    {'srt_kpt_id': 1, 'dst_kpt_id': 2, 'color': [224, 255, 255], 'thickness': 1},  # 右边眼睛-左边眼睛
    {'srt_kpt_id': 0, 'dst_kpt_id': 1, 'color': [47, 255, 173], 'thickness': 1},  # 鼻尖-左边眼睛
    {'srt_kpt_id': 0, 'dst_kpt_id': 2, 'color': [203, 192, 255], 'thickness': 1},  # 鼻尖-左边眼睛
    {'srt_kpt_id': 1, 'dst_kpt_id': 3, 'color': [196, 75, 255], 'thickness': 1},  # 右边眼睛-右边耳朵
    {'srt_kpt_id': 2, 'dst_kpt_id': 4, 'color': [86, 0, 25], 'thickness': 1},  # 左边眼睛-左边耳朵
    {'srt_kpt_id': 3, 'dst_kpt_id': 5, 'color': [255, 255, 0], 'thickness': 1},  # 右边耳朵-右边肩膀
    {'srt_kpt_id': 4, 'dst_kpt_id': 6, 'color': [255, 18, 200], 'thickness': 1}  # 左边耳朵-左边肩膀
]
# %%
for idx in range(num_bbox):  # 遍历每个框
    # 获取该框坐标
    bbox_xyxy = bboxes_xyxy[idx]
    # 获取框的预测类别(对于关键点检测,只有一个类别)
    bbox_label = results[0].names[0]
    # 画框
    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 = round(bbox_keypoints[srt_kpt_id][0])
        srt_kpt_y = round(bbox_keypoints[srt_kpt_id][1])
        srt_kpt_conf = bbox_keypoints[srt_kpt_id][2]  # 获取起始点置信度
        # print(srt_kpt_conf)
        # 获取终止点坐标
        dst_kpt_id = skeleton['dst_kpt_id']
        dst_kpt_x = round(bbox_keypoints[dst_kpt_id][0])
        dst_kpt_y = round(bbox_keypoints[dst_kpt_id][1])
        dst_kpt_conf = bbox_keypoints[dst_kpt_id][2]  # 获取终止点置信度
        # print(dst_kpt_conf)
        # 获取骨架连接颜色
        skeleton_color = skeleton['color']
        # 获取骨架连接线宽
        skeleton_thickness = skeleton['thickness']
        # 如果起始点和终止点的置信度都高于阈值,才画骨架连接
        if srt_kpt_conf > 0.5 and dst_kpt_conf > 0.5:
            # 画骨架连接
            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 = round(bbox_keypoints[kpt_id][0])
        kpt_y = round(bbox_keypoints[kpt_id][1])
        kpt_conf = bbox_keypoints[kpt_id][2]  # 获取该关键点置信度
        if kpt_conf > 0.5:
            # 画圆:图片、XY坐标、半径、颜色、线宽(-1为填充)
            img_bgr = cv2.circle(img_bgr, (kpt_x, kpt_y), kpt_radius, kpt_color, -1)
plt.imshow(img_bgr[:, :, ::-1])
plt.show()

运行结果如图:

在编写代码时出现了多种问题,解决过程可参考另一博客。

4、总结

可能是设备的问题,RTX1650在运行YOLOv8时处理速度有点过慢,视频掉帧严重,故在进行更进一步的筛选后,我将处理方式聚焦于轻量级的OpenPose模型。