import math
import os
import torch
import pickle
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import numpy.typing as npt
import fnmatch
import io
import seaborn as sns
import matplotlib.axes as Axes
import matplotlib.transforms as mtransforms
from PIL import Image
from functools import wraps
from typing import Sequence, Union, Optional
from tqdm import tqdm
from typing import List, Literal
from argparse import ArgumentParser
from scipy.ndimage.filters import gaussian_filter
from matplotlib.patches import FancyBboxPatch, Polygon, Rectangle, Circle
from matplotlib.collections import LineCollection
from torch_geometric.data import HeteroData, Dataset
from waymo_open_dataset.protos import scenario_pb2
from dev.utils.func import CONSOLE
from dev.modules.attr_tokenizer import Attr_Tokenizer
from dev.datasets.preprocess import TokenProcessor, cal_polygon_contour, AGENT_TYPE
from dev.datasets.scalable_dataset import WaymoTargetBuilder
__all__ = ['plot_occ_grid', 'plot_interact_edge', 'plot_map_edge', 'plot_insert_grid', 'plot_binary_map',
'plot_map_token', 'plot_prob_seed', 'plot_scenario', 'get_heatmap', 'draw_heatmap', 'plot_val', 'plot_tokenize']
def safe_run(func):
@wraps(func)
def wrapper1(*args, **kwargs):
try:
return func(*args, **kwargs)
except Exception as e:
print(e)
return
@wraps(func)
def wrapper2(*args, **kwargs):
return func(*args, **kwargs)
if int(os.getenv('DEBUG', 0)):
return wrapper2
else:
return wrapper1
@safe_run
def plot_occ_grid(scenario_id, occ, gt_occ=None, save_path='', mode='agent', prefix=''):
def generate_box_edges(matrix, find_value=1):
y, x = np.where(matrix == find_value)
edges = []
for xi, yi in zip(x, y):
edges.append([(xi - 0.5, yi - 0.5), (xi + 0.5, yi - 0.5)])
edges.append([(xi + 0.5, yi - 0.5), (xi + 0.5, yi + 0.5)])
edges.append([(xi + 0.5, yi + 0.5), (xi - 0.5, yi + 0.5)])
edges.append([(xi - 0.5, yi + 0.5), (xi - 0.5, yi - 0.5)])
return edges
os.makedirs(save_path, exist_ok=True)
n = int(math.sqrt(occ.shape[-1]))
plot_n = 3
plot_t = 5
occ_list = []
for i in range(plot_n):
for j in range(plot_t):
occ_list.append(occ[i, j].reshape(n, n))
occ_gt_list = []
if gt_occ is not None:
for i in range(plot_n):
for j in range(plot_t):
occ_gt_list.append(gt_occ[i, j].reshape(n, n))
row_labels = [f'n={n}' for n in range(plot_n)]
col_labels = [f't={t}' for t in range(plot_t)]
fig, axes = plt.subplots(plot_n, plot_t, figsize=(9, 6))
plt.subplots_adjust(wspace=0.1, hspace=0.1)
for i, ax in enumerate(axes.flat):
# NOTE: do not set vmin and vamx!
ax.imshow(occ_list[i], cmap='viridis', interpolation='nearest')
ax.axis('off')
if occ_gt_list:
gt_edges = generate_box_edges(occ_gt_list[i])
gts = LineCollection(gt_edges, colors='blue', linewidths=0.5)
ax.add_collection(gts)
insert_edges = generate_box_edges(occ_gt_list[i], find_value=-1)
inserts = LineCollection(insert_edges, colors='red', linewidths=0.5)
ax.add_collection(inserts)
ax.add_patch(plt.Rectangle((-0.5, -0.5), occ_list[i].shape[1], occ_list[i].shape[0],
linewidth=2, edgecolor='black', facecolor='none'))
for i, ax in enumerate(axes[:, 0]):
ax.annotate(row_labels[i], xy=(-0.1, 0.5), xycoords="axes fraction",
fontsize=12, ha="right", va="center", rotation=0)
for j, ax in enumerate(axes[0, :]):
ax.annotate(col_labels[j], xy=(0.5, 1.05), xycoords="axes fraction",
fontsize=12, ha="center", va="bottom")
plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_occ_{mode}.png'), dpi=500, bbox_inches='tight')
plt.close()
@safe_run
def plot_interact_edge(edge_index, scenario_ids, batch_sizes, num_seed, num_step, save_path='interact_edge_map',
**kwargs):
num_batch = len(scenario_ids)
batches = torch.cat([
torch.arange(num_batch).repeat_interleave(repeats=batch_sizes, dim=0),
torch.arange(num_batch).repeat_interleave(repeats=num_seed, dim=0),
], dim=0).repeat(num_step).numpy()
num_agent = batch_sizes.sum() + num_seed * num_batch
batch_sizes = torch.nn.functional.pad(batch_sizes, (1, 0), mode='constant', value=0)
ptr = torch.cumsum(batch_sizes, dim=0)
# assume difference scenarios and different timestep have the same number of seed agents
ptr_seed = torch.tensor(np.array([0] + [num_seed] * num_batch), device=ptr.device)
all_av_index = None
if 'av_index' in kwargs:
all_av_index = kwargs.pop('av_index').cpu() - ptr[:-1]
is_bos = np.zeros((batch_sizes.sum(), num_step)).astype(np.bool_)
if 'is_bos' in kwargs:
is_bos = kwargs.pop('is_bos').cpu().numpy()
src_index = torch.unique(edge_index[1])
for idx, src in enumerate(tqdm(src_index)):
src_batch = batches[src]
src_row = src % num_agent
if src_row // batch_sizes.sum() > 0:
seed_row = src_row % batch_sizes.sum() - ptr_seed[src_batch]
src_row = batch_sizes[src_batch + 1] + seed_row
else:
src_row = src_row - ptr[src_batch]
src_col = src // (num_agent)
src_mask = np.zeros((batch_sizes[src_batch + 1] + num_seed, num_step))
src_mask[src_row, src_col] = 1
tgt_mask = np.zeros((src_mask.shape[0], num_step))
tgt_index = edge_index[0, edge_index[1] == src]
for tgt in tgt_index:
tgt_batch = batches[tgt]
tgt_row = tgt % num_agent
if tgt_row // batch_sizes.sum() > 0:
seed_row = tgt_row % batch_sizes.sum() - ptr_seed[tgt_batch]
tgt_row = batch_sizes[tgt_batch + 1] + seed_row
else:
tgt_row = tgt_row - ptr[tgt_batch]
tgt_col = tgt // num_agent
tgt_mask[tgt_row, tgt_col] = 1
assert tgt_batch == src_batch
selected_step = tgt_mask.sum(axis=0) > 0
if selected_step.sum() > 1:
print(f"\nidx={idx}", src.item(), src_row.item(), src_col.item())
print(selected_step)
print(edge_index[:, edge_index[1] == src].tolist())
if all_av_index is not None:
kwargs['av_index'] = int(all_av_index[src_batch])
t = kwargs.get('t', src_col)
n = kwargs.get('n', 0)
is_bos_batch = is_bos[ptr[src_batch] : ptr[src_batch + 1]]
plot_binary_map(src_mask, tgt_mask, save_path, suffix=f'_{scenario_ids[src_batch]}_{t:02d}_{n:02d}_{idx:04d}',
is_bos=is_bos_batch, **kwargs)
@safe_run
def plot_map_edge(edge_index, pos_a, data, save_path='map_edge_map'):
map_points = data['map_point']['position'][:, :2].cpu().numpy()
token_pos = data['pt_token']['position'][:, :2].cpu().numpy()
token_heading = data['pt_token']['orientation'].cpu().numpy()
num_pt = token_pos.shape[0]
agent_index = torch.unique(edge_index[1])
for i in tqdm(agent_index):
xy = pos_a[i].cpu().numpy()
pt_index = edge_index[0, edge_index[1] == i].cpu().numpy()
pt_index = pt_index % num_pt
plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
_, ax = plt.subplots()
ax.set_axis_off()
plot_map_token(ax, map_points, token_pos[pt_index], token_heading[pt_index], colors='blue')
ax.scatter(xy[0], xy[1], s=0.5, c='red', edgecolors='none')
os.makedirs(save_path, exist_ok=True)
plt.savefig(os.path.join(save_path, f'map_{i}.png'), dpi=600, bbox_inches='tight')
plt.close()
def get_heatmap(x, y, prob, s=3, bins=1000):
heatmap, xedges, yedges = np.histogram2d(x, y, bins=bins, weights=prob, density=True)
heatmap = gaussian_filter(heatmap, sigma=s)
extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
return heatmap.T, extent
@safe_run
def draw_heatmap(vector, vector_prob, gt_idx):
fig, ax = plt.subplots(figsize=(10, 10))
vector_prob = vector_prob.cpu().numpy()
for j in range(vector.shape[0]):
if j in gt_idx:
color = (0, 0, 1)
else:
grey_scale = max(0, 0.9 - vector_prob[j])
color = (0.9, grey_scale, grey_scale)
# if lane[j, k, -1] == 0: continue
x0, y0, x1, y1, = vector[j, :4]
ax.plot((x0, x1), (y0, y1), color=color, linewidth=2)
return plt
@safe_run
def plot_insert_grid(scenario_id, prob, grid, ego_pos, map, save_path='', prefix='', inference=False, indices=None,
all_t_in_one=False):
"""
prob: float array of shape (num_step, num_grid)
grid: float array of shape (num_grid, 2)
"""
os.makedirs(save_path, exist_ok=True)
n = int(math.sqrt(prob.shape[1]))
# grid = grid[:, np.newaxis] + ego_pos[np.newaxis, ...]
for t in range(ego_pos.shape[0]):
plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
_, ax = plt.subplots()
# plot probability
prob_t = prob[t].reshape(n, n)
plt.imshow(prob_t, cmap='viridis', interpolation='nearest')
if indices is not None:
indice = indices[t]
if isinstance(indice, (int, float, np.int_)):
indice = [indice]
for _indice in indice:
if _indice == -1: continue
row = _indice // n
col = _indice % n
rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='red', facecolor='none', lw=2)
ax.add_patch(rect)
ax.grid(False)
ax.set_aspect('equal', adjustable='box')
plt.title('Prob of Rel Position Grid')
plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_heat_map_{t}.png'), dpi=300, bbox_inches='tight')
plt.close()
if all_t_in_one:
break
@safe_run
def plot_insert_grid(scenario_id, prob, indices=None, save_path='', prefix='', inference=False):
"""
prob: float array of shape (num_seed, num_step, num_grid)
grid: float array of shape (num_grid, 2)
"""
os.makedirs(save_path, exist_ok=True)
n = int(math.sqrt(prob.shape[-1]))
plot_n = 3
plot_t = 5
prob_list = []
for i in range(plot_n):
for j in range(plot_t):
prob_list.append(prob[i, j].reshape(n, n))
indice_list = []
if indices is not None:
for i in range(plot_n):
for j in range(plot_t):
indice_list.append(indices[i, j])
row_labels = [f'n={n}' for n in range(plot_n)]
col_labels = [f't={t}' for t in range(plot_t)]
fig, axes = plt.subplots(plot_n, plot_t, figsize=(9, 6))
fig.suptitle('Prob of Insert Position Grid')
plt.subplots_adjust(wspace=0.1, hspace=0.1)
for i, ax in enumerate(axes.flat):
ax.imshow(prob_list[i], cmap='viridis', interpolation='nearest')
ax.axis('off')
if indice_list:
row = indice_list[i] // n
col = indice_list[i] % n
rect = Rectangle((col - .5, row - .5), 1, 1, edgecolor='red', facecolor='none', lw=2)
ax.add_patch(rect)
ax.add_patch(plt.Rectangle((-0.5, -0.5), prob_list[i].shape[1], prob_list[i].shape[0],
linewidth=2, edgecolor='black', facecolor='none'))
for i, ax in enumerate(axes[:, 0]):
ax.annotate(row_labels[i], xy=(-0.1, 0.5), xycoords="axes fraction",
fontsize=12, ha="right", va="center", rotation=0)
for j, ax in enumerate(axes[0, :]):
ax.annotate(col_labels[j], xy=(0.5, 1.05), xycoords="axes fraction",
fontsize=12, ha="center", va="bottom")
ax.grid(False)
ax.set_aspect('equal', adjustable='box')
plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_insert_map.png'), dpi=500, bbox_inches='tight')
plt.close()
@safe_run
def plot_binary_map(src_mask, tgt_mask, save_path='', suffix='', av_index=None, is_bos=None, **kwargs):
from matplotlib.colors import ListedColormap
os.makedirs(save_path, exist_ok=True)
fig, axes = plt.subplots(1, 2, figsize=(10, 8))
title = []
if kwargs.get('t', None) is not None:
t = kwargs['t']
title.append(f't={t}')
if kwargs.get('n', None) is not None:
n = kwargs['n']
title.append(f'n={n}')
plt.title(' '.join(title))
cmap = ListedColormap(['white', 'green'])
axes[0].imshow(src_mask, cmap=cmap, interpolation='nearest')
cmap = ListedColormap(['white', 'orange'])
axes[1].imshow(tgt_mask, cmap=cmap, interpolation='nearest')
if av_index is not None:
rect = Rectangle((-0.5, av_index - 0.5), src_mask.shape[1], 1, edgecolor='red', facecolor='none', lw=2)
axes[0].add_patch(rect)
rect = Rectangle((-0.5, av_index - 0.5), tgt_mask.shape[1], 1, edgecolor='red', facecolor='none', lw=2)
axes[1].add_patch(rect)
if is_bos is not None:
rows, cols = np.where(is_bos)
for row, col in zip(rows, cols):
rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='blue', facecolor='none', lw=1)
axes[0].add_patch(rect)
rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='blue', facecolor='none', lw=1)
axes[1].add_patch(rect)
for ax in axes:
ax.set_xticks(range(src_mask.shape[1] + 1), minor=False)
ax.set_yticks(range(src_mask.shape[0] + 1), minor=False)
ax.grid(which='major', color='gray', linestyle='--', linewidth=0.5)
plt.savefig(os.path.join(save_path, f'map{suffix}.png'), dpi=300, bbox_inches='tight')
plt.close()
@safe_run
def plot_prob_seed(scenario_id, prob, save_path, prefix='', indices=None):
os.makedirs(save_path, exist_ok=True)
plt.figure(figsize=(8, 5))
plt.imshow(prob, cmap='viridis', aspect='auto')
plt.colorbar()
plt.title('Seed Probability')
if indices is not None:
for col in range(indices.shape[1]):
for row in indices[:, col]:
if row == -1: continue
rect = Rectangle((col - 0.5, row - 0.5), 1, 1, edgecolor='red', facecolor='none', lw=2)
plt.gca().add_patch(rect)
plt.tight_layout()
plt.savefig(os.path.join(save_path, f'{prefix}{scenario_id}_prob_seed.png'), dpi=300, bbox_inches='tight')
plt.close()
@safe_run
def plot_raw():
plt.figure(figsize=(30, 30))
plt.rcParams['axes.facecolor']='white'
data_path = '/u/xiuyu/work/dev4/data/waymo/scenario/training'
os.makedirs("data/vis/raw/0/", exist_ok=True)
file_list = os.listdir(data_path)
for cnt_file, file in enumerate(file_list):
file_path = os.path.join(data_path, file)
dataset = tf.data.TFRecordDataset(file_path, compression_type='')
for scenario_idx, data in enumerate(dataset):
scenario = scenario_pb2.Scenario()
scenario.ParseFromString(bytearray(data.numpy()))
tqdm.write(f"scenario id: {scenario.scenario_id}")
# draw maps
for i in range(len(scenario.map_features)):
# draw lanes
if str(scenario.map_features[i].lane) != '':
line_x = [z.x for z in scenario.map_features[i].lane.polyline]
line_y = [z.y for z in scenario.map_features[i].lane.polyline]
plt.scatter(line_x, line_y, c='g', s=5)
plt.text(line_x[0], line_y[0], str(scenario.map_features[i].id), fontdict={'family': 'serif', 'size': 20, 'color': 'green'})
# draw road_edge
if str(scenario.map_features[i].road_edge) != '':
road_edge_x = [polyline.x for polyline in scenario.map_features[i].road_edge.polyline]
road_edge_y = [polyline.y for polyline in scenario.map_features[i].road_edge.polyline]
plt.scatter(road_edge_x, road_edge_y)
plt.text(road_edge_x[0], road_edge_y[0], scenario.map_features[i].road_edge.type, fontdict={'family': 'serif', 'size': 20, 'color': 'black'})
if scenario.map_features[i].road_edge.type == 2:
plt.scatter(road_edge_x, road_edge_y, c='k')
elif scenario.map_features[i].road_edge.type == 3:
plt.scatter(road_edge_x, road_edge_y, c='purple')
print(scenario.map_features[i].road_edge)
else:
plt.scatter(road_edge_x, road_edge_y, c='k')
# draw road_line
if str(scenario.map_features[i].road_line) != '':
road_line_x = [j.x for j in scenario.map_features[i].road_line.polyline]
road_line_y = [j.y for j in scenario.map_features[i].road_line.polyline]
if scenario.map_features[i].road_line.type == 7:
plt.plot(road_line_x, road_line_y, c='y')
elif scenario.map_features[i].road_line.type == 8:
plt.plot(road_line_x, road_line_y, c='y')
elif scenario.map_features[i].road_line.type == 6:
plt.plot(road_line_x, road_line_y, c='y')
elif scenario.map_features[i].road_line.type == 1:
for i in range(int(len(road_line_x) / 7)):
plt.plot(road_line_x[i * 7 : 5 + i * 7], road_line_y[i * 7 : 5 + i * 7], color='w')
elif scenario.map_features[i].road_line.type == 2:
plt.plot(road_line_x, road_line_y, c='w')
else:
plt.plot(road_line_x, road_line_y, c='w')
# draw tracks
scenario_has_invalid_tracks = False
for i in range(len(scenario.tracks)):
traj_x = [center.center_x for center in scenario.tracks[i].states]
traj_y = [center.center_y for center in scenario.tracks[i].states]
head = [center.heading for center in scenario.tracks[i].states]
valid = [center.valid for center in scenario.tracks[i].states]
print(valid)
if i == scenario.sdc_track_index:
plt.scatter(traj_x[0], traj_y[0], s=140, c='r', marker='s')
plt.scatter([x for x, v in zip(traj_x, valid) if v],
[y for y, v in zip(traj_y, valid) if v], s=14, c='r')
plt.scatter([x for x, v in zip(traj_x, valid) if not v],
[y for y, v in zip(traj_y, valid) if not v], s=14, c='m')
else:
plt.scatter(traj_x[0], traj_y[0], s=140, c='k', marker='s')
plt.scatter([x for x, v in zip(traj_x, valid) if v],
[y for y, v in zip(traj_y, valid) if v], s=14, c='b')
plt.scatter([x for x, v in zip(traj_x, valid) if not v],
[y for y, v in zip(traj_y, valid) if not v], s=14, c='m')
if valid.count(False) > 0:
scenario_has_invalid_tracks = True
if scenario_has_invalid_tracks:
plt.savefig(f"scenario_{scenario_idx}_{scenario.scenario_id}.png")
plt.clf()
breakpoint()
break
colors = [
('#1f77b4', '#1a5a8a'), # blue
('#2ca02c', '#217721'), # green
('#ff7f0e', '#cc660b'), # orange
('#9467bd', '#6f4a91'), # purple
('#d62728', '#a31d1d'), # red
('#000000', '#000000'), # black
]
@safe_run
def plot_gif():
data_path = "/u/xiuyu/work/dev4/data/waymo_processed/training"
os.makedirs("data/vis/processed/0/gif", exist_ok=True)
file_list = os.listdir(data_path)
for scenario_idx, file in tqdm(enumerate(file_list), leave=False, desc="Scenario"):
fig, ax = plt.subplots()
ax.set_axis_off()
file_path = os.path.join(data_path, file)
data = pickle.load(open(file_path, "rb"))
scenario_id = data['scenario_id']
save_path = os.path.join("data/vis/processed/0/gif",
f"scenario_{scenario_idx}_{scenario_id}.gif")
if os.path.exists(save_path):
tqdm.write(f"Skipped {save_path}.")
continue
# draw maps
ax.scatter(data['map_point']['position'][:, 0],
data['map_point']['position'][:, 1], s=0.2, c='black', edgecolors='none')
# draw agents
agent_data = data['agent']
av_index = agent_data['av_index']
position = agent_data['position'] # (num_agent, 91, 3)
heading = agent_data['heading'] # (num_agent, 91)
shape = agent_data['shape'] # (num_agent, 91, 3)
category = agent_data['category'] # (num_agent,)
valid_mask = (position[..., 0] != 0) & (position[..., 1] != 0) # (num_agent, 91)
num_agent = valid_mask.shape[0]
num_timestep = position.shape[1]
is_av = np.arange(num_agent) == int(av_index)
is_blue = valid_mask.sum(axis=1) == num_timestep
is_green = ~valid_mask[:, 0] & valid_mask[:, -1]
is_orange = valid_mask[:, 0] & ~valid_mask[:, -1]
is_purple = (valid_mask.sum(axis=1) != num_timestep
) & (~is_green) & (~is_orange)
agent_colors = np.zeros((num_agent,))
agent_colors[is_blue] = 1
agent_colors[is_green] = 2
agent_colors[is_orange] = 3
agent_colors[is_purple] = 4
agent_colors[is_av] = 5
veh_mask = category == 1
ped_mask = category == 2
cyc_mask = category == 3
shape[veh_mask, :, 1] = 1.8
shape[veh_mask, :, 0] = 1.8
shape[ped_mask, :, 1] = 0.5
shape[ped_mask, :, 0] = 0.5
shape[cyc_mask, :, 1] = 1.0
shape[cyc_mask, :, 0] = 1.0
fig_paths = []
for tid in tqdm(range(num_timestep), leave=False, desc="Timestep"):
current_valid_mask = valid_mask[:, tid]
xs = position[current_valid_mask, tid, 0]
ys = position[current_valid_mask, tid, 1]
widths = shape[current_valid_mask, tid, 1]
lengths = shape[current_valid_mask, tid, 0]
angles = heading[current_valid_mask, tid]
current_agent_colors = agent_colors[current_valid_mask]
drawn_agents = []
contours = cal_polygon_contour(xs, ys, angles, widths, lengths) # (num_agent, 4, 2)
contours = np.concatenate([contours, contours[:, 0:1]], axis=1) # (num_agent, 5, 2)
for x, y, width, length, angle, color_type in zip(
xs, ys, widths, lengths, angles, current_agent_colors):
agent = plt.Rectangle((x, y), width, length, angle=((angle + np.pi / 2) / np.pi * 360) % 360,
linewidth=0.2,
facecolor=colors[int(color_type) - 1][0],
edgecolor=colors[int(color_type) - 1][1])
ax.add_patch(agent)
drawn_agents.append(agent)
plt.gca().set_aspect('equal', adjustable='box')
# for contour, color_type in zip(contours, agent_colors):
# drawn_agent = ax.plot(contour[:, 0], contour[:, 1])
# drawn_agents.append(drawn_agent)
fig_path = os.path.join("data/vis/processed/0/",
f"scenario_{scenario_idx}_{scenario_id}_{tid}.png")
plt.savefig(fig_path, dpi=600)
fig_paths.append(fig_path)
for drawn_agent in drawn_agents:
drawn_agent.remove()
plt.close()
# generate gif
import imageio.v2 as imageio
images = []
for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
images.append(imageio.imread(fig_path))
imageio.mimsave(save_path, images, duration=0.1)
@safe_run
def plot_map_token(ax: Axes, map_points: npt.NDArray, token_pos: npt.NDArray, token_heading: npt.NDArray, colors: Union[str, npt.NDArray]=None):
plot_map(ax, map_points)
x, y = token_pos[:, 0], token_pos[:, 1]
u = np.cos(token_heading)
v = np.sin(token_heading)
if colors is None:
colors = np.random.rand(x.shape[0], 3)
ax.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=0.2, color=colors, width=0.005,
headwidth=0.2, headlength=2)
ax.scatter(x, y, color='blue', s=0.2, edgecolors='none')
ax.axis("equal")
@safe_run
def plot_map(ax: Axes, map_points: npt.NDArray, color='black'):
ax.scatter(map_points[:, 0], map_points[:, 1], s=0.5, c=color, edgecolors='none')
xmin = np.min(map_points[:, 0])
xmax = np.max(map_points[:, 0])
ymin = np.min(map_points[:, 1])
ymax = np.max(map_points[:, 1])
ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)
@safe_run
def plot_agent(ax: Axes, xy: Sequence[float], heading: float, type: str, state, is_av: bool=False,
pl2seed_radius: float=25., attr_tokenizer: Attr_Tokenizer=None, enter_index: list=[], **kwargs):
if type == 'veh':
length = 4.3
width = 1.8
size = 1.0
elif type == 'ped':
length = 0.5
width = 0.5
size = 0.1
elif type == 'cyc':
length = 1.9
width = 0.5
size = 0.3
else:
raise ValueError(f"Unsupported agent type {type}")
if kwargs.get('label', None) is not None:
ax.text(
xy[0] + 1.5, xy[1] + 1.5,
kwargs['label'], fontsize=2, color="darkred", ha="center", va="center"
)
patch = FancyBboxPatch([-length / 2, -width / 2], length, width, linewidth=.2, **kwargs)
transform = (
mtransforms.Affine2D().rotate(heading).translate(xy[0], xy[1])
+ ax.transData
)
patch.set_transform(transform)
kwargs['label'] = None
angles = [0, 2 * np.pi / 3, np.pi, 4 * np.pi / 3]
pts = np.stack([size * np.cos(angles), size * np.sin(angles)], axis=-1)
center_patch = Polygon(pts, zorder=10., linewidth=.2, **kwargs)
center_patch.set_transform(transform)
ax.add_patch(patch)
ax.add_patch(center_patch)
if is_av:
if attr_tokenizer is not None:
circle_patch = Circle(
(xy[0], xy[1]), pl2seed_radius, linewidth=0.5, edgecolor='gray', linestyle='--', facecolor='none'
)
ax.add_patch(circle_patch)
grid = attr_tokenizer.get_grid(torch.tensor(np.array(xy)).float(),
torch.tensor(np.array([heading])).float()).numpy()[0] # (num_grid, 2)
ax.scatter(grid[:, 0], grid[:, 1], s=0.3, c='blue', edgecolors='none')
ax.text(grid[0, 0], grid[0, 1], 'Front', fontsize=2, color='darkred', ha='center', va='center')
ax.text(grid[-1, 0], grid[-1, 1], 'Back', fontsize=2, color='darkred', ha='center', va='center')
if enter_index:
for i in enter_index:
ax.plot(grid[int(i), 0], grid[int(i), 1], marker='x', color='red', markersize=1)
return patch, center_patch
@safe_run
def plot_all(map, xs, ys, angles, types, colors, is_avs, pl2seed_radius: float=25.,
attr_tokenizer: Attr_Tokenizer=None, enter_index: list=[], labels: list=[], **kwargs):
plt.subplots_adjust(left=0.3, right=0.7, top=0.7, bottom=0.3)
_, ax = plt.subplots()
ax.set_axis_off()
plot_map(ax, map)
if not labels:
labels = [None] * xs.shape[0]
for x, y, angle, type, color, label, is_av in zip(xs, ys, angles, types, colors, labels, is_avs):
assert type in ('veh', 'ped', 'cyc'), f"Unsupported type {type}."
plot_agent(ax, [x, y], angle.item(), type, None, is_av, facecolor=color, edgecolor='k', label=label,
pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, enter_index=enter_index)
ax.grid(False)
ax.set_aspect('equal', adjustable='box')
# ! set plot limit if need
if kwargs.get('limit_size', None):
cx = float(xs[is_avs])
cy = float(ys[is_avs])
lx, ly = kwargs['limit_size']
xmin, xmax = cx - lx, cx + lx
ymin, ymax = cy - ly, cy + ly
ax.set_xlim(xmin, xmax)
ax.set_ylim(ymin, ymax)
# ax.legend(loc='best', frameon=True)
pil_image = None
if kwargs.get('save_path', None):
plt.savefig(kwargs['save_path'], dpi=600, bbox_inches="tight")
else:
# ！convert to PIL image
buf = io.BytesIO()
plt.savefig(buf, format='png', dpi=600, bbox_inches='tight')
buf.seek(0)
pil_image = Image.open(buf).convert('RGB')
plt.close()
return pil_image
@safe_run
def plot_file(gt_folder: str,
folder: Optional[str] = None,
files: Optional[str] = None,
save_gif: bool = True,
batch_idx: Optional[int] = None,
time_idx: Optional[List[int]] = None,
limit_size: Optional[List[int]] = None,
**kwargs,
) -> List[Image.Image]:
from dev.metrics.compute_metrics import _unbatch
shift = 5
if files is None:
assert os.path.exists(folder), f'Path {folder} does not exist.'
files = list(fnmatch.filter(os.listdir(folder), 'idx_*_rollouts.pkl'))
CONSOLE.log(f'Found {len(files)} rollouts files from {folder}.')
if folder is None:
assert os.path.exists(files), f'Path {files} does not exist.'
folder = os.path.dirname(files)
files = [files]
parent, folder_name = os.path.split(folder.rstrip(os.sep))
if save_gif:
save_path = os.path.join(parent, f'{folder_name}_plots')
else:
save_path = None
file_outs = []
for file in (pbar := tqdm(files, leave=False, desc='Plotting files ...')):
pbar.set_postfix(file=file)
with open(os.path.join(folder, file), 'rb') as f:
preds = pickle.load(f)
scenario_ids = preds['_scenario_id']
agent_batch = preds['agent_batch']
agent_id = _unbatch(preds['agent_id'], agent_batch)
preds_traj = _unbatch(preds['pred_traj'], agent_batch)
preds_head = _unbatch(preds['pred_head'], agent_batch)
preds_type = _unbatch(preds['pred_type'], agent_batch)
if 'pred_state' in preds:
preds_state = _unbatch(preds['pred_state'], agent_batch)
else:
preds_state = tuple([torch.ones((*traj.shape[:2], traj.shape[2] // shift)) for traj in preds_traj]) # [n_agent, n_rollout, n_step2Hz]
preds_valid = _unbatch(preds['pred_valid'], agent_batch)
# ! fetch certain scenario
if batch_idx is not None:
scenario_ids = scenario_ids[batch_idx : batch_idx + 1]
agent_id = (agent_id[batch_idx],)
preds_traj = (preds_traj[batch_idx],)
preds_head = (preds_head[batch_idx],)
preds_type = (preds_type[batch_idx],)
preds_state = (preds_state[batch_idx],)
preds_valid = (preds_valid[batch_idx],)
scenario_outs = []
for i, scenario_id in enumerate(scenario_ids):
n_agent, n_rollouts = preds_traj[0].shape[:2]
rollout_outs = []
for j in range(n_rollouts): # 1
pred = dict(scenario_id=[scenario_id],
pred_traj=preds_traj[i][:, j],
pred_head=preds_head[i][:, j],
pred_state=(
torch.cat([torch.zeros(n_agent, 1), preds_state[i][:, j].repeat_interleave(repeats=shift, dim=-1)],
dim=1)
),
pred_type=preds_type[i][:, j],
)
# NOTE: hard code!!!
if 'av_id' in preds:
av_index = agent_id[i][:, 0].tolist().index(preds['av_id'])
else:
av_index = n_agent - 1
# ! load logged data
data_path = os.path.join(gt_folder, 'validation', f'{scenario_id}.pkl')
with open(data_path, 'rb') as f:
data = pickle.load(f)
rollout_outs.append(
plot_val(data, pred,
av_index=av_index,
save_path=save_path,
save_gif=save_gif,
time_idx=time_idx,
limit_size=limit_size,
**kwargs
)
)
scenario_outs.append(rollout_outs)
file_outs.append(scenario_outs)
return file_outs
@safe_run
def plot_val(data: Union[dict, str], pred: dict, av_index: int, save_path: str, suffix: str='',
pl2seed_radius: float=75., attr_tokenizer=None, **kwargs):
if isinstance(data, str):
assert data.endswith('.pkl'), f'Got invalid data path {data}.'
assert os.path.exists(data), f'Path {data} does not exist.'
with open(data, 'rb') as f:
data = pickle.load(f)
map_point = data['map_point']['position'].cpu().numpy()
scenario_id = pred['scenario_id'][0]
pred_traj = pred['pred_traj'].cpu().numpy() # (num_agent, num_future_step, 2)
pred_type = list(map(lambda i: AGENT_TYPE[i], pred['pred_type'].tolist()))
pred_state = pred['pred_state'].cpu().numpy()
pred_head = pred['pred_head'].cpu().numpy()
ids = np.arange(pred_traj.shape[0])
if 'agent_labels' in pred:
kwargs.update(agent_labels=pred['agent_labels'])
return plot_scenario(scenario_id, map_point, pred_traj, pred_head, pred_state, pred_type,
av_index=av_index, ids=ids, save_path=save_path, suffix=suffix,
pl2seed_radius=pl2seed_radius, attr_tokenizer=attr_tokenizer, **kwargs)
@safe_run
def plot_scenario(scenario_id: str,
map_data: npt.NDArray,
traj: npt.NDArray,
heading: npt.NDArray,
state: npt.NDArray,
types: List[str],
av_index: int,
color_type: Literal['state', 'type', 'seed', 'insert']='seed',
state_type: List[str]=['invalid', 'valid', 'enter', 'exit'],
plot_enter: bool=False,
suffix: str='',
pl2seed_radius: float=25.,
attr_tokenizer: Attr_Tokenizer=None,
enter_index: List[list] = [],
save_gif: bool=True,
tokenized: bool=False,
agent_labels: List[List[Optional[str]]] = [],
**kwargs):
num_historical_steps = 11
shift = 5
num_agent, num_timestep = traj.shape[:2]
if tokenized:
num_historical_steps = 2
shift = 1
if (
'save_path' in kwargs
and kwargs['save_path'] != ''
and kwargs['save_path'] != None
):
os.makedirs(kwargs['save_path'], exist_ok=True)
save_id = int(max([0] + list(map(lambda fname: int(fname.split("_")[-1]),
filter(lambda fname: fname.startswith(scenario_id)
and os.path.isdir(os.path.join(kwargs['save_path'], fname)),
os.listdir(kwargs['save_path'])))))) + 1
os.makedirs(f"{kwargs['save_path']}/{scenario_id}_{str(save_id).zfill(3)}", exist_ok=True)
if save_id > 1:
try:
import shutil
shutil.rmtree(f"{kwargs['save_path']}/{scenario_id}_{str(save_id - 1).zfill(3)}")
except:
pass
visible_mask = state != state_type.index('invalid')
if not plot_enter:
visible_mask &= (state != state_type.index('enter'))
last_valid_step = visible_mask.shape[1] - 1 - torch.argmax(torch.Tensor(visible_mask).flip(dims=[1]).long(), dim=1)
ids = None
if 'ids' in kwargs:
ids = kwargs['ids']
last_valid_step = {int(ids[i]): int(last_valid_step[i]) for i in range(len(ids))}
# agent colors
agent_colors = np.zeros((num_agent, num_timestep, 3))
agent_palette = sns.color_palette('husl', n_colors=7)
state_colors = {state: np.array(agent_palette[i]) for i, state in enumerate(state_type)}
seed_colors = {seed: np.array(agent_palette[i]) for i, seed in enumerate(['existing', 'entered', 'exited'])}
if color_type == 'state':
for t in range(state.shape[1]):
agent_colors[state[:, t] == state_type.index('invalid'), t * shift : (t + 1) * shift] = state_colors['invalid']
agent_colors[state[:, t] == state_type.index('valid'), t * shift : (t + 1) * shift] = state_colors['valid']
agent_colors[state[:, t] == state_type.index('enter'), t * shift : (t + 1) * shift] = state_colors['enter']
agent_colors[state[:, t] == state_type.index('exit'), t * shift : (t + 1) * shift] = state_colors['exit']
if color_type == 'seed':
agent_colors[:, :] = seed_colors['existing']
is_exited = np.any(state[:, num_historical_steps - 1:] == state_type.index('exit'), axis=-1)
is_entered = np.any(state[:, num_historical_steps - 1:] == state_type.index('enter'), axis=-1)
is_entered[av_index + 1:] = True # NOTE: hard code, need improvment
agent_colors[is_exited, :] = seed_colors['exited']
agent_colors[is_entered, :] = seed_colors['entered']
if color_type == 'insert':
agent_colors[:, :] = seed_colors['exited']
agent_colors[av_index + 1:] = seed_colors['existing']
agent_colors[av_index, :] = np.array(agent_palette[-1])
is_av = np.zeros_like(state[:, 0]).astype(np.bool_)
is_av[av_index] = True
# ! get timesteps to plot
timesteps = list(range(num_timestep))
if kwargs.get('time_idx', None) is not None:
time_idx = kwargs['time_idx']
assert set(time_idx).issubset(set(timesteps)), f'Got invalid time_idx: {time_idx=} v.s. {timesteps=}'
timesteps = sorted(time_idx)
# ! get plot limits
limit_size = kwargs.get('limit_size', None)
if limit_size is not None:
assert len(limit_size) == 2, f'Got invalid `limit_size`: {limit_size=}'
# ! plot all
pil_images = []
fig_paths = []
for tid in tqdm(timesteps, leave=False, desc="Plot ..."):
mask_t = visible_mask[:, tid]
xs = traj[mask_t, tid, 0]
ys = traj[mask_t, tid, 1]
angles = heading[mask_t, tid]
colors = agent_colors[mask_t, tid]
types_t = [types[i] for i, mask in enumerate(mask_t) if mask]
if ids is not None:
ids_t = ids[mask_t]
is_av_t = is_av[mask_t]
enter_index_t = enter_index[tid] if enter_index else None
labels = []
if agent_labels:
labels = [agent_labels[i][tid // shift] for i in range(len(agent_labels)) if mask_t[i]]
fig_path = None
if kwargs.get('save_path', None) is not None:
save_path = kwargs['save_path']
fig_path = os.path.join(f"{save_path}/{scenario_id}_{str(save_id).zfill(3)}", f"{tid}.png")
fig_paths.append(fig_path)
pil_images.append(
plot_all(map_data, xs, ys, angles, types_t,
colors=colors,
save_path=fig_path,
is_avs=is_av_t,
pl2seed_radius=pl2seed_radius,
attr_tokenizer=attr_tokenizer,
enter_index=enter_index_t,
labels=labels,
limit_size=limit_size,
)
)
# generate gif
if fig_paths and save_gif:
os.makedirs(os.path.join(save_path, 'gifs'), exist_ok=True)
images = []
gif_path = f"{save_path}/gifs/{scenario_id}_{str(save_id).zfill(3)}.gif"
for fig_path in tqdm(fig_paths, leave=False, desc="Generate gif ..."):
images.append(Image.open(fig_path))
try:
images[0].save(gif_path, save_all=True, append_images=images[1:], duration=100, loop=0)
tqdm.write(f"Saved gif at {gif_path}")
try:
import shutil
shutil.rmtree(f"{save_path}/{scenario_id}_{str(save_id).zfill(3)}")
os.remove(f"{save_path}/gifs/{scenario_id}_{str(save_id - 1).zfill(3)}.gif")
except:
pass
except Exception as e:
tqdm.write(f"{e}! Failed to save gif at {gif_path}")
return pil_images
def match_token_map(data):
# init map token
argmin_sample_len = 3
map_token_traj_path = '/u/xiuyu/work/dev4/dev/tokens/map_traj_token5.pkl'
map_token_traj = pickle.load(open(map_token_traj_path, 'rb'))
map_token = {'traj_src': map_token_traj['traj_src'], }
traj_end_theta = np.arctan2(map_token['traj_src'][:, -1, 1] - map_token['traj_src'][:, -2, 1],
map_token['traj_src'][:, -1, 0] - map_token['traj_src'][:, -2, 0])
indices = torch.linspace(0, map_token['traj_src'].shape[1]-1, steps=argmin_sample_len).long()
map_token['sample_pt'] = torch.from_numpy(map_token['traj_src'][:, indices]).to(torch.float)
map_token['traj_end_theta'] = torch.from_numpy(traj_end_theta).to(torch.float)
map_token['traj_src'] = torch.from_numpy(map_token['traj_src']).to(torch.float)
traj_pos = data['map_save']['traj_pos'].to(torch.float)
traj_theta = data['map_save']['traj_theta'].to(torch.float)
pl_idx_list = data['map_save']['pl_idx_list']
token_sample_pt = map_token['sample_pt'].to(traj_pos.device)
token_src = map_token['traj_src'].to(traj_pos.device)
max_traj_len = map_token['traj_src'].shape[1]
pl_num = traj_pos.shape[0]
pt_token_pos = traj_pos[:, 0, :].clone()
pt_token_orientation = traj_theta.clone()
cos, sin = traj_theta.cos(), traj_theta.sin()
rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
rot_mat[..., 0, 0] = cos
rot_mat[..., 0, 1] = -sin
rot_mat[..., 1, 0] = sin
rot_mat[..., 1, 1] = cos
traj_pos_local = torch.bmm((traj_pos - traj_pos[:, 0:1]), rot_mat.view(-1, 2, 2))
distance = torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1))
pt_token_id = torch.argmin(distance, dim=1)
noise = False
if noise:
topk_indices = torch.argsort(torch.sum((token_sample_pt[None] - traj_pos_local.unsqueeze(1)) ** 2, dim=(-2, -1)), dim=1)[:, :8]
sample_topk = torch.randint(0, topk_indices.shape[-1], size=(topk_indices.shape[0], 1), device=topk_indices.device)
pt_token_id = torch.gather(topk_indices, 1, sample_topk).squeeze(-1)
# cos, sin = traj_theta.cos(), traj_theta.sin()
# rot_mat = traj_theta.new_zeros(pl_num, 2, 2)
# rot_mat[..., 0, 0] = cos
# rot_mat[..., 0, 1] = sin
# rot_mat[..., 1, 0] = -sin
# rot_mat[..., 1, 1] = cos
# token_src_world = torch.bmm(token_src[None, ...].repeat(pl_num, 1, 1, 1).reshape(pl_num, -1, 2),
# rot_mat.view(-1, 2, 2)).reshape(pl_num, token_src.shape[0], max_traj_len, 2) + traj_pos[:, None, [0], :]
# token_src_world_select = token_src_world.view(-1, 1024, 11, 2)[torch.arange(pt_token_id.view(-1).shape[0]), pt_token_id.view(-1)].view(pl_num, max_traj_len, 2)
pl_idx_full = pl_idx_list.clone()
token2pl = torch.stack([torch.arange(len(pl_idx_list), device=traj_pos.device), pl_idx_full.long()])
count_nums = []
for pl in pl_idx_full.unique():
pt = token2pl[0, token2pl[1, :] == pl]
left_side = (data['pt_token']['side'][pt] == 0).sum()
right_side = (data['pt_token']['side'][pt] == 1).sum()
center_side = (data['pt_token']['side'][pt] == 2).sum()
count_nums.append(torch.Tensor([left_side, right_side, center_side]))
count_nums = torch.stack(count_nums, dim=0)
num_polyline = int(count_nums.max().item())
traj_mask = torch.zeros((int(len(pl_idx_full.unique())), 3, num_polyline), dtype=bool)
idx_matrix = torch.arange(traj_mask.size(2)).unsqueeze(0).unsqueeze(0)
idx_matrix = idx_matrix.expand(traj_mask.size(0), traj_mask.size(1), -1)
counts_num_expanded = count_nums.unsqueeze(-1)
mask_update = idx_matrix < counts_num_expanded
traj_mask[mask_update] = True
data['pt_token']['traj_mask'] = traj_mask
data['pt_token']['position'] = torch.cat([pt_token_pos, torch.zeros((data['pt_token']['num_nodes'], 1),
device=traj_pos.device, dtype=torch.float)], dim=-1)
data['pt_token']['orientation'] = pt_token_orientation
data['pt_token']['height'] = data['pt_token']['position'][:, -1]
data[('pt_token', 'to', 'map_polygon')] = {}
data[('pt_token', 'to', 'map_polygon')]['edge_index'] = token2pl # (2, num_points)
data['pt_token']['token_idx'] = pt_token_id
return data
@safe_run
def plot_tokenize(data, save_path: str):
shift = 5
token_size = 2048
pl2seed_radius = 75
# transformation
transform = WaymoTargetBuilder(num_historical_steps=11,
num_future_steps=80,
max_num=32,
training=False)
grid_range = 150.
grid_interval = 3.
angle_interval = 3.
attr_tokenizer = Attr_Tokenizer(grid_range=grid_range,
grid_interval=grid_interval,
radius=pl2seed_radius,
angle_interval=angle_interval)
# tokenization
token_processor = TokenProcessor(token_size,
training=False,
predict_motion=True,
predict_state=True,
predict_map=True,
state_token={'invalid': 0, 'valid': 1, 'enter': 2, 'exit': 3},
pl2seed_radius=pl2seed_radius)
CONSOLE.log(f"Loaded token processor with token_size: {token_size}")
# preprocess
data: HeteroData = transform(data)
tokenized_data = token_processor(data)
CONSOLE.log(f"Keys in tokenized data:\n{tokenized_data.keys()}")
# plot
agent_data = tokenized_data['agent']
map_data = tokenized_data['map_point']
# CONSOLE.log(f"Keys in agent data:\n{agent_data.keys()}")
av_index = agent_data['av_index']
raw_traj = agent_data['position'][..., :2].contiguous() # [n_agent, n_step, 2]
raw_heading = agent_data['heading'] # [n_agent, n_step]
traj = agent_data['traj_pos'][..., :2].contiguous() # [n_agent, n_step, 6, 2]
traj = traj[:, :, 1:, :].flatten(1, 2)
traj = torch.cat([raw_traj[:, :1], traj], dim=1)
heading = agent_data['traj_heading'] # [n_agent, n_step, 6]
heading = heading[:, :, 1:].flatten(1, 2)
heading = torch.cat([raw_heading[:, :1], heading], dim=1)
agent_state = agent_data['state_idx'].repeat_interleave(repeats=shift, dim=-1)
agent_state = torch.cat([torch.zeros_like(agent_state[:, :1]), agent_state], dim=1)
agent_type = agent_data['type']
ids = np.arange(raw_traj.shape[0])
return plot_scenario(
scenario_id=tokenized_data['scenario_id'],
map_data=tokenized_data['map_point']['position'].numpy(),
traj=raw_traj.numpy(),
heading=raw_heading.numpy(),
state=agent_state.numpy(),
types=list(map(lambda i: AGENT_TYPE[i], agent_type.tolist())),
av_index=av_index,
ids=ids,
save_path=save_path,
pl2seed_radius=pl2seed_radius,
attr_tokenizer=attr_tokenizer,
color_type='state',
)
def get_metainfos(folder: str):
import pandas as pd
assert os.path.exists(folder), f'Path {folder} does not exist.'
files = list(fnmatch.filter(os.listdir(folder), 'idx_*_rollouts.pkl'))
CONSOLE.log(f'Found {len(files)} rollouts files from {folder}.')
metainfos_path = f'{os.path.normpath(folder)}_metainfos.parquet'
csv_path = f'{os.path.normpath(folder)}_metainfos.csv'
if not os.path.exists(metainfos_path):
data = []
for file in tqdm(files):
pkl_data = pickle.load(open(os.path.join(folder, file), 'rb'))
data.extend((file, scenario_id, index) for index, scenario_id in enumerate(pkl_data['_scenario_id']))
df = pd.DataFrame(data, columns=('rollout_file', 'scenario_id', 'index'))
df.to_parquet(metainfos_path)
df.to_csv(csv_path)
CONSOLE.log(f'Successfully saved to {metainfos_path}.')
else:
CONSOLE.log(f'File {metainfos_path} already exists!')
return
def plot_comparison(methods: List[str], rollouts_paths: List[str], gt_folders: List[str],
save_path: str, scenario_ids: Optional[List[str]] = None):
import pandas as pd
from collections import defaultdict
# ! hyperparameter
fps = 10
plot_time = [1, 6, 12, 18, 24, 30]
# plot_time = [1, 5, 10, 15, 20, 25]
time_idx = [int(time * fps) for time in plot_time]
limit_size = [75, 60] # [width, height]
# ! load metainfos
metainfos = defaultdict(dict)
for method, rollout_path in zip(methods, rollouts_paths):
meta_info_path = f'{os.path.normpath(rollout_path)}_metainfos.parquet'
metainfos[method]['df'] = pd.read_parquet(meta_info_path)
CONSOLE.log(f'Loaded {method=} with {len(metainfos[method]["df"]["scenario_id"])=}.')
common_scenarios = set(metainfos['ours']['df']['scenario_id'])
for method, meta_info in metainfos.items():
if method == 'ours':
continue
common_scenarios &= set(meta_info['df']['scenario_id'])
for method, meta_info in metainfos.items():
df = metainfos[method]['df']
metainfos[method]['df'] = df[df['scenario_id'].isin(common_scenarios)]
CONSOLE.log(f'Filter and get {len(common_scenarios)=}.')
# ! load data and plot
if scenario_ids is None:
scenario_ids = metainfos['ours']['df']['scenario_id'].tolist()
CONSOLE.log(f'Plotting {len(scenario_ids)=} ...')
for scenario_id in (pbar := tqdm(scenario_ids)):
pbar.set_postfix(scenario_id=scenario_id)
figures = dict()
for method, rollout_path, gt_folder in zip(methods, rollouts_paths, gt_folders):
df = metainfos[method]['df']
_df = df.loc[df['scenario_id'] == scenario_id]
batch_idx = int(_df['index'].tolist()[0])
rollout_file = _df['rollout_file'].tolist()[0]
figures[method] = plot_file(
gt_folder=gt_folder,
files=os.path.join(rollout_path, rollout_file),
save_gif=False,
batch_idx=batch_idx,
time_idx=time_idx,
limit_size=limit_size,
color_type='insert',
)[0][0][0]
# ! plot figures
border = 5
padding_x = 20
padding_y = 50
img_width, img_height = figures['ours'][0].size
img_width = img_width + 2 * border
img_height = img_height + 2 * border
n_col = len(time_idx)
n_row = len(methods)
W = n_col * img_width + (n_col - 1) * padding_x
H = n_row * img_height + (n_row - 1) * padding_y
canvas = Image.new('RGB', (W, H), 'white')
for i_row, (method, method_figures) in enumerate(figures.items()):
for i_col, method_figure in enumerate(method_figures):
x = i_col * (img_width + padding_x)
y = i_row * (img_height + padding_y)
padded_figure = Image.new('RGB', (img_width, img_height), 'black')
padded_figure.paste(method_figure, (border, border))
canvas.paste(padded_figure, (x, y))
canvas.save(
os.path.join(save_path, f'{scenario_id}.png')
)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument('--data_path', type=str, default='/u/xiuyu/work/dev4/data/waymo_processed')
parser.add_argument('--tfrecord_dir', type=str, default='validation_tfrecords_splitted')
# plot tokenized data
parser.add_argument('--save_folder', type=str, default='plot_gt')
parser.add_argument('--split', type=str, default='validation')
parser.add_argument('--scenario_id', type=str, default=None)
parser.add_argument('--plot_tokenize', action='store_true')
# plot generated rollouts
parser.add_argument('--plot_file', action='store_true')
parser.add_argument('--folder_path', type=str, default=None)
parser.add_argument('--file_path', type=str, default=None)
# metainfos
parser.add_argument('--get_metainfos', action='store_true')
# plot comparison
parser.add_argument('--plot_comparison', action='store_true')
parser.add_argument('--comparison_folder', type=str, default='comparisons')
args = parser.parse_args()
if args.plot_tokenize:
scenario_id = "74ad7b76d5906d39"
# scenario_id = "1d60300bc06f4801"
data_path = os.path.join(args.data_path, args.split, f"{scenario_id}.pkl")
data = pickle.load(open(data_path, "rb"))
data['tfrecord_path'] = os.path.join(args.tfrecord_dir, f'{scenario_id}.tfrecords')
CONSOLE.log(f"Loaded scenario {scenario_id}")
save_path = os.path.join(args.data_path, args.save_folder, args.split)
os.makedirs(save_path, exist_ok=True)
plot_tokenize(data, save_path)
if args.plot_file:
plot_file(args.data_path, folder=args.folder_path, files=args.file_path)
if args.get_metainfos:
assert args.folder_path is not None, f'`folder_path` should not be None!'
get_metainfos(args.folder_path)
if args.plot_comparison:
methods = ['ours', 'smart']
gt_folders = [
'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/data/processed',
'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/data/processed',
'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/data/processed',
]
rollouts_paths = [
'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/validation_ours0',
'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/validation_smart',
'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/validation_cslft',
]
save_path = f'/baai-cwm-1/baai_cwm_ml/algorithm/xiuyu.yang/work/dev4/output/scalable_smart_long/{args.comparison_folder}/'
os.makedirs(save_path, exist_ok=True)
scenario_ids = ['72ff3e1540b28431','a16c927b1a1cca74','a504d55ea6658de7','639949ea1d16125b']
plot_comparison(methods, rollouts_paths, gt_folders,
save_path=save_path,
scenario_ids=scenario_ids)