Spaces:
Sleeping
Sleeping
File size: 3,078 Bytes
c66f90e |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
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
|