Spaces:
Sleeping
Sleeping
| 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('yolov8m-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') |