squat agent working

Modules/PoseEstimation/pose_agent.py

@@ -1,66 +1,86 @@
-from Modules.PoseEstimation.pose_estimator import calculate_angle, joints_id_dict, model
 from langchain.tools import tool
 from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.messages import HumanMessage
 from langchain_mistralai.chat_models import ChatMistralAI
+import os
+import sys
+import json
+sys.path.append(os.getcwd())
+from Modules.PoseEstimation.pose_estimator import model, compute_right_knee_angle, moving_average
 
 # If api_key is not passed, default behavior is to use the `MISTRAL_API_KEY` environment variable.
 llm = ChatMistralAI(model='mistral-large-latest', api_key="i5jSJkCFNGKfgIztloxTMjfckiFbYBj4")
 
 @tool
-def compute_right_knee_angle(pose: list) -> float:
-
+def get_keypoints_from_keypoints(video_path: str) -> str:
     """
-    Computes the knee angle.
+    Extracts keypoints from a video file.
 
     Args:
-        pose (list): list of keypoints
+        video_path (str): path to the video file
 
     Returns:
-        knee_angle (float): knee angle
+        file_path (str): path to the JSON file containing the keypoints
     """
 
-    right_hip = pose[joints_id_dict['right_hip']]
-    right_knee = pose[joints_id_dict['right_knee']]
-    right_ankle = pose[joints_id_dict['right_ankle']]
-
-    knee_angle = calculate_angle(right_hip, right_knee, right_ankle)
+    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 compute_right_knee_angle_list(json_path: str) -> list[float]:
+    """
+    Computes the knee angle from a list of keypoints.
 
-    print(knee_angle)
+    Args:
+        json_path (str): path to the JSON file containing the keypoints
 
-    return str(knee_angle)
+    Returns:
+        right_knee_angle_list (list[float]): list of knee angles
+    """
+    keypoints_list = json.load(open(json_path))
+    right_knee_angle_list = []
+    for keypoints in keypoints_list:
+        right_knee_angle = compute_right_knee_angle(keypoints['keypoints'])
+        right_knee_angle_list.append(right_knee_angle)
+    return moving_average(right_knee_angle_list, 10)
 
 @tool
-def get_keypoints_from_path(video_path: str):
+def check_knee_angle(json_path: str) -> bool:
     """
-    Get keypoints from a video.
+    Checks if the minimum knee angle is smaller than a threshold.
+    If the minimum knee angle is smaller than 90 degrees, the squat is considered correct, because it means the user is going deep enough.
 
     Args:
-        video_path (str): path to the video
-        model (YOLO): model to use
+        json_path (str): path to the JSON file containing the keypoints
 
     Returns:
-        keypoints (list): list of keypoints
+        is_correct (bool): True if the minimum knee angle is smaller than a threshold, False otherwise
     """
+    angles_list = compute_right_knee_angle_list(json_path)
+    for angle in angles_list:
+        if angle < 90:
+            return True
+    return False
 
-    keypoints = []
-    results = model(video_path, save=True, show_conf=False, show_boxes=False)
-    for frame in results:
-        tensor = frame.keypoints.xy[0]
-        keypoints.append(tensor.tolist())
-
-    return keypoints
-
-
-
-tools = [compute_right_knee_angle]
+tools = [get_keypoints_from_keypoints, check_knee_angle]
 
 prompt = ChatPromptTemplate.from_messages(
     [
         (
             "system",
-            "You are a helpful assistant. Make sure to use the compute_right_knee_angle tool for information.",
+            "You are a helpful assistant. Make sure to use the check_knee_angle tool if the user wants to check his movement.",
         ),
         ("placeholder", "{chat_history}"),
         ("human", "{input}"),
@@ -71,4 +91,6 @@ prompt = ChatPromptTemplate.from_messages(
 # Construct the Tools agent
 agent = create_tool_calling_agent(llm, tools, prompt)
 
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+response = agent_executor.invoke({"input": f"Is my squat correct ? The json file is in tmp/keypoints.json."})
+print(response["output"])

Modules/PoseEstimation/pose_estimator.py

@@ -1,5 +1,7 @@
 from ultralytics import YOLO
 import numpy as np
+import os
+import json
 
 id_joints_dict = {0: 'nose',
         1: 'left_eye',
@@ -22,14 +24,21 @@ 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):
-
+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 frame in results:
-        keypoints.append(frame.keypoints.xy)
+    for (i, frame) in enumerate(results):
+        frame_dict = {}
+        frame_dict['frame'] = i
+        frame_dict['keypoints'] = frame.keypoints.xy[0].tolist()
+        keypoints.append(frame_dict)
 
-    return keypoints
+    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):
 
@@ -53,45 +62,25 @@ def calculate_angle(a, b, c):
 
     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 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):