devbox
从前慢现在也慢
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

用python调用YOLOV8预测视频并解析结果----错误解决_yolov8 python api

作者：从前慢现在也慢 | 2024-04-24 10:17:15

赞

踩

yolov8 python api

1 相关引用

1 同济子豪兄关键点检测教程视频

2 同济子豪兄的GitHub代码参考

3 提出问题的小伙伴的博客

2 问题描述

本节调用了YOLOV8的预训练模型来对视频进行预测，采用的是python的API，并将关键点检测的结果可视化。在未更改代码之前，跑出来的效果如图所示。如果检测到的点数少于16，会被自动映射到原点。

 

要注意在同济子豪兄的源码中，以下这句代码要加上.data才能正常运行，否则会发生报错。

results[0].keypoints.data.cpu().numpy().astype('uint32')

3 问题解决 

对代码进行了解析，想到了一种解决方法。首先，映射到了原点可能是因为原点也作为关键点被检测了出来并且进行了链接绘制，假如我们只取检测到的关键点进行绘制连线，应该可以解决这个问题。所以我们加上一句代码，获取检查的关键点的数量。后面画框的关键点的时候只去画检测到的关键点。

  1. # 该框所有关键点坐标和置信度
  2. bbox_keypoints = bboxes_keypoints[idx]  
  3.  
  4. # 检查关键点的数量
  5. num_keypoints = bbox_keypoints.shape[0]
  1.         # 画该框的关键点
  2.         for kpt_id in range(num_keypoints):
  3.             # 获取该关键点的颜色、半径、XY坐标
  4.             kpt_color = kpt_color_map[kpt_id]['color']
  5.             kpt_radius = kpt_color_map[kpt_id]['radius']
  6.             kpt_x = int(bbox_keypoints[kpt_id][0])
  7.             kpt_y = int(bbox_keypoints[kpt_id][1])
  8.  
  9.             # 画圆：图片、XY坐标、半径、颜色、线宽（-1为填充）
  10.             img_bgr = cv2.circle(img_bgr, (kpt_x, kpt_y), kpt_radius, kpt_color, -1)
  11.  
  12.             # 写关键点类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
  13.             kpt_label = str(kpt_id)  # 写关键点类别 ID（二选一）
  14.             # kpt_label = str(kpt_color_map[kpt_id]['name']) # 写关键点类别名称（二选一）
  15.             img_bgr = cv2.putText(img_bgr, kpt_label,
  16.                                   (kpt_x + kpt_labelstr['offset_x'], kpt_y + kpt_labelstr['offset_y']),
  17.                                   cv2.FONT_HERSHEY_SIMPLEX, kpt_labelstr['font_size'], kpt_color,
  18.                                   kpt_labelstr['font_thickness'])
  19.  
  20.     return img_bgr

但是发现原点的关键点还是被检测了出来，于是思考能否加上一个置信度的判断条件，将置信度低的检测到的关键点给筛出掉。这里选择置信度大于0.5的关键点进行链接，

bbox_keypoints[srt_kpt_id][2] > 0.5 and
bbox_keypoints[dst_kpt_id][2] > 0.5

  1.         # 画该框的骨架连接
  2.         for skeleton in skeleton_map:
  3.             # 获取起始点坐标
  4.             srt_kpt_id = skeleton['srt_kpt_id']
  5.  
  6.             # 获取终止点坐标
  7.             dst_kpt_id = skeleton['dst_kpt_id']
  8.             if (
  9.                     srt_kpt_id < num_keypoints and
  10.                     dst_kpt_id < num_keypoints and
  11.                     bbox_keypoints[srt_kpt_id][2] > 0.5 and
  12.                     bbox_keypoints[dst_kpt_id][2] > 0.5
  13.             ):
  14.                 # 获取起始点坐标
  15.                 srt_kpt_x = int(bbox_keypoints[srt_kpt_id][0])
  16.                 srt_kpt_y = int(bbox_keypoints[srt_kpt_id][1])
  17.  
  18.                 # 获取终止点坐标
  19.                 dst_kpt_x = int(bbox_keypoints[dst_kpt_id][0])
  20.                 dst_kpt_y = int(bbox_keypoints[dst_kpt_id][1])
  21.  
  22.                 # 获取骨架连接颜色
  23.                 skeleton_color = skeleton['color']
  24.  
  25.                 # 获取骨架连接线宽
  26.                 skeleton_thickness = skeleton['thickness']
  27.  
  28.                 # 画骨架连接
  29.                 img_bgr = cv2.line(img_bgr, (srt_kpt_x, srt_kpt_y), (dst_kpt_x, dst_kpt_y), color=skeleton_color,
  30.                                    thickness=skeleton_thickness)

但是我们按照这个代码运行发现原点的关键点还是会被连接起来。这里的原因在于在前面results[0].keypoints.data.cpu().numpy().astype('uint32')这一句代码，我们将他转成整型的时候，连着置信度也转成整型了，因为置信度都小于1，转过去就都变成了0。所以无论设置什么阈值，都还是会画到原点的连接。所以在这里，我们进行代码的修改，使得置信度这个参数不化成整型。 最后再将各个参数量拼成整体bboxes_keypoints。

  1. # 关键点的 xy 坐标
  2.     bboxes_keypoints_position = results[0].keypoints.data[:, :, :2].cpu().numpy().astype('uint32')
  3.     confidence = results[0].keypoints.data[:, :, 2].cpu().numpy()
  4.     bboxes_keypoints = np.concatenate([bboxes_keypoints_position, confidence[:, :, None]], axis=2)

 注意，进行以上修改之后，后续画关键点和画连接，获取参数信息的时候，最好都强制转换为整型，否则容易报错。

  1. # 获取起始点坐标
  2.                 srt_kpt_x = int(bbox_keypoints[srt_kpt_id][0])
  3.                 srt_kpt_y = int(bbox_keypoints[srt_kpt_id][1])
  4.  
  5.                 # 获取终止点坐标
  6.                 dst_kpt_x = int(bbox_keypoints[dst_kpt_id][0])
  7.                 dst_kpt_y = int(bbox_keypoints[dst_kpt_id][1])
  1. kpt_x = int(bbox_keypoints[kpt_id][0])
  2. kpt_y = int(bbox_keypoints[kpt_id][1])

 进行以上修改后，就不会出现映射到原点并连接的问题了，如果还是出现了问题，可以尝试去调整一下前面提到的置信度阈值，我的是0.5，可以按实际情况进行调整。本人修改之后跑出来的效果是这样的：

 4 完整代码

  1. import cv2
  2. import numpy as np
  3. import time
  4. from tqdm import tqdm
  5.  
  6. from ultralytics import YOLO
  7.  
  8. import matplotlib.pyplot as plt
  9.  
  10. import torch
  11. # 有 GPU 就用 GPU，没有就用 CPU
  12. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  13. print('device:', device)
  14.  
  15. # 载入预训练模型
  16. # model = YOLO('yolov8n-pose.pt')
  17. # model = YOLO('yolov8s-pose.pt')
  18. # model = YOLO('yolov8m-pose.pt')
  19. # model = YOLO('yolov8l-pose.pt')
  20. # model = YOLO('yolov8x-pose.pt')
  21. model = YOLO('yolov8x-pose-p6.pt')
  22.  
  23. # 切换计算设备
  24. model.to(device)
  25. # model.cpu()  # CPU
  26. # model.cuda() # GPU
  27.  
  28. # 框（rectangle）可视化配置
  29. bbox_color = (150, 0, 0)             # 框的 BGR 颜色
  30. bbox_thickness = 2                   # 框的线宽
  31.  
  32. # 框类别文字
  33. bbox_labelstr = {
  34.     'font_size':1,         # 字体大小
  35.     'font_thickness':2,   # 字体粗细
  36.     'offset_x':0,          # X 方向，文字偏移距离，向右为正
  37.     'offset_y':-10,       # Y 方向，文字偏移距离，向下为正
  38. }
  39.  
  40. # 关键点 BGR 配色
  41. kpt_color_map = {
  42.     0:{'name':'Nose', 'color':[0, 0, 255], 'radius':6},                # 鼻尖
  43.     1:{'name':'Right Eye', 'color':[255, 0, 0], 'radius':6},           # 右边眼睛
  44.     2:{'name':'Left Eye', 'color':[255, 0, 0], 'radius':6},            # 左边眼睛
  45.     3:{'name':'Right Ear', 'color':[0, 255, 0], 'radius':6},           # 右边耳朵
  46.     4:{'name':'Left Ear', 'color':[0, 255, 0], 'radius':6},            # 左边耳朵
  47.     5:{'name':'Right Shoulder', 'color':[193, 182, 255], 'radius':6},  # 右边肩膀
  48.     6:{'name':'Left Shoulder', 'color':[193, 182, 255], 'radius':6},   # 左边肩膀
  49.     7:{'name':'Right Elbow', 'color':[16, 144, 247], 'radius':6},      # 右侧胳膊肘
  50.     8:{'name':'Left Elbow', 'color':[16, 144, 247], 'radius':6},       # 左侧胳膊肘
  51.     9:{'name':'Right Wrist', 'color':[1, 240, 255], 'radius':6},       # 右侧手腕
  52.     10:{'name':'Left Wrist', 'color':[1, 240, 255], 'radius':6},       # 左侧手腕
  53.     11:{'name':'Right Hip', 'color':[140, 47, 240], 'radius':6},       # 右侧胯
  54.     12:{'name':'Left Hip', 'color':[140, 47, 240], 'radius':6},        # 左侧胯
  55.     13:{'name':'Right Knee', 'color':[223, 155, 60], 'radius':6},      # 右侧膝盖
  56.     14:{'name':'Left Knee', 'color':[223, 155, 60], 'radius':6},       # 左侧膝盖
  57.     15:{'name':'Right Ankle', 'color':[139, 0, 0], 'radius':6},        # 右侧脚踝
  58.     16:{'name':'Left Ankle', 'color':[139, 0, 0], 'radius':6},         # 左侧脚踝
  59. }
  60.  
  61. # 点类别文字
  62. kpt_labelstr = {
  63.     'font_size':0.5,             # 字体大小
  64.     'font_thickness':1,       # 字体粗细
  65.     'offset_x':10,             # X 方向，文字偏移距离，向右为正
  66.     'offset_y':0,            # Y 方向，文字偏移距离，向下为正
  67. }
  68.  
  69. # 骨架连接 BGR 配色
  70. skeleton_map = [
  71.     {'srt_kpt_id':15, 'dst_kpt_id':13, 'color':[0, 100, 255], 'thickness':2},       # 右侧脚踝-右侧膝盖
  72.     {'srt_kpt_id':13, 'dst_kpt_id':11, 'color':[0, 255, 0], 'thickness':2},         # 右侧膝盖-右侧胯
  73.     {'srt_kpt_id':16, 'dst_kpt_id':14, 'color':[255, 0, 0], 'thickness':2},         # 左侧脚踝-左侧膝盖
  74.     {'srt_kpt_id':14, 'dst_kpt_id':12, 'color':[0, 0, 255], 'thickness':2},         # 左侧膝盖-左侧胯
  75.     {'srt_kpt_id':11, 'dst_kpt_id':12, 'color':[122, 160, 255], 'thickness':2},     # 右侧胯-左侧胯
  76.     {'srt_kpt_id':5, 'dst_kpt_id':11, 'color':[139, 0, 139], 'thickness':2},        # 右边肩膀-右侧胯
  77.     {'srt_kpt_id':6, 'dst_kpt_id':12, 'color':[237, 149, 100], 'thickness':2},      # 左边肩膀-左侧胯
  78.     {'srt_kpt_id':5, 'dst_kpt_id':6, 'color':[152, 251, 152], 'thickness':2},       # 右边肩膀-左边肩膀
  79.     {'srt_kpt_id':5, 'dst_kpt_id':7, 'color':[148, 0, 69], 'thickness':2},          # 右边肩膀-右侧胳膊肘
  80.     {'srt_kpt_id':6, 'dst_kpt_id':8, 'color':[0, 75, 255], 'thickness':2},          # 左边肩膀-左侧胳膊肘
  81.     {'srt_kpt_id':7, 'dst_kpt_id':9, 'color':[56, 230, 25], 'thickness':2},         # 右侧胳膊肘-右侧手腕
  82.     {'srt_kpt_id':8, 'dst_kpt_id':10, 'color':[0,240, 240], 'thickness':2},         # 左侧胳膊肘-左侧手腕
  83.     {'srt_kpt_id':1, 'dst_kpt_id':2, 'color':[224,255, 255], 'thickness':2},        # 右边眼睛-左边眼睛
  84.     {'srt_kpt_id':0, 'dst_kpt_id':1, 'color':[47,255, 173], 'thickness':2},         # 鼻尖-左边眼睛
  85.     {'srt_kpt_id':0, 'dst_kpt_id':2, 'color':[203,192,255], 'thickness':2},         # 鼻尖-左边眼睛
  86.     {'srt_kpt_id':1, 'dst_kpt_id':3, 'color':[196, 75, 255], 'thickness':2},        # 右边眼睛-右边耳朵
  87.     {'srt_kpt_id':2, 'dst_kpt_id':4, 'color':[86, 0, 25], 'thickness':2},           # 左边眼睛-左边耳朵
  88.     {'srt_kpt_id':3, 'dst_kpt_id':5, 'color':[255,255, 0], 'thickness':2},          # 右边耳朵-右边肩膀
  89.     {'srt_kpt_id':4, 'dst_kpt_id':6, 'color':[255, 18, 200], 'thickness':2}         # 左边耳朵-左边肩膀
  90. ]
  91.  
  92.  
  93. def process_frame(img_bgr):
  94.     '''
  95.     输入摄像头画面 bgr-array，输出图像 bgr-array
  96.     '''
  97.  
  98.     results = model(img_bgr, verbose=False)  # verbose设置为False，不单独打印每一帧预测结果
  99.  
  100.     # 预测框的个数
  101.     num_bbox = len(results[0].boxes.cls)
  102.  
  103.     # 预测框的 xyxy 坐标
  104.     bboxes_xyxy = results[0].boxes.xyxy.cpu().numpy().astype('uint32')
  105.  
  106.     # 关键点的 xy 坐标
  107.     bboxes_keypoints_position = results[0].keypoints.data[:, :, :2].cpu().numpy().astype('uint32')
  108.     confidence = results[0].keypoints.data[:, :, 2].cpu().numpy()
  109.     bboxes_keypoints = np.concatenate([bboxes_keypoints_position, confidence[:, :, None]], axis=2)
  110.  
  111.     for idx in range(num_bbox):  # 遍历每个框
  112.  
  113.         # 获取该框坐标
  114.         bbox_xyxy = bboxes_xyxy[idx]
  115.  
  116.         # 获取框的预测类别（对于关键点检测，只有一个类别）
  117.         bbox_label = results[0].names[0]
  118.  
  119.         # 画框
  120.         img_bgr = cv2.rectangle(img_bgr, (bbox_xyxy[0], bbox_xyxy[1]), (bbox_xyxy[2], bbox_xyxy[3]), bbox_color,
  121.                                 bbox_thickness)
  122.  
  123.         # 写框类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
  124.         img_bgr = cv2.putText(img_bgr, bbox_label,
  125.                               (bbox_xyxy[0] + bbox_labelstr['offset_x'], bbox_xyxy[1] + bbox_labelstr['offset_y']),
  126.                               cv2.FONT_HERSHEY_SIMPLEX, bbox_labelstr['font_size'], bbox_color,
  127.                               bbox_labelstr['font_thickness'])
  128.  
  129.         bbox_keypoints = bboxes_keypoints[idx]  # 该框所有关键点坐标和置信度
  130.         # 检查关键点的数量
  131.         num_keypoints = bbox_keypoints.shape[0]
  132.  
  133.         # 画该框的骨架连接
  134.         for skeleton in skeleton_map:
  135.             # 获取起始点坐标
  136.             srt_kpt_id = skeleton['srt_kpt_id']
  137.  
  138.             # 获取终止点坐标
  139.             dst_kpt_id = skeleton['dst_kpt_id']
  140.             if (
  141.                     srt_kpt_id < num_keypoints and
  142.                     dst_kpt_id < num_keypoints and
  143.                     bbox_keypoints[srt_kpt_id][2] > 0.5 and
  144.                     bbox_keypoints[dst_kpt_id][2] > 0.5
  145.             ):
  146.                 # 获取起始点坐标
  147.                 srt_kpt_x = int(bbox_keypoints[srt_kpt_id][0])
  148.                 srt_kpt_y = int(bbox_keypoints[srt_kpt_id][1])
  149.  
  150.                 # 获取终止点坐标
  151.                 dst_kpt_x = int(bbox_keypoints[dst_kpt_id][0])
  152.                 dst_kpt_y = int(bbox_keypoints[dst_kpt_id][1])
  153.  
  154.                 # 获取骨架连接颜色
  155.                 skeleton_color = skeleton['color']
  156.  
  157.                 # 获取骨架连接线宽
  158.                 skeleton_thickness = skeleton['thickness']
  159.  
  160.                 # 画骨架连接
  161.                 img_bgr = cv2.line(img_bgr, (srt_kpt_x, srt_kpt_y), (dst_kpt_x, dst_kpt_y), color=skeleton_color,
  162.                                    thickness=skeleton_thickness)
  163.  
  164.             # 画该框的关键点
  165.         for kpt_id in range(num_keypoints):
  166.             # 获取该关键点的颜色、半径、XY坐标
  167.             kpt_color = kpt_color_map[kpt_id]['color']
  168.             kpt_radius = kpt_color_map[kpt_id]['radius']
  169.             kpt_x = int(bbox_keypoints[kpt_id][0])
  170.             kpt_y = int(bbox_keypoints[kpt_id][1])
  171.  
  172.             # 画圆：图片、XY坐标、半径、颜色、线宽（-1为填充）
  173.             img_bgr = cv2.circle(img_bgr, (kpt_x, kpt_y), kpt_radius, kpt_color, -1)
  174.  
  175.             # 写关键点类别文字：图片，文字字符串，文字左上角坐标，字体，字体大小，颜色，字体粗细
  176.             kpt_label = str(kpt_id)  # 写关键点类别 ID（二选一）
  177.             # kpt_label = str(kpt_color_map[kpt_id]['name']) # 写关键点类别名称（二选一）
  178.             img_bgr = cv2.putText(img_bgr, kpt_label,
  179.                                   (kpt_x + kpt_labelstr['offset_x'], kpt_y + kpt_labelstr['offset_y']),
  180.                                   cv2.FONT_HERSHEY_SIMPLEX, kpt_labelstr['font_size'], kpt_color,
  181.                                   kpt_labelstr['font_thickness'])
  182.  
  183.     return img_bgr
  184.  
  185.  
  186. # 视频逐帧处理代码模板
  187. # 不需修改任何代码，只需定义process_frame函数即可
  188. # 同济子豪兄 2021-7-10
  189.  
  190. def generate_video(input_path='videos/robot.mp4'):
  191.     filehead = input_path.split('/')[-1]
  192.     output_path = "out-" + filehead
  193.  
  194.     print('视频开始处理', input_path)
  195.  
  196.     # 获取视频总帧数
  197.     cap = cv2.VideoCapture(input_path)
  198.     frame_count = 0
  199.     while (cap.isOpened()):
  200.         success, frame = cap.read()
  201.         frame_count += 1
  202.         if not success:
  203.             break
  204.     cap.release()
  205.     print('视频总帧数为', frame_count)
  206.  
  207.     # cv2.namedWindow('Crack Detection and Measurement Video Processing')
  208.     cap = cv2.VideoCapture(input_path)
  209.     frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  210.  
  211.     # fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
  212.     # fourcc = cv2.VideoWriter_fourcc(*'XVID')
  213.     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
  214.     fps = cap.get(cv2.CAP_PROP_FPS)
  215.  
  216.     out = cv2.VideoWriter(output_path, fourcc, fps, (int(frame_size[0]), int(frame_size[1])))
  217.  
  218.     # 进度条绑定视频总帧数
  219.     with tqdm(total=frame_count - 1) as pbar:
  220.         try:
  221.             while (cap.isOpened()):
  222.                 success, frame = cap.read()
  223.                 if not success:
  224.                     break
  225.  
  226.                 # 处理帧
  227.                 # frame_path = './temp_frame.png'
  228.                 # cv2.imwrite(frame_path, frame)
  229.                 try:
  230.                     frame = process_frame(frame)
  231.                 except:
  232.                     print('error')
  233.                     pass
  234.  
  235.                 if success == True:
  236.                     # cv2.imshow('Video Processing', frame)
  237.                     out.write(frame)
  238.  
  239.                     # 进度条更新一帧
  240.                     pbar.update(1)
  241.  
  242.                 # if cv2.waitKey(1) & 0xFF == ord('q'):
  243.                 # break
  244.         except:
  245.             print('中途中断')
  246.             pass
  247.  
  248.     cv2.destroyAllWindows()
  249.     out.release()
  250.     cap.release()
  251.     print('视频已保存', output_path)
  252. generate_video(input_path='videos/two.mp4')
  253. #generate_video(input_path='videos/lym.mp4')
  254. #generate_video(input_path='videos/cxk.mp4')

5 感谢

感谢同济子豪兄的教程视频以及GitHub代码教学。子豪兄csdn主页

感谢提出问题的小伙伴。又又土主页 

第一次写博客，如果有错误的地方，欢迎批评和指正，一起讨论交流学习。

谢谢大家！ 

本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号