Spaces:
Sleeping
Sleeping
File size: 2,899 Bytes
aa80475 5730896 9a8177e aa80475 5730896 aa80475 5730896 aa80475 5730896 aa80475 5730896 aa80475 9a8177e |
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 |
from ultralytics import YOLO
import numpy as np
import os
import json
import matplotlib.pyplot as plt
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(video_path):
save_folder='tmp'
os.makedirs(save_folder, exist_ok=True)
keypoints = []
results = model(video_path, save=True, show_conf=False, show_boxes=False)
for (i, frame) in enumerate(results):
frame_dict = {}
frame_dict['frame'] = i
frame_dict['keypoints'] = frame.keypoints.xy[0].tolist()
keypoints.append(frame_dict)
file_path = os.path.join(save_folder, 'keypoints.json')
with open(file_path, 'w') as f:
json.dump(keypoints, f)
return file_path
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_right_knee_angle(keypoints):
"""
Computes the knee angle.
Args:
keypoints (list): list of keypoints
Returns:
knee_angle (float): knee angle
"""
right_hip = keypoints[joints_id_dict['right_hip']]
right_knee = keypoints[joints_id_dict['right_knee']]
right_ankle = keypoints[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
def save_knee_angle_fig(angles):
os.makedirs('fig', exist_ok=True)
plt.plot(angles)
plt.xlabel('Frame')
plt.ylabel('Knee Angle')
plt.title('Evolution of the knee angle')
plt.savefig('fig/knee_angle.png') |