| 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 |