Spaces:
Sleeping
Sleeping
File size: 2,033 Bytes
a755c90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 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) |