def load_raw_data(root_dir: str, save_path: str = "./save/features"):
'''
Save raw csv to np.ndarray
@input root_dir (string): root directory contains .csv data
@input save_path (string): save features to this path, default "./save/features"
'''
afl = ArgoverseForecastingLoader(root_dir)
files = os.listdir(root_dir)
for f in files:
if not f.endswith(".csv"):
continue
seq_path = os.path.join(root_dir, f)
print("Processing " + seq_path)
id_list = afl.get(seq_path).track_id_list
agent_traj = afl.get(seq_path).agent_traj
df = afl.get(seq_path).seq_df
tarj_list = []
df['TIMESTAMP'] -= df['TIMESTAMP'].min()
for id in id_list:
subdf = df.loc[df['TRACK_ID'] == id]
tarj_list.append(
subdf.drop(columns='CITY_NAME').sort_values(
by=['TIMESTAMP']).to_numpy())
with open(os.path.join(
save_path,
f[:-4] + ".save",
), "wb") as f:
pickle.dump(tarj_list, f)
def save_results_single_pred(self):
print("running save results")
afl=ArgoverseForecastingLoader("data/val/data/")
checkpoint = torch.load(self.model_dir+'best-model.pt', map_location=lambda storage, loc: storage)
# self.model.load_state_dict(torch.load(self.model_dir+'best-model.pt')['model_state_dict'])
self.model.load_state_dict(checkpoint['model_state_dict'])
self.model.eval()
save_results_path=self.model_dir+"/results/"
# pdb.set_trace()
if not os.path.exists(save_results_path):
os.mkdir(save_results_path)
num_batches=len(self.val_loader.batch_sampler)
for i_batch,traj_dict in enumerate(self.val_loader):
print(f"Running {i_batch}/{num_batches}",end="\r")
gt_traj=traj_dict['gt_unnorm_agent'].numpy()
# output=self.model(traj_dict,mode='validate')
output=self.model(traj_dict)
output=self.val_loader.dataset.inverse_transform(output,traj_dict)
output=output.detach().cpu().numpy()
seq_paths=traj_dict['seq_path']
for index,seq_path in enumerate(seq_paths):
loader=afl.get(seq_path)
input_array=loader.agent_traj[0:20,:]
city=loader.city
del loader
seq_index=int(os.path.basename(seq_path).split('.')[0])
output_dict={'seq_path':seq_path,'seq_index':seq_index,'input':input_array,
'output':output[index],'target':gt_traj[index],'city':city}
with open(f"{save_results_path}/{seq_index}.pkl", 'wb') as f:
pickle.dump(output_dict,f)
class Argoverse(object):
def __init__(self, root : str):
"""
create the dataset manager object
- params:
- root : str : root path for the data
"""
self.root_path = root
self.manager = ArgoverseForecastingLoader(self.root_path)
self.ekf = EKF()
self.feature = PlainFeatures()
self.savgol = SavitzkyGolov(window_length=11,
poly=3)
all_vel_raw = []
all_vel_ekf = []
all_vel_savgol = []
all_vel_ekf_savgol = []
all_acc_raw = []
all_acc_ekf = []
all_acc_savgol = []
all_acc_ekf_savgol = []
all_j_raw = []
all_j_ekf = []
all_j_savgol = []
all_j_ekf_savgol = []
for index in range(0, len(self.manager.seq_list)):
print(index)
seq_path = self.manager.seq_list[index]
df = self.manager.get(seq_path).seq_df
traj = df[df['OBJECT_TYPE']=='AGENT']
traj = traj[['TIMESTAMP', 'X', 'Y']].values
self.savgol.set_window_size(traj.shape[0]//2)
#features
plain_features = self.feature.process(traj=traj)
ekf_features = self.ekf.process(traj=traj)
savgol_features = self.savgol.process(traj=traj)
ekf_savgol_features = self.savgol.filter(vector=ekf_features)
plain_features = np.squeeze(plain_features)
ekf_features = np.squeeze(ekf_features)
savgol_features = np.squeeze(savgol_features)
ekf_savgol_features = np.squeeze(ekf_savgol_features)
#path
plt.figure(1)
plt.plot(plain_features[:,0],plain_features[:,1], color='blue', label='raw', marker='.')
plt.plot(ekf_features[:,0],ekf_features[:,1], color='green', label='ekf', marker='.')
plt.plot(savgol_features[:,0],savgol_features[:,1], color='orange', label='savgol', marker='.')
plt.plot(ekf_savgol_features[:,0],ekf_savgol_features[:,1], color='red', label='ekf-savgol', marker='.')
plt.legend()
plt.savefig(f'path_{index}.png')
plt.clf()
#velocity
vel_plain = np.sqrt(np.power(plain_features[:,2],2) + np.power(plain_features[:,3],2))
vel_ekf = np.sqrt(np.power(ekf_features[:,2],2) + np.power(ekf_features[:,3],2))
vel_savgol = np.sqrt(np.power(savgol_features[:,2],2) + np.power(savgol_features[:,3],2))
vel_ekf_savgol = np.sqrt(np.power(ekf_savgol_features[:,2],2) + np.power(ekf_savgol_features[:,3],2))
plt.figure(2)
plt.plot(vel_plain, color='blue', label='raw', marker='.')
plt.plot(vel_ekf, color='green', label='ekf', marker='.')
plt.plot(vel_savgol, color='orange', label='savgol', marker='.')
plt.plot(vel_ekf_savgol, color='red', label='ekf-savgol', marker='.')
plt.legend()
plt.savefig(f'vel_{index}.png')
plt.clf()
acc_plain = np.sqrt(np.power(plain_features[:,4],2) + np.power(plain_features[:,5],2))
acc_ekf = np.sqrt(np.power(ekf_features[:,4],2) + np.power(ekf_features[:,5],2))
acc_savgol = np.sqrt(np.power(savgol_features[:,4],2) + np.power(savgol_features[:,5],2))
acc_ekf_savgol = np.sqrt(np.power(ekf_savgol_features[:,4],2) + np.power(ekf_savgol_features[:,5],2))
plt.figure(3)
plt.plot(acc_plain, color='blue', label='raw', marker='.')
plt.plot(acc_ekf, color='green', label='ekf', marker='.')
plt.plot(acc_savgol, color='orange', label='savgol', marker='.')
plt.plot(acc_ekf_savgol, color='red', label='ekf-savgol', marker='.')
plt.legend()
plt.savefig(f'acc_{index}.png')
plt.clf()
j_plain = np.sqrt(np.power(plain_features[:,6],2) + np.power(plain_features[:,7],2))
j_ekf = np.sqrt(np.power(ekf_features[:,6],2) + np.power(ekf_features[:,7],2))
j_savgol = np.sqrt(np.power(savgol_features[:,6],2) + np.power(savgol_features[:,7],2))
j_ekf_savgol = np.sqrt(np.power(ekf_savgol_features[:,6],2) + np.power(ekf_savgol_features[:,7],2))
plt.figure(3)
plt.plot(j_plain, color='blue', label='raw', marker='.')
plt.plot(j_ekf, color='green', label='ekf', marker='.')
plt.plot(j_savgol, color='orange', label='savgol', marker='.')
plt.plot(j_ekf_savgol, color='red', label='ekf-savgol', marker='.')
plt.legend()
plt.savefig(f'j_{index}.png')
plt.clf()
all_vel_raw.append(vel_plain)
all_vel_ekf.append(vel_ekf)
all_vel_savgol.append(vel_savgol)
all_vel_ekf_savgol.append(vel_ekf_savgol)
all_acc_raw.append(acc_plain)
all_acc_ekf.append(acc_ekf)
all_acc_savgol.append(acc_savgol)
all_acc_ekf_savgol.append(acc_ekf_savgol)
all_j_raw.append(j_plain)
all_j_ekf.append(j_ekf)
all_j_savgol.append(j_savgol)
all_j_ekf_savgol.append(j_ekf_savgol)
all_vel_raw = np.concatenate(all_vel_raw)
all_vel_ekf = np.concatenate(all_vel_ekf)
all_vel_savgol = np.concatenate(all_vel_savgol)
all_vel_ekf_savgol = np.concatenate(all_vel_ekf_savgol)
all_acc_raw = np.concatenate(all_acc_raw)
all_acc_ekf = np.concatenate(all_acc_ekf)
all_acc_savgol = np.concatenate(all_acc_savgol)
all_acc_ekf_savgol = np.concatenate(all_acc_ekf_savgol)
all_j_raw = np.concatenate(all_j_raw)
all_j_ekf = np.concatenate(all_j_ekf)
all_j_savgol = np.concatenate(all_j_savgol)
all_j_ekf_savgol = np.concatenate(all_j_ekf_savgol)
print('\033[92m PLAIN \033[0m')
self.stats(traj=[all_vel_raw, all_acc_raw, all_j_raw])
print('\033[92m EKF \033[0m')
self.stats(traj=[all_vel_ekf, all_acc_ekf, all_j_ekf])
print('\033[92m SAVGOL \033[0m')
self.stats(traj=[all_vel_savgol, all_acc_savgol, all_j_savgol])
print('\033[92m EKF-SAVGOL \033[0m')
self.stats(traj=[all_vel_ekf_savgol, all_acc_ekf_savgol, all_j_ekf_savgol])
plt.figure(1)
plt.boxplot([all_vel_raw,all_vel_ekf,all_vel_savgol,all_vel_ekf_savgol], labels=['raw','ekf', 'savgol', 'ekf-savgol'])
plt.savefig('box_plot_vel.png')
plt.clf()
plt.figure(2)
plt.boxplot([all_acc_raw,all_acc_ekf,all_acc_savgol,all_acc_ekf_savgol], labels=['raw','ekf', 'savgol', 'ekf-savgol'])
plt.savefig('box_plot_acc.png')
plt.clf()
plt.figure(3)
plt.boxplot([all_j_raw,all_j_ekf,all_j_savgol,all_j_ekf_savgol], labels=['raw','ekf', 'savgol', 'ekf-savgol'])
plt.savefig('box_plot_j.png')
plt.clf()
plt.figure(1)
plt.boxplot([all_vel_ekf,all_vel_savgol,all_vel_ekf_savgol], labels=['ekf', 'savgol', 'ekf-savgol'])
plt.savefig('2box_plot_vel.png')
plt.clf()
plt.figure(2)
plt.boxplot([all_acc_ekf,all_acc_savgol,all_acc_ekf_savgol], labels=['ekf', 'savgol', 'ekf-savgol'])
plt.savefig('2box_plot_acc.png')
plt.clf()
plt.figure(3)
plt.boxplot([all_j_ekf,all_j_savgol,all_j_ekf_savgol], labels=['ekf', 'savgol', 'ekf-savgol'])
plt.savefig('2box_plot_j.png')
plt.clf()
def stats(self, traj:np.ndarray) -> NoReturn:
#central tendency : mean
#dispersion : std
#bounds : min max
#quantile : 0.25, 0.5, 0.75
labels = ['vel', 'acc', 'jerk']
for t, l in zip(traj, labels):
_mean = round(np.mean(t),2)
_std = round(np.std(t),2)
_min = round(np.min(t),2)
_max = round(np.max(t),2)
_q25 = round(np.quantile(t, 0.25),2)
_q50 = round(np.quantile(t, 0.5),2)
_q75 = round(np.quantile(t, 0.75),2)
print (f'Feature: {l}')
print ('\tmean:{} | std:{} | min:{} | max:{} | q25:{} | q50:{} | q75:{}'.format(_mean,
_std, _min, _max, _q25, _q50, _q75))
def test_get(data_loader: ArgoverseForecastingLoader) -> None:
data_1 = data_loader.get("0")
data_2 = data_loader[0]
assert data_1.current_seq == data_2.current_seq
class BaseDataset(torch.utils.data.Dataset):
def __init__(self,
data_dict: Dict[str, Any],
args: Any,
mode: str,
base_dir="/work/vita/sadegh/argo/argoverse-api/",
use_history=True,
use_agents=True,
use_scene=True):
"""Initialize the Dataset.
Args:
data_dict: Dict containing all the data
args: Arguments passed to the baseline code
mode: train/val/test mode
"""
self.data_dict = data_dict
self.args = args
self.mode = mode
self.use_history = use_history
self.use_agents = use_agents
self.use_scene = use_scene
# Get input
self.input_data = data_dict["{}_input".format(mode)]
if mode != "test":
self.output_data = data_dict["{}_output".format(mode)]
self.data_size = self.input_data.shape[0]
# Get helpers
self.helpers = self.get_helpers()
self.helpers = list(zip(*self.helpers))
middle_dir = mode if mode != "test" else "test_obs"
self.root_dir = base_dir + middle_dir + "/data"
##set root_dir to the correct path to your dataset folder
self.afl = ArgoverseForecastingLoader(self.root_dir)
self.avm = ArgoverseMap()
self.mf = MapFeaturesUtils()
def __len__(self):
"""Get length of dataset.
Returns:
Length of dataset
"""
return self.data_size
def __getitem__(
self,
idx: int) -> Tuple[torch.FloatTensor, Any, Dict[str, np.ndarray]]:
"""Get the element at the given index.
Args:
idx: Query index
Returns:
A list containing input Tensor, Output Tensor (Empty if test) and viz helpers.
"""
helper = self.helpers[idx]
# hp=helper[0][:20]
# seq_to_find_lanes=np.concatenate([hp,[hp[-1]+i*(hp[-1]-hp[-2]) for i in range(1,10)]])
############################# find lanes
cnt_lines, img, cnt_lines_norm, world_to_image_space = self.mf.get_candidate_centerlines_for_trajectory(
helper[0][:20],
yaw_deg=helper[5],
city_name=helper[1][0],
avm=self.avm,
viz=True,
seq_len=60,
max_candidates=MAX_NUM_LANES,
)
#############################
# normalize history
traj = helper[0] if self.mode != "test" else helper[0][:20]
traj = transform_points(traj - helper[0][19],
yaw_as_rotation33(math.pi * helper[5] / 180))
path_type = []
path = []
history_agent_type = []
history_agent = []
agents_num = 0
normal_agents_hist = []
if self.use_history or self.use_agents:
if self.use_history:
ego_world_history = helper[0][:20]
history_xy = transform_points(ego_world_history,
world_to_image_space)
history_xy = crop_tensor(history_xy, (224, 224))
history_agent_type += [1]
history_agent += [history_xy]
if self.use_agents:
agents_history, normal_agents_hist, agents_num = self.get_agents(
idx, world_to_image_space, helper[0][19], helper[5])
history_agent_type += [2] * len(agents_history)
history_agent += agents_history
history_agent = torch.cat(
[linear_path_to_tensor(lane, -1) for lane in history_agent], 0)
if self.use_scene:
path_type = [0] * len(cnt_lines_norm)
path = torch.cat(
[linear_path_to_tensor(lane, -1) for lane in cnt_lines_norm],
0)
return {
"history_positions":
torch.FloatTensor(traj[:self.args.obs_len]),
"normal_agents_history":
normal_agents_hist,
"agents_num":
agents_num,
"target_positions":
torch.empty(1) if self.mode == "test" else torch.FloatTensor(
traj[self.args.obs_len:]),
"path":
path,
"path_type":
path_type,
"history_agent":
history_agent,
"history_agent_type":
history_agent_type,
"base_image":
img.transpose(2, 0, 1),
"centroid":
helper[0][19],
"yaw_deg":
helper[5],
"seq_id":
helper[8],
"world_to_image_space":
world_to_image_space,
}
def get_helpers(self) -> Tuple[Any]:
"""Get helpers for running baselines.
Returns:
helpers: Tuple in the format specified by LSTM_HELPER_DICT_IDX
Note: We need a tuple because DataLoader needs to index across all these helpers simultaneously.
"""
helper_df = self.data_dict[f"{self.mode}_helpers"]
candidate_centerlines = helper_df["CANDIDATE_CENTERLINES"].values
# print("ss",candidate_centerlines)
candidate_nt_distances = helper_df["CANDIDATE_NT_DISTANCES"].values
xcoord = np.stack(
helper_df["FEATURES"].values)[:, :,
FEATURE_FORMAT["X"]].astype("float")
ycoord = np.stack(
helper_df["FEATURES"].values)[:, :,
FEATURE_FORMAT["Y"]].astype("float")
centroids = np.stack((xcoord, ycoord), axis=2)
_DEFAULT_HELPER_VALUE = np.full((centroids.shape[0]), None)
city_names = np.stack(
helper_df["FEATURES"].values)[:, :, FEATURE_FORMAT["CITY_NAME"]]
seq_paths = helper_df["SEQUENCE"].values
translation = (helper_df["TRANSLATION"].values
if self.args.normalize else _DEFAULT_HELPER_VALUE)
rotation = (helper_df["ROTATION"].values
if self.args.normalize else _DEFAULT_HELPER_VALUE)
use_candidates = self.args.use_map and self.mode == "test"
candidate_delta_references = (
helper_df["CANDIDATE_DELTA_REFERENCES"].values
if self.args.use_map and use_candidates else _DEFAULT_HELPER_VALUE)
delta_reference = (helper_df["DELTA_REFERENCE"].values
if self.args.use_delta and not use_candidates else
_DEFAULT_HELPER_VALUE)
helpers = [None for i in range(len(LSTM_HELPER_DICT_IDX))]
# Name of the variables should be the same as keys in LSTM_HELPER_DICT_IDX
for k, v in LSTM_HELPER_DICT_IDX.items():
helpers[v] = locals()[k.lower()]
return tuple(helpers)
def get_agents(self, index, world_to_image_space, centroid, yaw_deg):
"""Get agents
"""
helper_df = self.data_dict[f"{self.mode}_helpers"]
seq_id = helper_df.iloc[index, 0]
seq_path = f"{self.root_dir}/{seq_id}.csv"
# print(seq_path)
df = self.afl.get(seq_path).seq_df
frames = df.groupby("TRACK_ID")
res = []
normal_agents_hist = np.full((MAX_AGENTS_NUM, self.args.obs_len, 2),
300)
# print(len(frames))
# Plot all the tracks up till current frame
num_selected = 0
rotation_mat = yaw_as_rotation33(math.pi * yaw_deg / 180)
for group_name, group_data in frames:
object_type = group_data["OBJECT_TYPE"].values[0]
# print(group_data[["X","Y"]].values.shape).
cor_xy = group_data[["X", "Y"]].to_numpy()
if cor_xy.shape[0] < 20:
continue
cor_xy = cor_xy[:self.args.obs_len]
if np.linalg.norm(centroid - cor_xy[-1]) > MIN_AGENTS_DIST:
continue
traj = transform_points(cor_xy - centroid, rotation_mat)
cor_xy = transform_points(cor_xy, world_to_image_space)
# print(cor_xy.shape)
cropped_vector = crop_tensor(cor_xy, (224, 224))
# print(cropped_vector.shape)
if len(cropped_vector) > 1:
normal_agents_hist[num_selected] = traj
res.append(cropped_vector)
num_selected += 1
if num_selected >= MAX_AGENTS_NUM:
break
# print(num_selected)
return res, normal_agents_hist, num_selected
class Argoverse_Data(Dataset):
def __init__(self,
root_dir='argoverse-data/forecasting_sample/data',
train_seq_size=20,
cuda=False,
test=False):
super(Argoverse_Data, self).__init__()
self.root_dir = root_dir
self.afl = ArgoverseForecastingLoader(self.root_dir)
self.seq_paths = glob.glob(f"{self.root_dir}/*.csv")
self.train_seq_size = train_seq_size
self.use_cuda = cuda
self.mode_test = test
def __len__(self):
return len(self.seq_paths)
def old_transform(self, trajectory):
def rotation_angle(x, y):
angle = np.arctan(abs(y / x))
direction = -1 * np.sign(x * y)
return direction * angle
translation = trajectory[0]
trajectory = trajectory - trajectory[0]
theta = rotation_angle(trajectory[19, 0], trajectory[19, 1])
c, s = np.cos(theta), np.sin(theta)
R = np.array([[c, -s], [s, c]])
trajectory = torch.tensor(trajectory)
trajectory = trajectory.permute(1, 0)
trajectory = np.matmul(R, trajectory)
trajectory = torch.tensor(trajectory)
trajectory = trajectory.permute(1, 0)
if self.mode_test:
return trajectory[0:self.train_seq_size].float(), R, translation
else:
return trajectory[0:self.train_seq_size].float(
), trajectory[self.train_seq_size:].float()
def transform(self, trajectory):
def rotation_angle(x, y):
angle = np.arctan(abs(y / x))
direction = -1 * np.sign(x * y)
return direction * angle
if self.mode_test:
translation = -trajectory[0]
train_trajectory = trajectory + translation
theta = rotation_angle(train_trajectory[19, 0],
train_trajectory[19, 1])
c, s = np.cos(theta), np.sin(theta)
R = torch.Tensor([[c, -s], [s, c]]).float()
train_trajectory = torch.tensor(train_trajectory).float()
train_trajectory = torch.matmul(R, train_trajectory.permute(
1, 0)).permute(1, 0)
return train_trajectory, R, torch.Tensor(translation).float()
else:
old_trajectory = trajectory
translation = -trajectory[0]
transformed_trajectory = trajectory + translation
theta = rotation_angle(transformed_trajectory[19, 0],
transformed_trajectory[19, 1])
c, s = np.cos(theta), np.sin(theta)
R = torch.Tensor([[c, -s], [s, c]]).float()
transformed_trajectory = torch.tensor(
transformed_trajectory).float()
transformed_trajectory = torch.matmul(
R, transformed_trajectory.permute(1, 0)).permute(1, 0)
train_trajectory = transformed_trajectory[:self.train_seq_size]
gt_transformed_trajectory = transformed_trajectory[self.
train_seq_size:]
actual_gt_trajectory = torch.Tensor(
trajectory[self.train_seq_size:]).float()
return train_trajectory, gt_transformed_trajectory, actual_gt_trajectory, R, torch.Tensor(
translation).float()
def inverse_transform_one(self, trajectory, R, t):
out = torch.matmul(R, trajectory.permute(1, 0)).permute(1, 0)
return out + t.reshape(1, 2)
def inverse_transform(self, trajectory, traj_dict):
R = traj_dict['rotation']
t = traj_dict['translation']
if self.use_cuda:
R = R.cuda()
t = t.cuda()
out = torch.matmul(R.permute(0, 2, 1),
trajectory.permute(0, 2, 1)).permute(0, 2, 1)
out = out - t.reshape(t.shape[0], 1, 2)
return out
def __getitem__(self, index):
'''
Obtain neighbour trajectories as well.
Obtain map parameters at the trajectories
Do it in the coordinates of the centerlines as well
'''
current_loader = self.afl.get(self.seq_paths[index])
agent_traj = current_loader.agent_traj
# if self.test:
# agent_train_traj,R,translation=self.transform(agent_traj)
# seq_index=int(os.path.basename(self.seq_paths[index]).split('.')[0])
# return {'seq_index': seq_index,'train_agent':agent_train_traj,'rotation':R,'translation':translation,'city':current_loader.city}
# else:
agent_train_traj, agent_gt_traj, agent_unnorm_gt_traj, R, translation = self.transform(
agent_traj)
return {
'seq_path': self.seq_paths[index],
'train_agent': agent_train_traj,
'gt_agent': agent_gt_traj,
'gt_unnorm_agent': agent_unnorm_gt_traj,
'rotation': R,
'translation': translation,
'city': current_loader.city
}
class TrajExtractor:
def __init__(self, root_dir, config, mode):
self.avl = ArgoverseForecastingLoader(root_dir)
self.avl.seq_list.sort()
self.config = config
self.train = mode == 'train'
logging.info(f'root_dir: {root_dir}, num of squence: {len(self.avl)}')
def __del__(self):
del self.avl
def extract(self, index=None, seq_path=None):
assert index is not None or Path(seq_path).exists(), 'cannot get sequence: no index is provided and no seq_path file'
if index is not None:
seq = self.avl[index]
else:
seq = self.avl.get(seq_path)
data = self.read_argo_data(seq)
data = self.get_obj_feats(data)
return data
def read_argo_data(self, seq):
city = copy.deepcopy(seq.city)
"""TIMESTAMP,TRACK_ID,OBJECT_TYPE,X,Y,CITY_NAME"""
df = copy.deepcopy(seq.seq_df)
agt_ts = np.sort(np.unique(df['TIMESTAMP'].values))
mapping = dict()
for i, ts in enumerate(agt_ts):
mapping[ts] = i
trajs = np.concatenate((
df.X.to_numpy().reshape(-1, 1),
df.Y.to_numpy().reshape(-1, 1)), 1)
steps = [mapping[x] for x in df['TIMESTAMP'].values]
steps = np.asarray(steps, np.int64)
objs = df.groupby(['TRACK_ID', 'OBJECT_TYPE']).groups
keys = list(objs.keys())
obj_type = [x[1] for x in keys]
agt_idx = obj_type.index('AGENT')
idcs = objs[keys[agt_idx]]
agt_traj = trajs[idcs]
agt_step = steps[idcs]
del keys[agt_idx]
ctx_trajs, ctx_steps = [], []
# ctx means neighbor (in my mind)
for key in keys:
idcs = objs[key]
ctx_trajs.append(trajs[idcs])
ctx_steps.append(steps[idcs])
data = dict()
data['file_id'] = self.avl.current_seq
data['city'] = city
data['trajs'] = [agt_traj] + ctx_trajs # [agent, other neighbor] traj
data['steps'] = [agt_step] + ctx_steps # [agent, other nrighbor] traj corresponding step index
return data
def get_obj_feats(self, data):
orig = data['trajs'][0][19].copy().astype(np.float32) # last point of agent.
# add noise at train phase
if self.train and self.config['rot_aug']:
theta = np.random.rand() * np.pi * 2.0
else:
pre = data['trajs'][0][18] - orig
theta = np.pi - np.arctan2(pre[1], pre[0]) # theta is the last 2 frame heading
rot = np.asarray([ # a sequence has the same rot mat
[np.cos(theta), -np.sin(theta)],
[np.sin(theta), np.cos(theta)]], np.float32)
feats, ctrs, gt_preds, has_preds, hist_trajs = [], [], [], [], []
for traj, step in zip(data['trajs'], data['steps']):
if 19 not in step:
continue # if an neighbor is not present at the last frame, omit it.
gt_pred = np.zeros((30, 2), np.float32)
has_pred = np.zeros(30, np.bool)
future_mask = np.logical_and(step >= 20, step < 50)
post_step = step[future_mask] - 20
post_traj = traj[future_mask]
gt_pred[post_step] = post_traj
has_pred[post_step] = 1
obs_mask = step < 20
step = step[obs_mask]
traj = traj[obs_mask]
idcs = step.argsort() # why argsort
step = step[idcs]
traj = traj[idcs]
# valid history step mask: must be continueous till last observed point
for i in range(len(step)): # clip gap
if step[i] == 19 - (len(step) - 1) + i:
break
step = step[i:]
traj = traj[i:]
feat = np.zeros((20, 3), np.float32)
feat[step, :2] = np.matmul(rot, (traj - orig.reshape(-1, 2)).T).T # route history
feat[step, 2] = 1.0
x_min, x_max, y_min, y_max = self.config['pred_range'] # clip extent
if feat[-1, 0] < x_min or feat[-1, 0] > x_max or feat[-1, 1] < y_min or feat[-1, 1] > y_max:
continue
ctrs.append(feat[-1, :2].copy()) # ctrs: abs xy coor (rotated, each agent's base pt)
feat[1:, :2] -= feat[:-1, :2]
feat[step[0], :2] = 0
feats.append(feat) # offset, rotated; valid flag
gt_preds.append(gt_pred) # abs xy coor, (no rotated, no based on base pt)
has_preds.append(has_pred) # loss flag
hist_trajs.append(traj)
feats = np.asarray(feats, np.float32)
ctrs = np.asarray(ctrs, np.float32)
gt_preds = np.asarray(gt_preds, np.float32)
has_preds = np.asarray(has_preds, np.bool)
data['feats'] = feats
data['ctrs'] = ctrs
data['orig'] = orig
data['theta'] = theta
data['rot'] = rot
data['gt_preds'] = gt_preds
data['has_preds'] = has_preds
# add by zhaoyi
data['hist_traj'] = hist_trajs
data['fur_traj'] = gt_preds
return data
def save_top_errors_accuracy_single_pred(self):
afl=ArgoverseForecastingLoader("data/val/data/")
self.model.load_state_dict(torch.load(self.model_dir+'best-model.pt')['model_state_dict'])
self.model.eval()
min_loss=np.inf
max_loss=0
num_images=10
loss_list_max=[]
input_max_list=[]
pred_max_list=[]
target_max_list=[]
city_name_max=[]
seq_path_list_max=[]
loss_list_min=[]
input_min_list=[]
pred_min_list=[]
target_min_list=[]
city_name_min=[]
seq_path_list_min=[]
num_batches=len(self.val_loader.batch_sampler)
for i_batch,traj_dict in enumerate(self.val_loader):
print(f"Running {i_batch}/{num_batches}",end="\r")
# pdb.set_trace()
# pred_traj=self.model(traj_dict)
# pred_traj=self.val_loader.dataset.inverse_transform(pred_traj,traj_dict)
gt_traj=traj_dict['gt_unnorm_traj']
output=self.model(traj_dict,mode='validate')
output=output.cpu()
loss=torch.norm(output.reshape(output.shape[0],-1)-gt_traj.reshape(gt_traj.shape[0],-1),dim=1)
min_loss,min_index=torch.min(loss,dim=0)
max_loss,max_index=torch.max(loss,dim=0)
print(f"Min loss: {min_loss}, Max loss: {max_loss}" )
seq_path_list_max.append(traj_dict['seq_path'][max_index])
seq_path_list_min.append(traj_dict['seq_path'][min_index])
loader_max=afl.get(traj_dict['seq_path'][max_index])
input_max_list.append(loader_max.agent_traj[0:20,:])
city_name_max.append(loader_max.city)
del loader_max
loader_min=afl.get(traj_dict['seq_path'][min_index])
input_min_list.append(loader_min.agent_traj[0:20,:])
city_name_min.append(loader_min.city)
del loader_min
pred_min_list.append(output[min_index])
target_min_list.append(gt_traj[min_index])
pred_max_list.append(output[max_index])
target_max_list.append(gt_traj[max_index])
# loss_list_max.append(min_loss.data)
# loss_list_min.append(max_loss.data)
loss_list_max.append(max_loss.data)
loss_list_min.append(min_loss.data)
# torch.cuda.empty_cache()
# pdb.set_trace()
loss_list_max_array=np.array(loss_list_max)
loss_list_max=list(loss_list_max_array.argsort()[-num_images:][::-1])
loss_list_min_array=np.array(loss_list_min)
loss_list_min=list(loss_list_min_array.argsort()[:num_images])
# pdb.set_trace()
avm=ArgoverseMap()
high_error_path=model_dir+"/visualization/high_errors/"
low_error_path=model_dir+"/visualization/low_errors/"
if not os.path.exists(high_error_path):
os.makedirs(high_error_path)
if not os.path.exists(low_error_path):
os.makedirs(low_error_path)
# if self.use_cuda:
# input_=input_.cpu()
# output=output.cpu()
input_max=[]
pred_max=[]
target_max=[]
city_max=[]
# import pdb;pdb.set_trace()
# seq_path_max=[]
centerlines_max=[]
for i,index in enumerate(loss_list_max):
print(f"Max: {i}")
# pdb.set_trace()
input_max.append(input_max_list[index])
pred_max.append([pred_max_list[index].detach().numpy()])
print(f"Difference in predicted and input_traj for maximum at {i} is {np.linalg.norm(input_max[-1]-pred_max[-1][0][0:20,:])}")
target_max.append(target_max_list[index].detach().numpy())
city_max.append(city_name_max[index])
viz_sequence(df=pd.read_csv(seq_path_list_max[index]) ,save_path=f"{high_error_path}/dataframe_{i}.png",show=True,avm=avm)
centerlines_max.append(avm.get_candidate_centerlines_for_traj(input_max[-1], city_max[-1],viz=False))
print("Created max array")
input_min=[]
pred_min=[]
target_min=[]
city_min=[]
# seq_path_min=[]
centerlines_min=[]
for i,index in enumerate(loss_list_min):
# pdb.set_trace()
print(f"Min: {i}")
input_min.append(input_min_list[index])
pred_min.append([pred_min_list[index].detach().numpy()])
print(f"Difference in predicted and input_traj for minimum at {i} is {np.linalg.norm(input_min[-1]-pred_min[-1][0][0:20,:])}")
target_min.append(target_min_list[index].detach().numpy())
city_min.append(city_name_min[index])
# seq_path_min.append(seq_path_list_min[index])
viz_sequence(df=pd.read_csv(seq_path_list_min[index]) ,save_path=f"{low_error_path}/dataframe_{i}.png",show=True,avm=avm)
centerlines_min.append(avm.get_candidate_centerlines_for_traj(input_min[-1], city_min[-1],viz=False))
print("Created min array")
print(f"Saving max visualizations at {high_error_path}")
# import pdb;pdb.set_trace()
viz_predictions(input_=np.array(input_max), output=pred_max,target=np.array(target_max),centerlines=centerlines_max,city_names=np.array(city_max),avm=avm,save_path=high_error_path)
print(f"Saving min visualizations at {low_error_path}")
# viz_predictions(input_=pred_min[0], output=pred_min,target=np.array(target_min),centerlines=centerlines_min,city_names=np.array(city_min),avm=avm,save_path=low_error_path)
# viz_predictions(input_=np.expand_dims(input_min[0],axis=0),output=pred_min[0],target=np.expand_dims(target_min[0],axis=0),centerlines=[centerlines_min], city_names=np.expand_dims(city_min[0],axis=0),avm=avm,save_path=low_error_path)
viz_predictions(input_=np.array(input_min), output=pred_min,target=np.array(target_min),centerlines=centerlines_min,city_names=np.array(city_min),avm=avm,save_path=low_error_path)
class ArgoverseForecastDataset(torch.utils.data.Dataset):
def __init__(self, cfg):
super().__init__()
self.am = ArgoverseMap()
self.axis_range = self.get_map_range(self.am)
self.city_halluc_bbox_table, self.city_halluc_tableidx_to_laneid_map = self.am.build_hallucinated_lane_bbox_index(
)
self.laneid_map = self.process_laneid_map()
self.vector_map, self.extra_map = self.generate_vector_map()
# am.draw_lane(city_halluc_tableidx_to_laneid_map['PIT']['494'], 'PIT')
# self.save_vector_map(self.vector_map)
self.last_observe = cfg['last_observe']
##set root_dir to the correct path to your dataset folder
self.root_dir = cfg['data_locate']
self.afl = ArgoverseForecastingLoader(self.root_dir)
self.map_feature = dict(PIT=[], MIA=[])
self.city_name, self.center_xy, self.rotate_matrix = dict(), dict(
), dict()
def __len__(self):
return len(self.afl)
def __getitem__(self, index):
# self.am.find_local_lane_polygons()
self.trajectory, city_name, extra_fields = self.get_trajectory(index)
traj_id = extra_fields['trajectory_id'][0]
self.city_name.update({str(traj_id): city_name})
center_xy = self.trajectory[self.last_observe - 1][1]
self.center_xy.update({str(traj_id): center_xy})
trajectory_feature = (self.trajectory -
np.array(center_xy).reshape(1, 1, 2)).reshape(
-1, 4)
rotate_matrix = get_rotate_matrix(
trajectory_feature[self.last_observe, :]) # rotate coordinate
self.rotate_matrix.update({str(traj_id): rotate_matrix})
trajectory_feature = ((trajectory_feature.reshape(-1, 2)).dot(
rotate_matrix.T)).reshape(-1, 4)
trajectory_feature = self.normalize_coordinate(
trajectory_feature, city_name) # normalize to [-1, 1]
self.traj_feature = torch.from_numpy(
np.hstack((
trajectory_feature,
extra_fields['TIMESTAMP'].reshape(-1, 1),
# extra_fields['OBJECT_TYPE'].reshape(-1, 1),
extra_fields['trajectory_id'].reshape(-1, 1)))).float()
map_feature_dict = dict(PIT=[], MIA=[])
for city in ['PIT', 'MIA']:
for i in range(len(self.vector_map[city])):
map_feature = (self.vector_map[city][i] -
np.array(center_xy).reshape(1, 1, 2)).reshape(
-1, 2)
map_feature = (map_feature.dot(rotate_matrix.T)).reshape(-1, 4)
map_feature = self.normalize_coordinate(map_feature, city)
tmp_tensor = torch.from_numpy(
np.hstack((
map_feature,
self.extra_map[city]['turn_direction'][i],
self.extra_map[city]['in_intersection'][i],
self.extra_map[city]['has_traffic_control'][i],
# self.extra_map[city]['OBJECT_TYPE'][i],
self.extra_map[city]['lane_id'][i])))
map_feature_dict[city].append(tmp_tensor.float())
# self.map_feature[city] = np.array(self.map_feature[city])
self.map_feature[city] = map_feature_dict[city]
self.map_feature['city_name'] = city_name
return self.traj_feature, self.map_feature
def get_trajectory(self, index):
seq_path = self.afl.seq_list[index]
data = self.afl.get(seq_path).seq_df
data = data[data['OBJECT_TYPE'] == 'AGENT']
extra_fields = dict(TIMESTAMP=[], OBJECT_TYPE=[], trajectory_id=[])
polyline = []
j = int(str(seq_path).split('/')[-1].split('.')[0])
flag = True
city_name = ''
for _, row in data.iterrows():
if flag:
xlast = row['X']
ylast = row['Y']
tlast = row['TIMESTAMP']
city_name = row['CITY_NAME']
flag = False
continue
startpoint = np.array([xlast, ylast])
endpoint = np.array([row['X'], row['Y']])
xlast = row['X']
ylast = row['Y']
extra_fields['TIMESTAMP'].append(tlast)
extra_fields['OBJECT_TYPE'].append(0) # 'AGENT'
extra_fields['trajectory_id'].append(j) # 'AGENT'
tlast = row['TIMESTAMP']
polyline.append([startpoint, endpoint])
extra_fields['TIMESTAMP'] = np.array(extra_fields['TIMESTAMP'])
extra_fields['TIMESTAMP'] -= np.min(
extra_fields['TIMESTAMP']) # adjust time stamp
extra_fields['OBJECT_TYPE'] = np.array(extra_fields['OBJECT_TYPE'])
extra_fields['trajectory_id'] = np.array(extra_fields['trajectory_id'])
return np.array(polyline), city_name, extra_fields
def generate_vector_map(self):
vector_map = {'PIT': [], 'MIA': []}
extra_map = {
'PIT':
dict(OBJECT_TYPE=[],
turn_direction=[],
lane_id=[],
in_intersection=[],
has_traffic_control=[]),
'MIA':
dict(OBJECT_TYPE=[],
turn_direction=[],
lane_id=[],
in_intersection=[],
has_traffic_control=[])
}
polyline = []
# index = 1
pbar = tqdm(total=17326)
pbar.set_description("Generating Vector Map")
for city_name in ['PIT', 'MIA']:
for key in self.laneid_map[city_name]:
pts = self.am.get_lane_segment_polygon(key, city_name)
turn_str = self.am.get_lane_turn_direction(key, city_name)
if turn_str == 'LEFT':
turn = -1
elif turn_str == 'RIGHT':
turn = 1
else:
turn = 0
pts_len = pts.shape[0] // 2
positive_pts = pts[:pts_len, :2]
negative_pts = pts[pts_len:2 * pts_len, :2]
polyline.clear()
for i in range(pts_len - 1):
v1 = np.array([positive_pts[i], positive_pts[i + 1]])
v2 = np.array([
negative_pts[pts_len - 1 - i],
negative_pts[pts_len - i - 2]
])
polyline.append(v1)
polyline.append(v2)
# extra_field['table_index'] = self.laneid_map[city_name][key]
repeat_t = 2 * (pts_len - 1)
vector_map[city_name].append(np.array(polyline).copy())
extra_map[city_name]['turn_direction'].append(
np.repeat(turn, repeat_t, axis=0).reshape(-1, 1))
extra_map[city_name]['OBJECT_TYPE'].append(
np.repeat(-1, repeat_t, axis=0).reshape(-1, 1)) #HD Map
extra_map[city_name]['lane_id'].append(
np.repeat(int(key), repeat_t, axis=0).reshape(-1, 1))
extra_map[city_name]['in_intersection'].append(
np.repeat(1 *
self.am.lane_is_in_intersection(key, city_name),
repeat_t,
axis=0).reshape(-1, 1))
extra_map[city_name]['has_traffic_control'].append(
np.repeat(1 * self.am.lane_has_traffic_control_measure(
key, city_name),
repeat_t,
axis=0).reshape(-1, 1))
# if index > 10:
# break
# index = index + 1
pbar.update(1)
pbar.close()
print("Generate Vector Map Successfully!")
return vector_map, extra_map #vector_map:list
def process_laneid_map(self):
laneid_map = {}
tmp_map = {}
tmp1_map = {}
for key in self.city_halluc_tableidx_to_laneid_map['PIT']:
tmp_map[self.city_halluc_tableidx_to_laneid_map['PIT'][key]] = key
laneid_map['PIT'] = tmp_map
for key in self.city_halluc_tableidx_to_laneid_map['MIA']:
tmp1_map[self.city_halluc_tableidx_to_laneid_map['MIA'][key]] = key
laneid_map['MIA'] = tmp1_map
return laneid_map
def get_map_range(self, am):
map_range = dict(PIT={}, MIA={})
for city_name in ['PIT', 'MIA']:
poly = am.get_vector_map_lane_polygons(city_name)
poly_modified = (np.vstack(poly))[:, :2]
max_coordinate = np.max(poly_modified, axis=0)
min_coordinate = np.min(poly_modified, axis=0)
map_range[city_name].update({'max': max_coordinate})
map_range[city_name].update({'min': min_coordinate})
return map_range
def normalize_coordinate(self, array, city_name):
max_coordinate = self.axis_range[city_name]['max']
min_coordinate = self.axis_range[city_name]['min']
array = (10. * (array.reshape(-1, 2)) /
(max_coordinate - min_coordinate)).reshape(-1, 4)
return array
def save_vector_map(self, vector_map):
save_path = "./data/vector_map/"
for city_name in ['PIT', 'MIA']:
tmp_map = np.vstack(vector_map[city_name]).reshape(-1, 4)
np.save(save_path + city_name + "_vectormap", tmp_map)
import re
import pickle
# %matplotlib inline
if __name__ == "__main__":
am = ArgoverseMap()
for folder in os.listdir(DATA_DIR):
if not re.search(r'val', folder):
# FIXME: modify the target folder by hand ('val|train|sample|test')
# if not re.search(r'test', folder):
continue
print(f"folder: {folder}")
afl = ArgoverseForecastingLoader(os.path.join(DATA_DIR, folder))
norm_center_dict = {}
for name in tqdm(afl.seq_list):
afl_ = afl.get(name)
path, name = os.path.split(name)
name, ext = os.path.splitext(name)
agent_feature, obj_feature_ls, lane_feature_ls, norm_center = compute_feature_for_one_seq(
afl_.seq_df,
am,
OBS_LEN,
LANE_RADIUS,
OBJ_RADIUS,
viz=False,
mode='nearby')
df = encoding_features(agent_feature, obj_feature_ls,
lane_feature_ls)
save_features(
df, name,
class Argoverse(object):
def __init__(self, root : str):
"""
create the dataset manager object
- params:
- root : str : root path for the data
"""
self.root_path = root
# - seq_list -> list with paths for each sequence file
self.manager = ArgoverseForecastingLoader(self.root_path)
self.ekf = EKF()
self.savgol = SavitzkyGolov(window_length=5,
poly=3)
self.compute_features = PlainFeatures()
def _get_trajectories_by_track_id(self,
sequence:str,
track_id:str,
obs_len: int=2)->List[pd.DataFrame]:
"""
return a segment of trajectory for each vehicle
defined by its track_id and sequence file
- params:
- sequence : str : sequence file
- track_id : str : id of the agent
- obs_len : int : length of the trajectory segment in seconds
- return:
- segments : List[pd.DataFrame] : list of trajectories with
length 'obs_len' in seconds
"""
obs_len = obs_len*10 #10Hz
df = self.manager.get(sequence).seq_df
if track_id == 'AGENT':
traj = df[df['OBJECT_TYPE']==track_id]
else:
traj = df[df['TRACK_ID']==track_id]
traj = traj[['TIMESTAMP', 'X', 'Y']]
#print(np.shape(traj))
segments = []
if traj.shape[0] >= obs_len:
segments = [traj.iloc[i*obs_len:(i+1)*obs_len].values \
for i in range(0, traj.shape[0]//obs_len)]
return np.asarray(segments)
def _get_trajectories_by_sequence(self,
sequence:str,
target_name:str='AGENT',
obs_len:int=2)->List[pd.DataFrame]:
"""
return all segments of trajectory with length
obs_len (in seconds) from all traffic agents
of a sequence
- params:
- sequence: str : name of the sequence (file)
- obs_len : int : length of the segments (in seconds)
- target_name : str : 'AGENT' whether to return the path of the
agent of each sequence only
'ALL' whether to return the path of all
vehicles of the sequence
- return:
- segments: np.ndarray : List with all segments
of trajectory with length 'obs_len'
"""
assert target_name in ['AGENT', 'ALL'],\
('[Argoverse Manager][_get_trajectories_by_sequence]'
' unkown target_name ({})').format(target_name)
def _get_segments(track_ids: List,
start_idx:int,
sequence:str,
batch_size:int=10,
obs_len:int=2)->List:
result = []
for _v_id in track_ids[start_idx:start_idx+batch_size]:
t = self._get_trajectories_by_track_id(sequence=sequence,
obs_len=obs_len,
track_id=_v_id)
if np.shape(t)[0]>0:
result.append(t)
return result
# main function
_r = []
if target_name=='ALL':
track_ids = self.manager.get(sequence).track_id_list
batch_size = max(1,round(len(track_ids)*0.5))
_r = Parallel(n_jobs=-2)(delayed(_get_segments)(
track_ids=track_ids,
start_idx=i,
sequence=sequence,
batch_size=batch_size,
obs_len=obs_len
) for i in range(0, len(track_ids), batch_size))
else: #target_name == AGENT
_r = [[self._get_trajectories_by_track_id(sequence=sequence,
obs_len=obs_len,
track_id=target_name)]]
segments = []
if np.shape(_r)[0] > 0:
segments = [np.concatenate(p)\
for p in _r \
if np.shape(p)[0] > 0]
#print(np.concatenate(segments).shape, np.shape(_r))
return np.concatenate(segments) if np.shape(segments)[0] > 0 else []
def _aux_get_traj_segments(self,
start_idx:int,
sequences:List,
temp_dir:str,
mode:str,
target_name:str,
batch_size:int=10,
obs_len:int=2)->NoReturn:
"""
auxiliary function to get trajectories segments
and save them into a temp folder
"""
result = []
for seq in sequences[start_idx:start_idx+batch_size]:
_seq = self._get_trajectories_by_sequence(sequence = seq,
obs_len = obs_len,
target_name = target_name)
if np.shape(_seq)[0]>0:
result.append(_seq)
if np.shape(result)[0]>0:
result = np.concatenate(result)
print(('[Argoverse][get_trajectories][{}]'
' saving sequence [{}/{}][{}]').format(target_name,
start_idx,
start_idx+batch_size,
len(sequences)))
_path = os.path.join(temp_dir,'trajectory_{}_{}_{}.npy'.format(mode,
start_idx, start_idx+batch_size))
np.save(_path, result)
else:
print(('[Argoverse][get_trajectories][{}]'
' empty sequence [{}/{}][{}]').format(target_name,
start_idx,
start_idx+batch_size,
len(sequences)))
def get_trajectories(self,
save_dir:str,
mode:str,
batch_size:int,
merge_files:bool=True,
target_name:str='AGENT',
obs_len:int=2)->str:
"""
return all segments of trajectories with length
'obs_len' in seconds of the whole dataset
- params:
- obs_len : int : length in seconds
- save_dir: str : path to save the result
- mode : str : train/test
- merge_files: bool : True, if all files should be merger
into a single file
False, otherwise
- target_name : str : 'AGENT' whether to return the path of the
agent of each sequence only
'ALL' whether to return the path of all
vehicles of the sequence
- return:
- traj_path: path in which the trajectories
files were saved
"""
assert target_name in ['AGENT', 'ALL'],\
('[Argoverse Manager][get_trajectories]'
' unkown target_name ({})').format(target_name)
sequences = self.manager.seq_list#[:500]
batch_size = min(batch_size, int(len(sequences)*0.2))
temp_save_dir = tempfile.mkdtemp()
# get trajectories and save them into a temp dir
Parallel(n_jobs=-2)(delayed(self._aux_get_traj_segments)(
start_idx=i,
sequences=sequences,
batch_size=batch_size,
obs_len=obs_len,
temp_dir=temp_save_dir,
mode=mode,
target_name=target_name
) for i in range(0, len(sequences), batch_size))
if merge_files:
# merge all trajectory files saved into the temp dir
# remove all temp files after the merge
print('[Argoverse][get_trajectories] merging files...')
self.merge_saved_features(temp_dir=temp_save_dir,
save_dir=save_dir,
mode=mode,
prefix=f'trajectory_{obs_len}s_{target_name}')
else:
# return the path in which trajectory segments files
# were stored
return temp_save_dir
def _savgol_traj(self, traj:np.ndarray)->np.ndarray:
'''
applied a Savitzky-Golov filter to the trajectory
- path, velocity, acceleration and jerk
- params:
traj : nd.array : path [[x,y]] (m,2)
- return:
filtered : nd.array : [[x,y,v_x, v_y, a_x, a_y, j_x, j_y]]
'''
self.savgol.set_window_size(traj.shape[0]//2)
return self.savgol.process(traj)
def _ekf_savgol_traj(self, traj:np.ndarray)->np.ndarray:
'''
applied a Savitzky-Golov filter to the trajectory after
compute features using ekf
- path, velocity, acceleration and jerk
- params:
traj : nd.array : path [[x,y]] (m,2)
- return:
filtered : nd.array : [[x,y,v_x, v_y, a_x, a_y, j_x, j_y]]
'''
ekf_traj = self.ekf.process(traj)
self.savgol.set_window_size(traj.shape[0]//2)
return self.savgol.filter(ekf_traj)
def _ekf_traj(self, traj:np.ndarray)->np.ndarray:
'''
compute features using ekf
- path, velocity, acceleration and jerk
- params:
traj : nd.array : path [[x,y]] (m,2)
- return:
filtered : nd.array : [[x,y,v_x, v_y, a_x, a_y, j_x, j_y]]
'''
return self.ekf.process(traj)
def _none_traj(self, traj:np.ndarray)->np.ndarray:
'''
compute features using the raw data
- path, velocity, acceleration and jerk
- params:
traj : nd.array : path [[x,y]] (m,2)
- return:
filtered : nd.array : [[x,y,v_x, v_y, a_x, a_y, j_x, j_y]]
'''
return self.compute_features.process(traj)
def _aux_get_filtered_traj(self,
start_idx:int,
trajectories:np.ndarray,
batch_size:int,
mode:str,
temp_dir:str,
index_file:int,
filter_name:str)->NoReturn:
if filter_name == 'ekf':
seq_traj = np.concatenate([self._ekf_traj(traj=p)\
for p in trajectories[start_idx:start_idx+batch_size]])
elif filter_name == 'savgol':
seq_traj = np.concatenate([self._savgol_traj(traj=p)\
for p in trajectories[start_idx:start_idx+batch_size]])
elif filter_name == 'ekf-savgol':
seq_traj = np.concatenate([self._ekf_savgol_traj(traj=p)\
for p in trajectories[start_idx:start_idx+batch_size]])
elif filter_name == 'none':
seq_traj = np.concatenate([self._none_traj(traj=p)\
for p in trajectories[start_idx:start_idx+batch_size]])
else:
assert False, ('[Argoverse Manager][get_filtered_trajectories]',
' unknown filter_name ({})').format(filter_name)
print (("[Argoverse][get_filtered_trajectories][{}][{}]"
" saving sequence [{}/{}][{}]").format(filter_name,
index_file,
start_idx,
start_idx+batch_size,
len(trajectories)))
_path = os.path.join(temp_dir,\
'{}_features_{}_{}_{}.npy'.format(index_file, mode, start_idx,\
start_idx+batch_size))
np.save(_path, seq_traj)
def get_filtered_trajectories(self,
mode:str,
save_dir:str,
filter_name:str,
batch_size:int=100,
obs_len:int=2)->NoReturn:
"""
get all trajectories filtered by Extended Kalman Filter
- params:
mode : str : train/test
save_dir : str : folder in which the result
should be stored
filter_name: str : name of the filter to handle the data noise
batch_size : int : size of the batch
obs_len : int : length of each trajectory in seconds
"""
if mode == 'test':
obs_len = 2
print ("[Argoverse][get_filtered_trajectories] getting trajectories...")
traj_dir = self.get_trajectories(obs_len=obs_len,
save_dir=save_dir,
mode=mode,
merge_files=False,
batch_size=batch_size,
target_name='AGENT')
temp_save_dir = tempfile.mkdtemp()
index = 1
num_files = len(os.listdir(traj_dir))
print (("[Argoverse][get_filtered_trajectories] getting filtered"
" trajectories [{}]...").format(num_files))
for traj_file in os.listdir(traj_dir):
if (not traj_file.endswith(".npy")) or\
(mode not in traj_file) or\
('trajectory' not in traj_file):
continue
trajectories = np.load(os.path.join(traj_dir,traj_file))
batch_size = max(1, min(batch_size, int(trajectories.shape[0]*0.2)))
print (("[Argoverse][get_filtered_trajectories][{}/{}] filtering"
" trajectories...").format(index, num_files))
Parallel(n_jobs=-2)(delayed(self._aux_get_filtered_traj)(
start_idx=i,
trajectories=trajectories,
batch_size=batch_size,
mode=mode,
temp_dir=temp_save_dir,
index_file=index,
filter_name=filter_name
) for i in range(0, trajectories.shape[0], batch_size))
index = index + 1
print("[Argoverse][get_filtered_trajectories] merging files...")
self.merge_saved_features(temp_dir=temp_save_dir,
save_dir=save_dir,
mode=mode,
prefix=f'filtered_{obs_len}s_{filter_name}',
delete_after_process=False)
self.merge_saved_features(temp_dir=traj_dir,
save_dir=save_dir,
mode=mode,
prefix=f'trajectory_{obs_len}s',
delete_after_process=False)
print("[Argoverse][get_filtered_trajectories] done!")
def merge_saved_features(self,
temp_dir:str,
save_dir:str,
mode:str,
prefix:str,
delete_after_process:bool=True)->NoReturn:
"""
merge files (.npy) into a single file (.npy)
- params:
- temp_dir : str : path where the files were stored
- save_dir : str : path where the result shoudl be saved
- mode : str : train/test
- prefix : str : name for the result file
- delete_after_process : bool : true -> delete unit file after process
false -> otherwise
"""
all_features = []
for feature_file in os.listdir(temp_dir):
if (not feature_file.endswith(".npy")) or\
(mode not in feature_file):
continue
_p_file = os.path.join(temp_dir,feature_file)
all_features.append(np.load(_p_file))
if delete_after_process:
os.remove(_p_file)
all_features = np.squeeze(np.concatenate(all_features))
file_name = 'sequences_{}_{}.npy'.format(mode, prefix)
print (('[Argoverse][merge_saved_features]'
' file: {} | shape {}').format(file_name,
np.shape(all_features)))
_path = os.path.join(save_dir,file_name)
np.save(_path, all_features)
def main():
input_root, output_root, do_trajectories, do_maps, seed = get_command_line_args()
print("Preprocessing Script for Argoverse Dataset.")
print("Trajectories: {:s}, Maps: {:s}, Random Seed: {:d}".format("Y" if do_trajectories else "N", "Y" if do_maps else "N", seed))
np.random.seed(seed)
random.seed(seed)
if do_trajectories:
train_all = None
for partition in ['train', 'train_val', 'val', 'test_obs']:
print("Start Processing {:s} set.".format(partition))
if 'train' in partition:
if train_all is None:
partition_path = P(input_root).joinpath('train', 'data')
afl = ArgoverseForecastingLoader(partition_path)
seq_list = afl.seq_list
seq_list.sort()
np.random.shuffle(seq_list)
train_all = seq_list
if partition == 'train':
seq_list = train_all[int(len(train_all) * FRAC_TRAIN_VAL):]
if partition == 'train_val':
seq_list = train_all[:int(len(train_all) * FRAC_TRAIN_VAL)]
else:
partition_path = P(input_root).joinpath(partition, 'data')
afl = ArgoverseForecastingLoader(partition_path)
seq_list = afl.seq_list
pool = Pool(cpu_count())
callback = Counting_Callback(task_name="Trajectories", num_data=len(seq_list))
for seq in seq_list:
scn_code = int(seq.stem)
afl.get(seq)
df = afl.seq_df
obsv_path = P(output_root).joinpath(partition, 'observation', '{:06d}-{:03d}.pkl'.format(scn_code, REFERENCE_FRAME))
obsv_path.parent.mkdir(parents=True, exist_ok=True)
pred_path = None
if partition is not 'test_obs':
pred_path = P(output_root).joinpath(partition, 'prediction', '{:06d}-{:03d}.pkl'.format(scn_code, REFERENCE_FRAME))
pred_path.parent.mkdir(parents=True, exist_ok=True)
# generate_trajectories(df.copy(), obsv_path, pred_path)
pool.apply_async(generate_trajectories, (df.copy(), obsv_path, pred_path), callback=callback)
pool.close()
pool.join()
# Create train_all set using symbolic link.
print("Making symlinks to form train_all split... ", end="", flush=True)
trainall_dirname = 'train_all'
trainall_obsv_path = P(output_root).joinpath('{:s}/observation'.format(trainall_dirname))
trainall_obsv_path.mkdir(parents=True, exist_ok=True)
trainall_pred_path = P(output_root).joinpath('{:s}/prediction'.format(trainall_dirname))
trainall_pred_path.mkdir(parents=True, exist_ok=True)
train_path = P(output_root).joinpath('train')
train_obsv_pkl = list(train_path.glob('observation/*.pkl'))
train_pred_pkl = list(train_path.glob('prediction/*.pkl'))
trainval_path = P(output_root).joinpath('train_val')
trainval_obsv_pkl = list(trainval_path.glob('observation/*.pkl'))
trainval_pred_pkl = list(trainval_path.glob('prediction/*.pkl'))
obsv_pkl_list = train_obsv_pkl + trainval_obsv_pkl
pred_pkl_list = train_pred_pkl + trainval_pred_pkl
for obsv_pkl, pred_pkl in zip(obsv_pkl_list, pred_pkl_list):
obsv_filename, obsv_split = obsv_pkl.name, obsv_pkl.parent.parent.stem
pred_filename, pred_split = pred_pkl.name, pred_pkl.parent.parent.stem
obsv_relpath = P('../../{:s}/observation/'.format(obsv_split)).joinpath(obsv_filename)
obsv_link = trainall_obsv_path.joinpath(obsv_filename)
obsv_link.symlink_to(obsv_relpath)
pred_relpath = P('../../{:s}/prediction/'.format(pred_split)).joinpath(pred_filename)
pred_link = trainall_pred_path.joinpath(pred_filename)
pred_link.symlink_to(pred_relpath)
print(" Done.")
if do_maps:
am = ArgoverseMap()
for city_name in ["MIA", "PIT"]:
print("Generating maps for {:s}.".format(city_name))
mask_path = P(output_root).joinpath('raw_map', '{:s}_mask.pkl'.format(city_name))
dt_path = P(output_root).joinpath('raw_map', '{:s}_dt.pkl'.format(city_name))
mask_vis_path = P(output_root).joinpath('raw_map_visualization', '{:s}_mask_vis.png'.format(city_name))
dt_vis_path = P(output_root).joinpath('raw_map_visualization', '{:s}_dt_vis.png'.format(city_name))
mask_vis_path.parent.mkdir(parents=True, exist_ok=True)
mask_path.parent.mkdir(parents=True, exist_ok=True)
map_mask, image_to_city = am.get_rasterized_driveable_area(city_name)
print("Calculating Signed Distance Transform... ", end="", flush=True)
image = map_mask.astype(np.int32)
invert_image = 1-image
dt = np.where(invert_image, -distance_transform_edt(invert_image), distance_transform_edt(image))
print("Done.")
print("Saving Results... ", end="", flush=True)
dump({'map': map_mask, 'image_to_city': image_to_city}, mask_path)
dump({'map': dt, 'image_to_city': image_to_city}, dt_path)
mask_vis = (map_mask*255).astype(np.uint8)
dt_max = dt.max()
dt_min = dt.min()
dt_vis = ((dt - dt_min)/(dt_max - dt_min)*255).astype(np.uint8)
cv2.imwrite(str(mask_vis_path), mask_vis)
cv2.imwrite(str(dt_vis_path), dt_vis)
print("Done. Saved {:s}, {:s}, {:s}, and {:s}.".format(str(mask_path), str(mask_vis_path), str(dt_path), str(dt_vis_path)))