File size: 1,434 Bytes
cca6a52 cfbaa51 cca6a52 29f3df9 89a012c cca6a52 1ca9b28 cca6a52 89a012c cca6a52 89a012c cca6a52 89a012c cca6a52 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import numpy as np
from ultralytics import YOLO
from CPR_Module.Common.logging_config import cpr_logger
class PoseEstimator:
"""Human pose estimation using YOLO"""
def __init__(self, model_path="CPR_Module/Common/yolo11n-pose.pt", min_confidence=0.2):
self.model = YOLO(model_path)
self.min_confidence = min_confidence
def detect_poses(self, frame):
"""Detect human poses in a frame"""
try:
results = self.model(frame, verbose=False, conf=self.min_confidence, show=False)
if not results or len(results[0].keypoints.xy) == 0:
return None
return results[0]
except Exception as e:
cpr_logger.error(f"Pose detection error: {e}")
return None
def get_keypoints(self, results, person_idx=0):
"""Extract keypoints for a detected person"""
try:
if not results or len(results.keypoints.xy) <= person_idx:
return None
return results.keypoints.xy[person_idx].cpu().numpy()
except Exception as e:
cpr_logger.error(f"Keypoint extraction error: {e}")
return None
def draw_keypoints(self, frame, results):
"""Draw detected keypoints on frame"""
try:
return results.plot()
except Exception as e:
cpr_logger.error(f"Keypoint drawing error: {e}")
return frame |