def viz_predictions(
input_: np.ndarray,
output: np.ndarray,
target: np.ndarray,
centerlines: np.ndarray,
city_names: np.ndarray,
idx=None,
show: bool = True,
) -> None:
"""Visualize predicted trjectories.
Args:
input_ (numpy array): Input Trajectory with shape (num_tracks x obs_len x 2)
output (numpy array of list): Top-k predicted trajectories, each with shape (num_tracks x pred_len x 2)
target (numpy array): Ground Truth Trajectory with shape (num_tracks x pred_len x 2)
centerlines (numpy array of list of centerlines): Centerlines (Oracle/Top-k) for each trajectory
city_names (numpy array): city names for each trajectory
show (bool): if True, show
"""
num_tracks = input_.shape[0]
obs_len = input_.shape[1]
pred_len = target.shape[1]
plt.figure(0, figsize=(8, 7))
avm = ArgoverseMap()
for i in range(num_tracks):
plt.plot(
input_[i, :, 0],
input_[i, :, 1],
color="#ECA154",
label="Observed",
alpha=1,
linewidth=3,
zorder=15,
)
plt.plot(
input_[i, -1, 0],
input_[i, -1, 1],
"o",
color="#ECA154",
label="Observed",
alpha=1,
linewidth=3,
zorder=15,
markersize=9,
)
plt.plot(
target[i, :, 0],
target[i, :, 1],
color="#d33e4c",
label="Target",
alpha=1,
linewidth=3,
zorder=20,
)
plt.plot(
target[i, -1, 0],
target[i, -1, 1],
"o",
color="#d33e4c",
label="Target",
alpha=1,
linewidth=3,
zorder=20,
markersize=9,
)
for j in range(len(centerlines[i])):
plt.plot(
centerlines[i][j][:, 0],
centerlines[i][j][:, 1],
"--",
color="grey",
alpha=1,
linewidth=1,
zorder=0,
)
for j in range(len(output[i])):
plt.plot(
output[i][j][:, 0],
output[i][j][:, 1],
color="#007672",
label="Predicted",
alpha=1,
linewidth=3,
zorder=15,
)
plt.plot(
output[i][j][-1, 0],
output[i][j][-1, 1],
"o",
color="#007672",
label="Predicted",
alpha=1,
linewidth=3,
zorder=15,
markersize=9,
)
for k in range(pred_len):
lane_ids = avm.get_lane_ids_in_xy_bbox(
output[i][j][k, 0],
output[i][j][k, 1],
city_names[i],
query_search_range_manhattan=2.5,
)
for j in range(obs_len):
lane_ids = avm.get_lane_ids_in_xy_bbox(
input_[i, j, 0],
input_[i, j, 1],
city_names[i],
query_search_range_manhattan=2.5,
)
[avm.draw_lane(lane_id, city_names[i]) for lane_id in lane_ids]
for j in range(pred_len):
lane_ids = avm.get_lane_ids_in_xy_bbox(
target[i, j, 0],
target[i, j, 1],
city_names[i],
query_search_range_manhattan=2.5,
)
[avm.draw_lane(lane_id, city_names[i]) for lane_id in lane_ids]
plt.axis("equal")
plt.xticks([])
plt.yticks([])
handles, labels = plt.gca().get_legend_handles_labels()
by_label = OrderedDict(zip(labels, handles))
if show:
plt.show()
for k in range(pred_len):
lane_ids = avm.get_lane_ids_in_xy_bbox(
output[i][j][k, 0],
output[i][j][k, 1],
city_names[i],
query_search_range_manhattan=2.5,
)
for j in range(obs_len):
lane_ids = avm.get_lane_ids_in_xy_bbox(
input_[i, j, 0],
input_[i, j, 1],
city_names[i],
query_search_range_manhattan=2.5,
)
[avm.draw_lane(lane_id, city_names[i]) for lane_id in lane_ids]
for j in range(pred_len):
lane_ids = avm.get_lane_ids_in_xy_bbox(
target[i, j, 0],
target[i, j, 1],
city_names[i],
query_search_range_manhattan=2.5,
)
[avm.draw_lane(lane_id, city_names[i]) for lane_id in lane_ids]
plt.axis("equal")
plt.xticks([])
plt.yticks([])
handles, labels = plt.gca().get_legend_handles_labels()
by_label = OrderedDict(zip(labels, handles))
plt.legend()