Spaces:
Sleeping
Sleeping
File size: 2,721 Bytes
aa80475 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
from ultralytics import YOLO
import numpy as np
id_joints_dict = {0: 'nose',
1: 'left_eye',
2: 'right_eye',
3: 'left_ear',
4: 'right_ear',
5: 'left_shoulder',
6: 'right_shoulder',
7: 'left_elbow',
8: 'right_elbow',
9: 'left_wrist',
10: 'right_wrist',
11: 'left_hip',
12: 'right_hip',
13: 'left_knee',
14: 'right_knee',
15: 'left_ankle',
16: 'right_ankle'}
joints_id_dict = {v: k for k, v in id_joints_dict.items()}
model = YOLO('yolov8n-pose.pt')
def get_keypoints_from_keypoints(model, video_path):
keypoints = []
results = model(video_path, save=True, show_conf=False, show_boxes=False)
for frame in results:
keypoints.append(frame.keypoints.xy)
return keypoints
keypoints = get_keypoints_from_keypoints(model, '../../data/pose/squat.mp4')
def calculate_angle(a, b, c):
"""
Calculates the angle between three joints.
Args:
a (tuple): coordinates of the first joint
b (tuple): coordinates of the second joint
c (tuple): coordinates of the third joint
Returns:
angle (float): angle between the three joints
"""
ba = np.array(a) - np.array(b)
bc = np.array(c) - np.array(b)
cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
angle = np.arccos(cosine_angle)
return np.degrees(angle)
def compute_left_knee_angle(pose):
"""
Computes the knee angle.
Args:
pose (dict): pose dictionary
Returns:
knee_angle (float): knee angle
"""
left_hip = pose[0][joints_id_dict['left_hip']]
left_knee = pose[0][joints_id_dict['left_knee']]
left_ankle = pose[0][joints_id_dict['left_ankle']]
knee_angle = calculate_angle(left_hip, left_knee, left_ankle)
return knee_angle
def compute_right_knee_angle(pose):
"""
Computes the knee angle.
Args:
pose (dict): pose dictionary
Returns:
knee_angle (float): knee angle
"""
right_hip = pose[0][joints_id_dict['right_hip']]
right_knee = pose[0][joints_id_dict['right_knee']]
right_ankle = pose[0][joints_id_dict['right_ankle']]
knee_angle = calculate_angle(right_hip, right_knee, right_ankle)
return knee_angle
def moving_average(data, window_size):
"""
Computes the moving average of a list.
Args:
data (list): list of values
window_size (int): size of the window
Returns:
avg (list): list of moving average values
"""
avg = []
for i in range(len(data) - window_size + 1):
avg.append(sum(data[i:i + window_size]) / window_size)
return avg
|