目录效果更好rtmpose:效果更好rtmpose:import os import cv2 import argparse import numpy as np from rtmlib import Wholebody def process_video(video_path, output_pathNone): 处理视频进行人体姿态检测 if not os.path.exists(video_path): raise FileNotFoundError(f视频路径无效: {video_path}) # 初始化模型 wholebody Wholebody(modeperformance, backendonnxruntime, devicecuda) # 打开视频 cap cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f无法打开视频: {video_path}) # 获取视频信息 fps int(cap.get(cv2.CAP_PROP_FPS)) width int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 设置输出视频 out None if output_path: fourcc cv2.VideoWriter_fourcc(*mp4v) out cv2.VideoWriter(output_path, fourcc, fps, (width, height)) frame_count 0 print(f开始处理视频: {video_path}) while True: ret, frame cap.read() if not ret: break frame_count 1 frame_rgbcv2.cvtColor(frame,cv2.COLOR_BGR2RGB) # 推理获取关键点 keypoints, scores wholebody(frame_rgb) person_detected False # 如果检测到人体 if keypoints is not None and len(keypoints) 0: for kpts in keypoints: # 过滤有效关键点 valid_kpts [pt for pt in kpts if pt[0] 0 and pt[1] 0] if len(valid_kpts) 5: # 至少5个关键点才认为是有效人体 continue pts np.array(valid_kpts) x_min, y_min np.min(pts, axis0) x_max, y_max np.max(pts, axis0) area (x_max - x_min) * (y_max - y_min) if area 900: continue # 计算边界框扩展20% valid_kpts [pt for pt in kpts if pt[0] 0 and pt[1] 0] pts np.array(valid_kpts) x_min, y_min np.min(pts, axis0) x_max, y_max np.max(pts, axis0) # 扩展框 w, h x_max - x_min, y_max - y_min expand_w, expand_h int(w * 0.01), int(h * 0.01) x1 max(0, int(x_min - expand_w)) y1 max(0, int(y_min - expand_h)) x2 min(width, int(x_max expand_w)) y2 min(height, int(y_max expand_h)) # 绘制边界框 cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(frame, ren, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) # 绘制关键点 for pt in kpts: if pt[0] 0 and pt[1] 0: cv2.circle(frame, (int(pt[0]), int(pt[1])), 2, (0, 255, 255), -1) # 显示状态信息 status Person Detected if person_detected else No Person color (0, 255, 0) if person_detected else (0, 0, 255) cv2.putText(frame, f{status} | Frame: {frame_count}, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) # 保存或显示 # if out: # out.write(frame) cv2.imshow(Pose Detection, frame) if cv2.waitKey(1) 0xFF ord(q): break # 释放资源 cap.release() if out: out.release() cv2.destroyAllWindows() print(f处理完成共 {frame_count} 帧) if output_path: print(f视频已保存: {output_path}) if __name__ __main__: parser argparse.ArgumentParser(description视频人体姿态检测) parser.add_argument(--video_path, typestr,defaultrC:\Users\Administrator\Videos\yumao\yumao.mp4, help视频路径) parser.add_argument(--output_path, typestr, defaultNone, help输出视频路径可选) args parser.parse_args() process_video(args.video_path, args.output_path)pip install ultralyticsfrom ultralytics import YOLO import cv2 import numpy as np import time from pathlib import Path import json class PoseVideoProcessor: def __init__(self, model_path, conf_threshold0.5): self.model YOLO(model_path) self.conf_threshold conf_threshold self.keypoints_names [nose, left_eye, right_eye, left_ear, right_ear, left_shoulder, right_shoulder, left_elbow, right_elbow, left_wrist, right_wrist, left_hip, right_hip, left_knee, right_knee, left_ankle, right_ankle] self.skeleton [[5, 7], [7, 9], # 左臂 [6, 8], [8, 10], # 右臂 [5, 6], [5, 11], [6, 12], # 躯干 [11, 12], [11, 13], [13, 15], # 左腿 [12, 14], [14, 16] # 右腿 ] def draw_pose(self, frame, keypoints, scores, show_confidenceTrue): 在帧上绘制姿态关键点和骨骼 h, w frame.shape[:2] img_copy frame.copy() # 绘制骨骼 for pair in self.skeleton: p1 keypoints[pair[0]] p2 keypoints[pair[1]] s1 scores[pair[0]] s2 scores[pair[1]] if s1 self.conf_threshold and s2 self.conf_threshold: cv2.line(img_copy, (int(p1[0]), int(p1[1])), (int(p2[0]), int(p2[1])), (0, 255, 0), 2) # 绘制关键点 for i, (kp, score) in enumerate(zip(keypoints, scores)): if score self.conf_threshold: x, y int(kp[0]), int(kp[1]) # 绘制圆点 cv2.circle(img_copy, (x, y), 5, (0, 0, 255), -1) # 显示置信度 if show_confidence: cv2.putText(img_copy, f{score:.2f}, (x 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1) return img_copy def process_video_one(self, input_path, output_pathNone, show_previewTrue, save_keypointsFalse, keypoints_fileNone): cap cv2.VideoCapture(str(input_path)) if not cap.isOpened(): print(f无法打开视频: {input_path}) return # 获取视频属性 fps int(cap.get(cv2.CAP_PROP_FPS)) width int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_frames int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 设置输出视频 if output_path is None: input_path_obj Path(input_path) output_path input_path_obj.parent / f{input_path_obj.stem}_pose_advanced.mp4 fourcc cv2.VideoWriter_fourcc(*mp4v) out cv2.VideoWriter(str(output_path), fourcc, fps, (width, height)) # 准备保存关键点数据 all_keypoints_data [] print(f开始处理视频: {input_path}) print(f总帧数: {total_frames}) frame_count 0 start_time time.time() processing_times [] while True: ret, frame cap.read() if not ret: break # 推理 frame_start time.time() results self.model.predict(frame, confself.conf_threshold, verboseFalse) frame_time time.time() - frame_start processing_times.append(frame_time) # 获取关键点 keypoints_data None annotated_frame frame.copy() if len(results) 0 and results[0].keypoints is not None: keypoints results[0].keypoints.data.cpu().numpy() if len(keypoints) 0: keypoints keypoints[0] # 取第一个检测对象 kps keypoints[:, :2] # x, y坐标 scores keypoints[:, 2] # 置信度 # 绘制姿态 annotated_frame self.draw_pose(frame, kps, scores, show_confidenceFalse) # 保存关键点数据 if save_keypoints: frame_data {frame: frame_count, keypoints: kps.tolist(), scores: scores.tolist()} all_keypoints_data.append(frame_data) # 添加帧信息 cv2.putText(annotated_frame, fFrame: {frame_count}/{total_frames}, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) cv2.putText(annotated_frame, fFPS: {1 / frame_time:.1f}, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # 写入输出视频 out.write(annotated_frame) # 显示预览 if show_preview: # 缩放显示以适应屏幕 display_frame cv2.resize(annotated_frame, (width // 2, height // 2)) cv2.imshow(Pose Estimation, display_frame) key cv2.waitKey(1) 0xFF if key ord(q) or key 27: # q或ESC退出 break elif key ord( ): # 空格暂停 cv2.waitKey(0) frame_count 1 if frame_count % 30 0: avg_time np.mean(processing_times[-30:]) print(f已处理: {frame_count}/{total_frames} 帧, 平均推理时间: {avg_time * 1000:.1f}ms) # 释放资源 cap.release() out.release() cv2.destroyAllWindows() # 保存关键点数据 if save_keypoints and keypoints_file: with open(keypoints_file, w) as f: json.dump(all_keypoints_data, f, indent2) print(f关键点数据保存至: {keypoints_file}) # 统计信息 total_time time.time() - start_time avg_time np.mean(processing_times) if processing_times else 0 avg_fps 1 / avg_time if avg_time 0 else 0 print(\n 处理完成 ) print(f总处理时间: {total_time:.2f}s) print(f平均推理时间: {avg_time * 1000:.1f}ms) print(f平均FPS: {avg_fps:.1f}) print(f结果保存至: {output_path}) return str(output_path) def process_video(self, input_path, output_pathNone, show_previewTrue, save_keypointsFalse, keypoints_fileNone, max_peopleNone): cap cv2.VideoCapture(str(input_path)) if not cap.isOpened(): print(f无法打开视频: {input_path}) return # 获取视频属性 fps int(cap.get(cv2.CAP_PROP_FPS)) width int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) total_frames int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # 设置输出视频 if output_path is None: input_path_obj Path(input_path) output_path input_path_obj.parent / f{input_path_obj.stem}_pose_advanced.mp4 fourcc cv2.VideoWriter_fourcc(*mp4v) out cv2.VideoWriter(str(output_path), fourcc, fps, (width, height)) # 准备保存关键点数据 all_keypoints_data [] print(f开始处理视频: {input_path}) print(f总帧数: {total_frames}) frame_count 0 start_time time.time() processing_times [] while True: ret, frame cap.read() if not ret: break frameframe[150:-50, 230:-150, :] framecv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 推理 frame_start time.time() results self.model.predict(frame, confself.conf_threshold, verboseFalse) frame_time time.time() - frame_start processing_times.append(frame_time) # 获取关键点 annotated_frame frame.copy() frame_people_data [] if len(results) 0 and results[0].keypoints is not None: keypoints_all results[0].keypoints.data.cpu().numpy() # 限制检测人数 if max_people is not None and len(keypoints_all) max_people: keypoints_all keypoints_all[:max_people] print(f帧 {frame_count}: 检测到 {len(keypoints_all)} 个人) # 调试信息 # 处理每个人 for person_idx, keypoints in enumerate(keypoints_all): kps keypoints[:, :2] # x, y坐标 scores keypoints[:, 2] # 置信度 # 为每个人使用不同的颜色 color self.get_person_color(person_idx) # 绘制姿态 annotated_frame self.draw_pose_with_color(annotated_frame, kps, scores, show_confidenceFalse, colorcolor) # 保存关键点数据 if save_keypoints: person_data {person_id: person_idx, keypoints: kps.tolist(), scores: scores.tolist()} frame_people_data.append(person_data) # 保存帧数据 if save_keypoints and frame_people_data: all_keypoints_data.append({frame: frame_count, people: frame_people_data}) # 添加帧信息 cv2.putText(annotated_frame, fFrame: {frame_count}/{total_frames}, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 2) # cv2.putText(annotated_frame, # fPeople: {len(keypoints_all) if len(results) 0 and results[0].keypoints is not None else 0}, # (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # 写入输出视频 # out.write(annotated_frame) # 显示预览 if show_preview: # 缩放显示以适应屏幕 display_frame cv2.resize(annotated_frame, (width // 1, height // 1)) cv2.imshow(Pose Estimation, display_frame) key cv2.waitKey(1) 0xFF if key ord(q) or key 27: # q或ESC退出 break elif key ord( ): # 空格暂停 cv2.waitKey(0) frame_count 1 if frame_count % 30 0: avg_time np.mean(processing_times[-30:]) print(f已处理: {frame_count}/{total_frames} 帧, 平均推理时间: {avg_time * 1000:.1f}ms) # 释放资源 cap.release() out.release() cv2.destroyAllWindows() # 保存关键点数据 if save_keypoints and keypoints_file: with open(keypoints_file, w) as f: json.dump(all_keypoints_data, f, indent2) print(f关键点数据保存至: {keypoints_file}) # 统计信息 total_time time.time() - start_time avg_time np.mean(processing_times) if processing_times else 0 avg_fps 1 / avg_time if avg_time 0 else 0 print(\n 处理完成 ) print(f总处理时间: {total_time:.2f}s) print(f平均推理时间: {avg_time * 1000:.1f}ms) print(f平均FPS: {avg_fps:.1f}) print(f结果保存至: {output_path}) return str(output_path) def get_person_color(self, person_idx): 为不同的人分配不同的颜色 colors [(0, 255, 0), # 绿色 (255, 0, 0), # 蓝色 (0, 0, 255), # 红色 (255, 255, 0), # 青色 (255, 0, 255), # 品红 (0, 255, 255), # 黄色 (128, 0, 128), # 紫色 (0, 128, 128), # 深青色 ] return colors[person_idx % len(colors)] def draw_pose_with_color(self, frame, keypoints, scores, show_confidenceTrue, color(0, 255, 0)): 在帧上绘制姿态关键点和骨骼使用指定颜色 img_copy frame.copy() # 绘制骨骼 for pair in self.skeleton: p1 keypoints[pair[0]] p2 keypoints[pair[1]] s1 scores[pair[0]] s2 scores[pair[1]] if s1 self.conf_threshold and s2 self.conf_threshold: cv2.line(img_copy, (int(p1[0]), int(p1[1])), (int(p2[0]), int(p2[1])), color, 2) # 绘制关键点 for i, (kp, score) in enumerate(zip(keypoints, scores)): if score self.conf_threshold: x, y int(kp[0]), int(kp[1]) # 绘制圆点 cv2.circle(img_copy, (x, y), 2, color, -1) # 显示置信度 if show_confidence: cv2.putText(img_copy, f{score:.2f}, (x 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1) return img_copy # 使用示例 if __name__ __main__: # 初始化处理器 processor PoseVideoProcessor(model_pathyolo26x-pose.pt, conf_threshold0.05) # 处理视频 processor.process_video(input_pathrC:\Users\Administrator\Videos\yumao\yumao.mp4, output_pathoutput_video.mp4, show_previewTrue, save_keypointsTrue, keypoints_filekeypoints_data.json)