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)))