POSE AGENT WORKING

Modules/PoseEstimation/pose_agent.py

@@ -77,7 +77,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 +89,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 +107,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 +118,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

@@ -7,10 +7,10 @@ from langchain_core.prompts import ChatPromptTemplate
 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.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
@@ -39,6 +39,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"]):
@@ -52,45 +55,53 @@ with col1:
         with st.chat_message("assistant"):
             # Build answer from LLM
             direction = router_chain.invoke({"question":prompt})
+
             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
-            # elif direction =='movement_analysis':
-            #     response = agent_executor.invoke(
-            #         {"input" : instruction}
-            #     )["output"]
-            print(type(response))
+                with st.spinner("Thinking..."):
+                    response = base_chain.invoke(
+                        {"question":prompt}
+                    ).content
+            elif direction =='movement_analysis':
+                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"]
             st.session_state.messages.append({"role": "assistant", "content": response})
             st.markdown(response)
 
-    st.subheader("Movement Analysis")
-        # 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()
+        if video_uploaded.split("/")[-1] != st.session_state.file_name:
+            shutil.rmtree('fig', ignore_errors=True)
+            shutil.rmtree('runs', ignore_errors=True)
+            st.session_state.file_name = None
+        st.session_state.file_name = video_uploaded.split("/")[-1]
         _left, mid, _right = st.columns(3)
         with mid:
             if os.path.exists('runs'):
-                st.video(os.path.join('runs', 'pose', 'predict', 'squat.mp4'), loop=True)
+                predict_list = os.listdir(os.path.join('runs', 'pose'))
+                predict_list.sort()
+                predict_dir = predict_list[-1]
+                st.video(os.path.join('runs', 'pose', predict_dir, 'squat.mp4'), 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))