目录
效果更好：
rtmpose:
效果更好：
rtmpose:
import os import cv2 import argparse import numpy as np from rtmlib import Wholebody def process_video(video_path, output_path=None): """处理视频，进行人体姿态检测""" if not os.path.exists(video_path): raise FileNotFoundError(f"视频路径无效: {video_path}") # 初始化模型 wholebody = Wholebody(mode='performance', backend='onnxruntime', device='cuda') # 打开视频 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_rgb=cv2.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, axis=0) x_max, y_max = np.max(pts, axis=0) 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, axis=0) x_max, y_max = np.max(pts, axis=0) # 扩展框 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", type=str,default=r"C:\Users\Administrator\Videos\yumao\yumao.mp4", help="视频路径") parser.add_argument("--output_path", type=str, default=None, help="输出视频路径（可选）") args = parser.parse_args() process_video(args.video_path, args.output_path)pip install ultralytics
from 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_threshold=0.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_confidence=True): """在帧上绘制姿态关键点和骨骼""" 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_path=None, show_preview=True, save_keypoints=False, keypoints_file=None): 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, conf=self.conf_threshold, verbose=False) 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_confidence=False) # 保存关键点数据 if save_keypoints: frame_data = {'frame': frame_count, 'keypoints': kps.tolist(), 'scores': scores.tolist()} all_keypoints_data.append(frame_data) # 添加帧信息 cv2.putText(annotated_frame, f"Frame: {frame_count}/{total_frames}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) cv2.putText(annotated_frame, f"FPS: {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, indent=2) 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_path=None, show_preview=True, save_keypoints=False, keypoints_file=None, max_people=None): 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=frame[150:-50, 230:-150, :] frame=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 推理 frame_start = time.time() results = self.model.predict(frame, conf=self.conf_threshold, verbose=False) 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_confidence=False, color=color) # 保存关键点数据 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, f"Frame: {frame_count}/{total_frames}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 2) # cv2.putText(annotated_frame, # f"People: {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, indent=2) 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_confidence=True, 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_path="yolo26x-pose.pt", conf_threshold=0.05) # 处理视频 processor.process_video(input_path=r"C:\Users\Administrator\Videos\yumao\yumao.mp4", output_path="output_video.mp4", show_preview=True, save_keypoints=True, keypoints_file="keypoints_data.json")
)
