Spaces:
Sleeping
Sleeping
File size: 4,286 Bytes
e051030 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
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]}"))
|