File size: 1,923 Bytes
d4cbafd | 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 | import numpy as np
import pdb
def prediction_output_to_trajectories(prediction_output_dict,
dt,
max_h,
ph,
map=None,
prune_ph_to_future=False):
prediction_timesteps = prediction_output_dict.keys()
output_dict = dict()
histories_dict = dict()
futures_dict = dict()
for t in prediction_timesteps:
histories_dict[t] = dict()
output_dict[t] = dict()
futures_dict[t] = dict()
prediction_nodes = prediction_output_dict[t].keys()
for node in prediction_nodes:
predictions_output = prediction_output_dict[t][node]
position_state = {'position': ['x', 'y']}
history = node.get(np.array([t - max_h, t]), position_state) # History includes current pos
history = history[~np.isnan(history.sum(axis=1))]
#pdb.set_trace()
future = node.get(np.array([t + 1, t + ph]), position_state)
# replace nan to 0
#future[np.isnan(future)] = 0
future = future[~np.isnan(future.sum(axis=1))]
if prune_ph_to_future:
predictions_output = predictions_output[:, :, :future.shape[0]]
if predictions_output.shape[2] == 0:
continue
trajectory = predictions_output
if map is None:
histories_dict[t][node] = history
output_dict[t][node] = trajectory
futures_dict[t][node] = future
else:
histories_dict[t][node] = map.to_map_points(history)
output_dict[t][node] = map.to_map_points(trajectory)
futures_dict[t][node] = map.to_map_points(future)
return output_dict, histories_dict, futures_dict
|