# adapted from Trajectron++ online_server.py import logging from multiprocessing import Queue import os import time import json import torch import dill import random import pathlib import numpy as np from trajectron.utils import prediction_output_to_trajectories from trajectron.model.online.online_trajectron import OnlineTrajectron from trajectron.model.model_registrar import ModelRegistrar from trajectron.environment import Environment, Scene import matplotlib.pyplot as plt import zmq logger = logging.getLogger("trajpred.inference") # if not torch.cuda.is_available() or self.config.device == 'cpu': # self.config.device = torch.device('cpu') # else: # if torch.cuda.device_count() == 1: # # If you have CUDA_VISIBLE_DEVICES set, which you should, # # then this will prevent leftover flag arguments from # # messing with the device allocation. # self.config.device = 'cuda:0' # self.config.device = torch.device(self.config.device) def create_online_env(env, hyperparams, scene_idx, init_timestep): test_scene = env.scenes[scene_idx] online_scene = Scene(timesteps=init_timestep + 1, map=test_scene.map, dt=test_scene.dt) online_scene.nodes = test_scene.get_nodes_clipped_at_time( timesteps=np.arange(init_timestep - hyperparams['maximum_history_length'], init_timestep + 1), state=hyperparams['state']) online_scene.robot = test_scene.robot online_scene.calculate_scene_graph(attention_radius=env.attention_radius, edge_addition_filter=hyperparams['edge_addition_filter'], edge_removal_filter=hyperparams['edge_removal_filter']) return Environment(node_type_list=env.node_type_list, standardization=env.standardization, scenes=[online_scene], attention_radius=env.attention_radius, robot_type=env.robot_type) def get_maps_for_input(input_dict, scene, hyperparams): scene_maps = list() scene_pts = list() heading_angles = list() patch_sizes = list() nodes_with_maps = list() for node in input_dict: if node.type in hyperparams['map_encoder']: x = input_dict[node] me_hyp = hyperparams['map_encoder'][node.type] if 'heading_state_index' in me_hyp: heading_state_index = me_hyp['heading_state_index'] # We have to rotate the map in the opposit direction of the agent to match them if type(heading_state_index) is list: # infer from velocity or heading vector heading_angle = -np.arctan2(x[-1, heading_state_index[1]], x[-1, heading_state_index[0]]) * 180 / np.pi else: heading_angle = -x[-1, heading_state_index] * 180 / np.pi else: heading_angle = None scene_map = scene.map[node.type] map_point = x[-1, :2] patch_size = hyperparams['map_encoder'][node.type]['patch_size'] scene_maps.append(scene_map) scene_pts.append(map_point) heading_angles.append(heading_angle) patch_sizes.append(patch_size) nodes_with_maps.append(node) if heading_angles[0] is None: heading_angles = None else: heading_angles = torch.Tensor(heading_angles) maps = scene_maps[0].get_cropped_maps_from_scene_map_batch(scene_maps, scene_pts=torch.Tensor(scene_pts), patch_size=patch_sizes[0], rotation=heading_angles) maps_dict = {node: maps[[i]] for i, node in enumerate(nodes_with_maps)} return maps_dict class InferenceServer: def __init__(self, config: dict): self.config = config context = zmq.Context() self.trajectory_socket: zmq.Socket = context.socket(zmq.SUB) self.trajectory_socket.connect(config.zmq_trajectory_addr) self.trajectory_socket.setsockopt(zmq.SUBSCRIBE, b'') self.prediction_socket: zmq.Socket = context.socket(zmq.PUB) self.prediction_socket.bind(config.zmq_prediction_addr) print(self.prediction_socket) def run(self): if self.config.seed is not None: random.seed(self.config.seed) np.random.seed(self.config.seed) torch.manual_seed(self.config.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(self.config.seed) # Choose one of the model directory names under the experiment/*/models folders. # Possibilities are 'vel_ee', 'int_ee', 'int_ee_me', or 'robot' # model_dir = os.path.join(self.config.log_dir, 'int_ee') # model_dir = 'models/models_04_Oct_2023_21_04_48_eth_vel_ar3' # Load hyperparameters from json config_file = os.path.join(self.config.model_dir, self.config.conf) if not os.path.exists(config_file): raise ValueError('Config json not found!') with open(config_file, 'r') as conf_json: hyperparams = json.load(conf_json) # Add hyperparams from arguments hyperparams['dynamic_edges'] = self.config.dynamic_edges hyperparams['edge_state_combine_method'] = self.config.edge_state_combine_method hyperparams['edge_influence_combine_method'] = self.config.edge_influence_combine_method hyperparams['edge_addition_filter'] = self.config.edge_addition_filter hyperparams['edge_removal_filter'] = self.config.edge_removal_filter hyperparams['batch_size'] = self.config.batch_size hyperparams['k_eval'] = self.config.k_eval hyperparams['offline_scene_graph'] = self.config.offline_scene_graph hyperparams['incl_robot_node'] = self.config.incl_robot_node hyperparams['edge_encoding'] = not self.config.no_edge_encoding hyperparams['use_map_encoding'] = self.config.map_encoding output_save_dir = os.path.join(self.config.output_dir, 'pred_figs') pathlib.Path(output_save_dir).mkdir(parents=True, exist_ok=True) with open(self.config.eval_data_dict, 'rb') as f: eval_env = dill.load(f, encoding='latin1') if eval_env.robot_type is None and hyperparams['incl_robot_node']: eval_env.robot_type = eval_env.NodeType[0] # TODO: Make more general, allow the user to specify? for scene in eval_env.scenes: scene.add_robot_from_nodes(eval_env.robot_type) logger.info('Loaded data from %s' % (self.config.eval_data_dict,)) # Creating a dummy environment with a single scene that contains information about the world. # When using this code, feel free to use whichever scene index or initial timestep you wish. scene_idx = 0 # You need to have at least acceleration, so you want 2 timesteps of prior data, e.g. [0, 1], # so that you can immediately start incremental inference from the 3rd timestep onwards. init_timestep = 1 eval_scene = eval_env.scenes[scene_idx] online_env = create_online_env(eval_env, hyperparams, scene_idx, init_timestep) model_registrar = ModelRegistrar(self.config.model_dir, self.config.eval_device) model_registrar.load_models(iter_num=100) trajectron = OnlineTrajectron(model_registrar, hyperparams, self.config.eval_device) # If you want to see what different robot futures do to the predictions, uncomment this line as well as # related "... += adjustment" lines below. # adjustment = np.stack([np.arange(13)/float(i*2.0) for i in range(6, 12)], axis=1) # Here's how you'd incrementally run the model, e.g. with streaming data. trajectron.set_environment(online_env, init_timestep) for timestep in range(init_timestep + 1, eval_scene.timesteps): input_dict = eval_scene.get_clipped_input_dict(timestep, hyperparams['state']) maps = None if hyperparams['use_map_encoding']: maps = get_maps_for_input(input_dict, eval_scene, hyperparams) robot_present_and_future = None if eval_scene.robot is not None and hyperparams['incl_robot_node']: robot_present_and_future = eval_scene.robot.get(np.array([timestep, timestep + hyperparams['prediction_horizon']]), hyperparams['state'][eval_scene.robot.type], padding=0.0) robot_present_and_future = np.stack([robot_present_and_future, robot_present_and_future], axis=0) # robot_present_and_future += adjustment start = time.time() dists, preds = trajectron.incremental_forward(input_dict, maps, prediction_horizon=6, num_samples=51, robot_present_and_future=robot_present_and_future, full_dist=True) end = time.time() logger.info("t=%d: took %.2f s (= %.2f Hz) w/ %d nodes and %d edges" % (timestep, end - start, 1. / (end - start), len(trajectron.nodes), trajectron.scene_graph.get_num_edges())) # unsure what this bit from online_prediction.py does: # detailed_preds_dict = dict() # for node in eval_scene.nodes: # if node in preds: # detailed_preds_dict[node] = preds[node] #adapted from trajectron.visualization # prediction_dict provides the actual predictions # histories_dict provides the trajectory used for prediction # futures_dict is the Ground Truth, which is unvailable in an online setting prediction_dict, histories_dict, futures_dict = prediction_output_to_trajectories({timestep: preds}, eval_scene.dt, hyperparams['maximum_history_length'], hyperparams['prediction_horizon'] ) assert(len(prediction_dict.keys()) <= 1) if len(prediction_dict.keys()) == 0: return ts_key = list(prediction_dict.keys())[0] prediction_dict = prediction_dict[ts_key] histories_dict = histories_dict[ts_key] futures_dict = futures_dict[ts_key] response = {} for node in histories_dict: history = histories_dict[node] # future = futures_dict[node] predictions = prediction_dict[node] if np.isnan(history[-1]).any(): continue response[node.id] = { 'id': node.id, 'history': history.tolist(), 'predictions': predictions[0].tolist() # use batch 0 } data = json.dumps(response) self.prediction_socket.send_string(data) # time.sleep(1) # print(prediction_dict) # print(histories_dict) # print(futures_dict) def run_inference_server(config): s = InferenceServer(config) s.run()