YOLOv8-pose针对视频实时提取打印对应关节点序号及坐标

因为我在找如何提取YOLOv8-pose的关键点的时候,大多都是针对静态图像,视频直接套用不太行,因此就改进了一下,如下:

初步代码:

复制代码
import torch  # 导入PyTorch库
import cv2 as cv  # 导入OpenCV库并重命名为cv
import numpy as np  # 导入NumPy库并重命名为np
from ultralytics.data.augment import LetterBox  # 从ultralytics.data.augment中导入LetterBox类
from ultralytics.utils import ops  # 从ultralytics.utils中导入ops模块
from ultralytics.engine.results import Results  # 从ultralytics.engine.results中导入Results类
import copy  # 导入copy模块

# 视频路径
video_path = 'D:/cs/yolov8_2/ultralytics/ceshi1.mp4'  # 将此处路径改为你的视频文件路径
device = 'cuda:0'  # 设备类型,此处使用CUDA
conf = 0.25  # 置信度阈值
iou = 0.7  # IoU(交并比)阈值

# 加载视频
cap = cv.VideoCapture(video_path)

# 检查视频是否成功打开
if not cap.isOpened():
    print("Error: Could not open video.")  # 打印错误消息
    exit()  # 退出程序

# 加载模型
ckpt = torch.load('yolov8n-pose.pt', map_location='cpu')  # 加载模型参数
model = ckpt['model'].to(device).float()  # 将模型加载到指定设备(CPU或GPU)并转换为浮点数类型
model.eval()  # 将模型设置为评估模式

results = []  # 存储结果的列表

while True:
    ret, frame = cap.read()  # 读取视频帧

    # 如果没有读取到帧或者视频结束,则退出循环
    if not ret:
        break

    orig_img = frame  # 原始图像

    # 预处理
    im = [orig_img]  # 图像列表
    im = [LetterBox([640, 640], auto=True, stride=32)(image=x) for x in im]  # 对图像进行LetterBox缩放
    im = im[0][None]  # 转换为数组形式
    im = im[..., ::-1].transpose((0, 3, 1, 2))  # BGR转RGB,BHWC转BCHW
    im = np.ascontiguousarray(im)  # 转换为连续的内存布局
    im = torch.from_numpy(im)  # 将数组转换为PyTorch张量
    img = im.to(device)  # 将张量移动到指定设备
    img = img.float()  # 转换为浮点数类型
    img /= 255  # 归一化

    # 推理
    preds = model(img)  # 模型推理
    prediction = ops.non_max_suppression(preds, conf, iou, agnostic=False, max_det=300, classes=None, nc=len(model.names))  # 非最大抑制得到预测结果

    for i, pred in enumerate(prediction):
        shape = orig_img.shape  # 图像形状
        pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round()  # 缩放边界框坐标
        pred_kpts = pred[:, 6:].view(len(pred), *model.kpt_shape) if len(pred) else pred[:, 6:]  # 获取关键点坐标

        pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, shape)  # 缩放关键点坐标

        results.append(
            Results(orig_img=orig_img,  # 原始图像
                    path=video_path,  # 视频路径
                    names=model.names,  # 类别名称
                    boxes=pred[:, :6],  # 边界框
                    keypoints=pred_kpts))  # 关键点
            # 获取关键点坐标及其对应的序号
        for j, kpts in enumerate(pred_kpts):
            keypoints = kpts.cpu().numpy()  # 将关键点转换为NumPy数组
            # keypoints 包含了所有关键点的坐标,每一行是一个关键点的坐标
            # 根据模型的结构,关键点序号可能是从0到N-1,N是关键点的总数
            # 您可以在这里使用 keypoints 获取关键点的坐标和对应的序号
            for k, keypoint in enumerate(keypoints):
                x, y = keypoint[:2]  # 关键点的坐标
                keypoint_index = k  # 关键点的序号
                # 这里可以对每个关键点的坐标和序号执行您需要的操作
                print("关键点序号:", keypoint_index, "关键点坐标:", (x, y))

    # 显示帧
    plot_args = {'line_width': None, 'boxes': True, 'conf': True, 'labels': True}  # 绘图参数
    plot_args['im_gpu'] = img[0]  # 图像张量
    plotted_img = results[-1].plot(**plot_args)  # 显示处理后的最后一帧结果
    cv.imshow('plotted_img', plotted_img)  # 显示图像

    # 按 'q' 键退出
    if cv.waitKey(1) & 0xFF == ord('q'):
        break

# 释放视频捕获并关闭所有窗口
cap.release()
cv.destroyAllWindows()

效果:

相关推荐
骄傲的心别枯萎1 小时前
RV1126 NO.46:RV1126+OPENCV对视频流进行视频膨胀操作
人工智能·opencv·计算机视觉·音视频·rv1126
B站_计算机毕业设计之家4 小时前
深度学习:Yolo水果检测识别系统 深度学习算法 pyqt界面 训练集测试集 深度学习 数据库 大数据 (建议收藏)✅
数据库·人工智能·python·深度学习·算法·yolo·pyqt
B站_计算机毕业设计之家8 小时前
深度学习:python人脸表情识别系统 情绪识别系统 深度学习 神经网络CNN算法 ✅
python·深度学习·神经网络·算法·yolo·机器学习·cnn
王道长服务器 | 亚马逊云8 小时前
AWS + 发财CMS:高效采集站的新形态
服务器·网络·云计算·音视频·aws
EasyDSS9 小时前
无人机RTMP推流平台EasyDSS:构建新一代智能交通视频监控系统
音视频·无人机
城南皮卡丘9 小时前
【数据集+源码】基于yolov11+streamlit的玉米叶片病虫害检测系统
yolo·目标检测·计算机视觉·智慧农业·玉米病害检测
Vue102410 小时前
uniapp 使用renderjs 封装 video-player 视频播放器, html5视频播放器-解决视频层级、覆盖、播放卡顿
uni-app·音视频·html5
王哈哈^_^10 小时前
【完整源码+数据集】蓝莓数据集,yolo11蓝莓成熟度检测数据集 3023 张,蓝莓成熟度数据集,目标检测蓝莓识别算法系统实战教程
人工智能·算法·yolo·目标检测·计算机视觉·ai·视觉检测
盘古开天166610 小时前
通俗易懂:YOLO模型原理详解,从零开始理解目标检测
人工智能·yolo·目标检测
王哈哈^_^11 小时前
【完整源码+数据集】高空作业数据集,yolo高空作业检测数据集 2076 张,人员高空作业数据集,目标检测高空作业识别系统实战教程
人工智能·算法·yolo·目标检测·计算机视觉·目标跟踪·视觉检测