Merge branch 'main' of https://huggingface.co/spaces/EntrepreneurFirst/FitnessEquation

Modules/PoseEstimation/pose_agent.py

@@ -7,7 +7,7 @@ import os
 import sys
 import json
 sys.path.append(os.getcwd())
-from Modules.PoseEstimation.pose_estimator import model, compute_right_knee_angle, moving_average
+from Modules.PoseEstimation.pose_estimator import model, compute_right_knee_angle, moving_average, save_knee_angle_fig
 
 # 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")
@@ -54,9 +54,11 @@ def compute_right_knee_angle_list(json_path: str) -> list[float]:
     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
+    right_knee_angle_list = moving_average(right_knee_angle_list, 10)
+    save_knee_angle_fig(right_knee_angle_list)
+    return right_knee_angle_list
+
 def check_knee_angle(json_path: str) -> bool:
     """
     Checks if the minimum knee angle is smaller than a threshold.
@@ -74,13 +76,29 @@ def check_knee_angle(json_path: str) -> bool:
             return True
     return False
 
-tools = [get_keypoints_from_keypoints, check_knee_angle]
+@tool
+def check_squat(video_path: str) -> bool:
+    """
+    Checks if the squat is correct.
+    This function uses the check_knee_angle tool to check if the squat is correct. 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 file
+
+    Returns:
+        is_correct (bool): True if the squat is correct, False otherwise
+    """
+    json_path = get_keypoints_from_keypoints(video_path)
+    return check_knee_angle(json_path)
+
+tools = [check_squat]
 
 prompt = ChatPromptTemplate.from_messages(
     [
         (
             "system",
-            "You are a helpful assistant. Make sure to use the check_knee_angle tool if the user wants to check his movement.",
+            "You are a helpful assistant. Make sure to use the check_knee_angle tool if the user wants to check his movement. Also explain your response",
         ),
         ("placeholder", "{chat_history}"),
         ("human", "{input}"),
@@ -92,5 +110,5 @@ prompt = ChatPromptTemplate.from_messages(
 agent = create_tool_calling_agent(llm, tools, prompt)
 
 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."})
+response = agent_executor.invoke({"input": f"Is my squat correct ? The video file is in data/pose/squat.mp4."})
 print(response["output"])

Modules/PoseEstimation/pose_estimation.ipynb

@@ -1062,7 +1062,12 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "plt.plot(angles)"
+    "os.makedirs('fig', exist_ok=True)\n",
+    "plt.plot(angles)\n",
+    "plt.xlabel('Frame')\n",
+    "plt.ylabel('Knee Angle')\n",
+    "plt.title('Evolution of the knee angle')\n",
+    "plt.savefig('fig/knee_angle.png')"
    ]
   },
   {

Modules/PoseEstimation/pose_estimator.py

@@ -2,6 +2,7 @@ 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',
@@ -99,4 +100,12 @@ def moving_average(data, window_size):
     for i in range(len(data) - window_size + 1):
         avg.append(sum(data[i:i + window_size]) / window_size)
 
-    return avg
+    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')