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

Modules/PoseEstimation/curl_agent.py

@@ -0,0 +1,147 @@
+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
+
+from operator import itemgetter
+from typing import Dict, List, Union
+
+from langchain_core.messages import AIMessage
+from langchain_core.runnables import (
+    Runnable,
+    RunnableLambda,
+    RunnableMap,
+    RunnablePassthrough,
+)
+
+import numpy as np
+
+# 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 shoulder_angle(pose: list) -> float:
+    
+    """
+    Computes the shoulder angle.
+
+    Args:
+        pose (list): list of keypoints
+
+    Returns:
+        arm_angle (float): arm angle with chest
+    """
+    right_elbow = pose[joints_id_dict['right_elbow']]
+    right_shoulder = pose[joints_id_dict['right_shoulder']]
+    right_hip = pose[joints_id_dict['right_hip']]   
+
+    left_elbow = pose[joints_id_dict['left_elbow']]
+    left_shoulder = pose[joints_id_dict['left_shoulder']]
+    left_hip = pose[joints_id_dict['left_hip']] 
+
+    right_arm_angle = calculate_angle(right_elbow, right_shoulder, right_hip)
+    left_arm_angle = calculate_angle(left_elbow, left_shoulder, left_hip)
+
+    return right_arm_angle
+
+
+@tool
+def elbow_angle(pose):
+    """
+    Computes the elbow angle.
+
+    Args:
+        pose (list): list of keypoints
+
+    Returns:
+        elbow_angle (float): elbow angle with chest
+    """
+    right_elbow = pose[joints_id_dict['right_elbow']]
+    right_shoulder = pose[joints_id_dict['right_shoulder']]
+    right_wrist = pose[joints_id_dict['right_wrist']]   
+
+    left_elbow = pose[joints_id_dict['left_elbow']]
+    left_shoulder = pose[joints_id_dict['left_shoulder']]
+    left_wrist = pose[joints_id_dict['left_wrist']] 
+
+    right_elbow_angle = calculate_angle(right_shoulder, right_elbow, right_wrist)
+    left_elbow_angle = calculate_angle(left_shoulder, left_elbow, left_wrist)
+
+    return right_elbow_angle
+
+
+tools = [shoulder_angle, elbow_angle]
+
+llm_with_tools = llm.bind_tools(tools)
+tool_map = {tool.name: tool for tool in tools}
+
+# prompt = ChatPromptTemplate.from_messages(
+#     [
+#         (
+#             "system",
+#             "You are a helpful assistant. Make sure to use the compute_right_knee_angle tool for information.",
+#         ),
+#         ("placeholder", "{chat_history}"),
+#         ("human", "{input}"),
+#         ("placeholder", "{agent_scratchpad}"),
+#     ]
+# )
+
+# Construct the Tools agent
+# curl_agent = create_tool_calling_agent(llm, tools, prompt)
+
+
+pose_sequence = [
+    # Pose 1
+    [
+        # Head
+        [50, 50],
+        # Shoulders
+        [40, 80], [60, 80],
+        # Elbows
+        [30, 110], [70, 110],
+        # Wrists
+        [25, 140], [75, 140],
+        # Hips
+        [45, 180], [55, 180],
+        # Knees
+        [40, 220], [60, 220],
+        # Ankles
+        [35, 250], [65, 250],
+    ],
+    # Pose 2
+    [
+        # Head
+        [60, 60],
+        # Shoulders
+        [50, 90], [70, 90],
+        # Elbows
+        [40, 120], [80, 120],
+        # Wrists
+        [35, 150], [85, 150],
+        # Hips
+        [55, 180], [65, 180],
+        # Knees
+        [50, 220], [70, 220],
+        # Ankles
+        [45, 250], [75, 250],
+    ]]
+
+    # Create an agent executor by passing in the agent and tools
+# agent_executor = AgentExecutor(agent=curl_agent, tools=tools, verbose=True)
+# agent_executor.invoke({"input": f"Compute shoulder and elbow angle and display them given the following pose estimation: {pose_sequence[0]}"})
+
+def call_tools(msg: AIMessage) -> Runnable:
+    """Simple sequential tool calling helper."""
+    tool_map = {tool.name: tool for tool in tools}
+    tool_calls = msg.tool_calls.copy()
+    for tool_call in tool_calls:
+        tool_call["output"] = tool_map[tool_call["name"]].invoke(tool_call["args"])
+    return tool_calls
+
+
+chain = llm_with_tools | call_tools
+
+print(chain.invoke(f"What is the shoulder angle and elbow angle given the following pose estimation: {pose_sequence[0]}"))

Modules/PoseEstimation/pose_agent.py

@@ -0,0 +1,74 @@
+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
+
+# 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:
+
+    """
+    Computes the knee angle.
+
+    Args:
+        pose (list): list of keypoints
+
+    Returns:
+        knee_angle (float): knee angle
+    """
+
+    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)
+
+    print(knee_angle)
+
+    return str(knee_angle)
+
+@tool
+def get_keypoints_from_path(video_path: str):
+    """
+    Get keypoints from a video.
+
+    Args:
+        video_path (str): path to the video
+        model (YOLO): model to use
+
+    Returns:
+        keypoints (list): list of keypoints
+    """
+
+    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]
+
+prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            "You are a helpful assistant. Make sure to use the compute_right_knee_angle tool for information.",
+        ),
+        ("placeholder", "{chat_history}"),
+        ("human", "{input}"),
+        ("placeholder", "{agent_scratchpad}"),
+    ]
+)
+
+# Construct the Tools agent
+agent = create_tool_calling_agent(llm, tools, prompt)
+
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

Modules/PoseEstimation/pose_estimator.py

@@ -31,8 +31,6 @@ def get_keypoints_from_keypoints(model, video_path):
 
     return keypoints
 
-keypoints = get_keypoints_from_keypoints(model, '../../data/pose/squat.mp4')
-
 def calculate_angle(a, b, c):
 
     """
@@ -112,5 +110,4 @@ 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

app.py

@@ -10,6 +10,8 @@ import os
 from Modules.rag import rag_chain
 
 mistral_api_key = os.getenv("MISTRAL_API_KEY")
+from Modules.PoseEstimation import pose_estimator
+from utils import save_uploaded_file
 
 def format_messages(messages):
     formatted_messages = ""
@@ -72,12 +74,17 @@ with col2:
     if video_uploaded is None:
         video_uploaded = ask_video.file_uploader("Choose a video file", type=["mp4", "ogg", "webm"])
     if video_uploaded:
+        video_uploaded = save_uploaded_file(video_uploaded)
         ask_video.empty()
-        with st.spin("Processing video"):
-            pass # TO DO
         _left, mid, _right = st.columns(3)
         with mid:
             st.video(video_uploaded)
+            apply_pose = st.button("Apply Pose Estimation")
+
+        if apply_pose:
+            with st.spinner("Processing video"):
+                keypoints = pose_estimator.get_keypoints_from_keypoints(pose_estimator.model, video_uploaded)
+            
 
     st.subheader("Graph Displayer")
         # TO DO

utils.py

@@ -0,0 +1,12 @@
+import streamlit as st
+import os
+
+def save_uploaded_file(uploaded_file):
+    try:
+        file_path = os.path.join('uploaded', uploaded_file.name)
+        with open(file_path, 'wb') as f:
+            f.write(uploaded_file.getvalue())
+        return file_path
+    except Exception as e:
+        st.error(f"Error: {e}")
+        return None