Spaces:
Sleeping
Sleeping
| import numpy as np | |
| from ultralytics.engine.results import Results | |
| KEYPOINT_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"] | |
| def get_keypoints(result: Results): | |
| keypoints = None | |
| for i, box in enumerate(result.boxes): | |
| if box.cls != 0.: # Only consider the person class | |
| continue | |
| person_conf = box.conf.item() | |
| k = result.keypoints.data[i] | |
| x = k[:, 0].tolist() | |
| y = k[:, 1].tolist() | |
| score = k[:, 2] | |
| visible = (score > 0.5).tolist() | |
| # keypoints = {'x': x, 'y': y, 'visible': visible} | |
| keypoints = {key_name: (x_, y_) if v_ else None for key_name,x_,y_,v_ in zip(KEYPOINT_NAMES, x, y, visible)} | |
| break | |
| return keypoints | |
| def calculate_angle(p1, p2, p3): | |
| v1 = np.array([p1[0] - p2[0], p1[1] - p2[1]]) | |
| v2 = np.array([p3[0] - p2[0], p3[1] - p2[1]]) | |
| angle_rad = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))) | |
| angle_deg = np.degrees(angle_rad) | |
| return angle_deg | |
| def calculate_angle_to_horizontal(vector): | |
| angle_rad = np.arctan2(vector[1], vector[0]) | |
| angle_deg = np.degrees(angle_rad) | |
| # Adjust the angle to be within -90 to +90 degrees | |
| if angle_deg > 90: | |
| angle_deg = 180 - angle_deg | |
| elif angle_deg < -90: | |
| angle_deg = -180 - angle_deg | |
| return -angle_deg | |
| def extend_line(start, end, extend_factor=3): | |
| vector = np.array(end) - np.array(start) | |
| length = np.linalg.norm(vector) | |
| unit_vector = vector / np.linalg.norm(vector) | |
| new_point = end + unit_vector * length * extend_factor | |
| new_point = new_point.tolist() | |
| return (new_point[0], new_point[1]) | |
| def get_elbow_angles(keypoints: dict): | |
| left_elbow_angle = None | |
| right_elbow_angle = None | |
| if keypoints['left_shoulder'] and keypoints['left_elbow'] and keypoints['left_wrist']: | |
| left_elbow_angle = calculate_angle(keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']) | |
| if keypoints['right_shoulder'] and keypoints['right_elbow'] and keypoints['right_wrist']: | |
| right_elbow_angle = calculate_angle(keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']) | |
| return left_elbow_angle, right_elbow_angle | |
| def get_eye_angles(keypoints: dict): | |
| left_eye_angle = None | |
| right_eye_angle = None | |
| if keypoints['left_ear'] and keypoints['left_eye']: | |
| left_vector = (keypoints['left_eye'][0] - keypoints['left_ear'][0], keypoints['left_eye'][1] - keypoints['left_ear'][1]) | |
| left_eye_angle = calculate_angle_to_horizontal(left_vector) | |
| if keypoints['right_ear'] and keypoints['right_eye']: | |
| right_vector = (keypoints['right_eye'][0] - keypoints['right_ear'][0], keypoints['right_eye'][1] - keypoints['right_ear'][1]) | |
| right_eye_angle = calculate_angle_to_horizontal(right_vector) | |
| return left_eye_angle, right_eye_angle | |