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

Modules/PoseEstimation/pose_agent.py

@@ -3,6 +3,7 @@ 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 torch
 import os
 import sys
 import json
@@ -11,6 +12,7 @@ from Modules.PoseEstimation.pose_estimator import model, compute_right_knee_angl
 
 # 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")
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 @tool
 def get_keypoints_from_keypoints(video_path: str) -> str:
@@ -27,7 +29,7 @@ def get_keypoints_from_keypoints(video_path: str) -> str:
     save_folder='tmp'
     os.makedirs(save_folder, exist_ok=True)
     keypoints = []
-    results = model(video_path, save=True, show_conf=False, show_boxes=False)
+    results = model(video_path, save=True, show_conf=False, show_boxes=False, device=device)
     for (i, frame) in enumerate(results):
         frame_dict = {}
         frame_dict['frame'] = i
@@ -77,7 +79,7 @@ def check_knee_angle(json_path: str) -> bool:
     return False
 
 @tool
-def check_squat(video_path: str) -> bool:
+def check_squat(file_name: str) -> str:
     """
     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.
@@ -89,8 +91,17 @@ def check_squat(video_path: str) -> bool:
     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)
+
+    video_path = os.path.join('uploaded', file_name)
+    if os.path.exists(video_path):
+        json_path = get_keypoints_from_keypoints(video_path)
+        is_correct = check_knee_angle(json_path)
+        if is_correct:
+            return "The squat is correct because your knee angle is smaller than 90 degrees."
+        else:
+            return "The squat is incorrect because your knee angle is greater than 90 degrees."
+    else:
+        return "The video file does not exist."
 
 tools = [check_squat]
 
@@ -98,7 +109,7 @@ 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. Also explain your response",
+            "You are a helpful assistant. Make sure to use the check_squat tool if the user wants to check his movement. Also explain your response",
         ),
         ("placeholder", "{chat_history}"),
         ("human", "{input}"),
@@ -109,6 +120,4 @@ 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)
-response = agent_executor.invoke({"input": f"Is my squat correct ? The video file is in data/pose/squat.mp4."})
-print(response["output"])
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

app.py

@@ -9,15 +9,15 @@ import json
 from dotenv import load_dotenv
 load_dotenv() # load .env api keys 
 import os
-
+import shutil
 from Modules.rag import rag_chain
 from Modules.router import router_chain
 from Modules.workout_plan import workout_chain
-# from Modules.PoseEstimation.pose_agent import agent_executor
+from Modules.PoseEstimation.pose_agent import agent_executor
 
 mistral_api_key = os.getenv("MISTRAL_API_KEY")
 from Modules.PoseEstimation import pose_estimator
-from utils import save_uploaded_file
+from utils import save_uploaded_file, encode_video_H264
 
 
 st.set_page_config(layout="wide", initial_sidebar_state="collapsed")
@@ -44,6 +44,9 @@ with col1:
 
     if "messages" not in st.session_state:
         st.session_state.messages = []
+    
+    if "file_name" not in st.session_state:
+        st.session_state.file_name = None
 
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
@@ -54,25 +57,28 @@ with col1:
         with st.chat_message("user"):
             st.markdown(prompt)
 
-        with st.chat_message("assistant"):
+        with st.chat_message("assistant", avatar="data/AI_Bro.png"):
             # Build answer from LLM
             direction = router_chain.invoke({"question":prompt})
             print(type(direction))
             print(direction)
             if direction=='fitness_advices':
-                response = rag_chain.invoke(
-                            prompt
-                            )
+                with st.spinner("Thinking..."):
+                    response = rag_chain.invoke(
+                                prompt
+                                )
             elif direction=='smalltalk':
-                response = base_chain.invoke(
-                    {"question":prompt}
-                ).content
+                with st.spinner("Thinking..."):
+                    response = base_chain.invoke(
+                        {"question":prompt}
+                    ).content
             elif direction =='movement_analysis':
-                response = "I can't do that for the moment"
-            #     response = agent_executor.invoke(
-            #         {"input" : instruction}
-            #     )["output"]
-            # elif direction == 'workout_plan':
+                if st.session_state.file_name is not None:
+                    prompt += "the file name is " + st.session_state.file_name
+                with st.spinner("Analyzing movement..."):
+                    response = agent_executor.invoke(
+                        {"input" : prompt}
+                    )["output"]
             else:
                 response = "Sure! I just made a workout for you. Check on the table I just provided you."
                 json_output = workout_chain.invoke({"query":prompt})
@@ -87,28 +93,35 @@ with col1:
     if display_workout:
         st.subheader("Workout")
         st.data_editor(workout_df)
-        # TO DO 
 # Second column containers
 with col2:
-    st.subheader("Sports Agenda")
+    # st.subheader("Sports Agenda")
         # TO DO
     st.subheader("Video Analysis")
-    ask_video = st.empty()
-    if video_uploaded is None:
-        video_uploaded = ask_video.file_uploader("Choose a video file", type=["mp4", "ogg", "webm"])
+
+    video_uploaded = st.file_uploader("Choose a video file", type=["mp4", "ogg", "webm", "MOV"])
     if video_uploaded:
         video_uploaded = save_uploaded_file(video_uploaded)
-        ask_video.empty()
-        _left, mid, _right = st.columns(3)
+        if video_uploaded.split("/")[-1] != st.session_state.file_name:
+            shutil.rmtree('fig', ignore_errors=True)
+            shutil.rmtree('/home/user/.pyenv/runs', ignore_errors=True)
+            st.session_state.file_name = None
+        st.session_state.file_name = video_uploaded.split("/")[-1]
+        _left, mid, _right = st.columns([1, 3, 1])
         with mid:
-            if os.path.exists('runs'):
-                st.video(os.path.join('runs', 'pose', 'predict', 'squat.mp4'), loop=True)
+            if os.path.exists('/home/user/.pyenv/runs'):
+                predict_list = os.listdir(os.path.join('/home/user/.pyenv/runs', 'pose'))
+                predict_list.sort()
+                predict_dir = predict_list[-1]
+                file_name = os.listdir(os.path.join('/home/user/.pyenv/runs', 'pose', predict_dir))[0]
+                file_path =os.path.join('/home/user/.pyenv/runs', 'pose', predict_dir, file_name)
+                file_path = encode_video_H264(file_path, remove_original=True)
+                st.video(file_path, loop=True)
             else :
                 st.video(video_uploaded)
             
-
-    st.subheader("Graph Displayer")
     if os.path.exists('fig'):
+        st.subheader("Graph Displayer")
         file_list = os.listdir('fig')
         for file in file_list:
             st.image(os.path.join('fig', file))

requirements.txt

@@ -1,20 +1,23 @@
 transformers
 torch
-langchain-core 
-langchain-mistralai
 pandas
-langchain-community
 streamlit-audiorec
 openai-whisper
 tiktoken
-langchain
 bs4
 chromadb
-langgraph
-langchainhub
 pypdf
 duckduckgo-search
 python-dotenv
 pypdf
 chromadb
-ultralytics
+moviepy
+ultralytics
+langchain==0.1.16
+langchain-community==0.0.34
+langchain-core==0.1.45
+langchain-experimental==0.0.57
+langchain-mistralai==0.1.2
+langchain-text-splitters==0.0.1
+langchainhub==0.1.15
+langsmith==0.1.50

utils.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import os
+from moviepy.editor import VideoFileClip
 
 def save_uploaded_file(uploaded_file):
     try:
@@ -9,4 +10,24 @@ def save_uploaded_file(uploaded_file):
         return file_path
     except Exception as e:
         st.error(f"Error: {e}")
-        return None
+        return None
+    
+def encode_video_H264(video_path, remove_original=False):
+    """
+    Encode video to H264 codec
+    
+    Args:
+        video_path (str): path to video to be encoded
+        remove_original (bool): whether to remove original video after encoding
+    Returns:
+        output_path (str): path to encoded video
+    """
+
+    output_path = video_path.split('.')[0] + '_H264.mp4'
+    clip = VideoFileClip(video_path)
+    clip.write_videofile(output_path, codec='libx264')
+    if remove_original:
+        os.remove(video_path)
+    clip.close()
+
+    return output_path