73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
|
import numpy as np
|
||
|
|
from scipy.integrate import cumtrapz
|
||
|
|
|
||
|
|
|
||
|
|
def integrate(f, dx, F0=0.):
|
||
|
|
N = f.shape[0]
|
||
|
|
return F0 + np.hstack((np.zeros((N, 1)), cumtrapz(f, axis=1, dx=dx)))
|
||
|
|
|
||
|
|
|
||
|
|
def integrate_trajectory(v, x0, dt):
|
||
|
|
xd_ = integrate(v[..., 0], dx=dt, F0=x0[0])
|
||
|
|
yd_ = integrate(v[..., 1], dx=dt, F0=x0[1])
|
||
|
|
integrated = np.stack([xd_, yd_], axis=2)
|
||
|
|
return integrated
|
||
|
|
|
||
|
|
|
||
|
|
def prediction_output_to_trajectories(prediction_output_dict,
|
||
|
|
dt,
|
||
|
|
max_h,
|
||
|
|
ph,
|
||
|
|
map=None,
|
||
|
|
gmm_agg='mean',
|
||
|
|
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']}
|
||
|
|
velocity_state = {'velocity': ['x', 'y']}
|
||
|
|
acceleration_state = {'acceleration': ['m']}
|
||
|
|
history = node.get(np.array([t - max_h, t]), position_state) # History includes current pos
|
||
|
|
history = history[~np.isnan(history.sum(axis=1))]
|
||
|
|
|
||
|
|
future = node.get(np.array([t + 1, t + ph]), position_state)
|
||
|
|
future = future[~np.isnan(future.sum(axis=1))]
|
||
|
|
|
||
|
|
current_pos = node.get(t, position_state)[0] # List with single item
|
||
|
|
current_vel = node.get(t, velocity_state)[0] # List with single item
|
||
|
|
|
||
|
|
predictions_output = getattr(predictions_output, gmm_agg)(axis=1)
|
||
|
|
|
||
|
|
if prune_ph_to_future:
|
||
|
|
predictions_output = predictions_output[:, :future.shape[0]]
|
||
|
|
if predictions_output.shape[1] == 0:
|
||
|
|
continue
|
||
|
|
|
||
|
|
vel_broad = np.expand_dims(np.broadcast_to(current_vel,
|
||
|
|
(predictions_output.shape[0],
|
||
|
|
current_vel.shape[-1])), axis=-2)
|
||
|
|
vel = np.concatenate((vel_broad, predictions_output), axis=1)
|
||
|
|
trajectory = integrate_trajectory(vel, current_pos, dt=dt)[:, 1:]
|
||
|
|
|
||
|
|
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
|