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