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')