File size: 1,646 Bytes
1ca9b28 cfbaa51 1ca9b28 29f3df9 1ca9b28 cfbaa51 1ca9b28 cfbaa51 1ca9b28 |
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 41 42 43 44 45 |
from ultralytics import YOLO
from CPR_Module.Common.logging_config import cpr_logger
class PoseEstimator:
"""Human pose estimation using YOLO"""
def __init__(self, min_confidence, model_path="CPR_Module/Common/yolo11n-pose.pt"):
self.model = YOLO(model_path).to("cuda")
if next(self.model.model.parameters()).is_cuda:
cpr_logger.info("YOLO model loaded on CUDA (GPU).")
else:
cpr_logger.warning("YOLO model is not on CUDA. Check your setup.")
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, iou=0.2)
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 |