def test_raises_error_when_seconds_negative(self): mock_samples = [{'token': '1', 'timestamp': 0, 'anns': ['1', '1b']}] nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) with self.assertRaises(ValueError): helper.get_future_for_agent('1', '1', -1, False) with self.assertRaises(ValueError): helper.get_past_for_agent('1', '1', -1, False) with self.assertRaises(ValueError): helper.get_past_for_sample('1', -1, False) with self.assertRaises(ValueError): helper.get_future_for_sample('1', -1, False)
def test_get_past_for_agent_in_frame(self, ): mock_samples = [{ 'token': '5', 'timestamp': 0 }, { 'token': '4', 'timestamp': -1e6 }, { 'token': '3', 'timestamp': -2e6 }, { 'token': '2', 'timestamp': -3e6 }, { 'token': '1', 'timestamp': -4e6 }] # Testing we can get the exact amount of past seconds available nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) past = helper.get_past_for_agent('1', '5', 3, True) np.testing.assert_allclose(past, np.array([[1., -1.], [2., -2.], [3., -3.]]))
def test_get_past_for_last_returns_nothing(self): mock_samples = [{'token': '1', 'timestamp': 0}] # Testing we get nothing if we're at the last annotation nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) past = helper.get_past_for_agent('1', '1', 3, False) np.testing.assert_equal(past, np.array([]))
def test_get_no_data_when_seconds_0(self): mock_samples = [{'token': '1', 'timestamp': 0, 'anns': ['1']}] nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) np.testing.assert_equal( helper.get_future_for_agent('1', '1', 0, False), np.array([])) np.testing.assert_equal(helper.get_past_for_agent('1', '1', 0, False), np.array([])) np.testing.assert_equal(helper.get_future_for_sample('1', 0, False), np.array([])) np.testing.assert_equal(helper.get_past_for_sample('1', 0, False), np.array([]))
def test_get_past_for_agent_no_data_to_get(self, ): mock_samples = [{ 'token': '5', 'timestamp': 0 }, { 'token': '4', 'timestamp': -3.5e6 }] # Testing we get nothing if the first sample annotation is past our threshold nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) past = helper.get_past_for_agent('1', '5', 3, False) np.testing.assert_equal(past, np.array([]))
def test_get_past_for_agent_within_buffer(self, ): mock_samples = [{ 'token': '5', 'timestamp': 0 }, { 'token': '4', 'timestamp': -1e6 }, { 'token': '3', 'timestamp': -3.05e6 }, { 'token': '2', 'timestamp': -3.2e6 }] # Testing we get data if it is after future seconds but within buffer nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) past = helper.get_past_for_agent('1', '5', 3, False) np.testing.assert_equal(past, np.array([[3, 3], [2, 2]]))
def test_get_past_for_agent_less_amount(self, ): mock_samples = [{ 'token': '5', 'timestamp': 0 }, { 'token': '4', 'timestamp': -1e6 }, { 'token': '3', 'timestamp': -2.6e6 }, { 'token': '2', 'timestamp': -4e6 }, { 'token': '1', 'timestamp': -5.5e6 }] # Testing we do not include data after the past seconds nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) past = helper.get_past_for_agent('1', '5', 3, False) np.testing.assert_equal(past, np.array([[3, 3], [2, 2]]))
def test_get_past_for_agent_exact_amount(self, ): mock_samples = [{ 'token': '5', 'timestamp': 0 }, { 'token': '4', 'timestamp': -1e6 }, { 'token': '3', 'timestamp': -2e6 }, { 'token': '2', 'timestamp': -3e6 }, { 'token': '1', 'timestamp': -4e6 }] # Testing we can get the exact amount of past seconds available nusc = MockNuScenes(self.mock_annotations, mock_samples) helper = PredictHelper(nusc) past = helper.get_past_for_agent('1', '5', 3, False) np.testing.assert_equal(past, np.array([[3, 3], [2, 2], [1, 1]]))
def compute_metrics(predictions: List[Dict[str, Any]], helper: PredictHelper, config: PredictionConfig) -> Dict[str, Any]: """ Computes metrics from a set of output. :param predictions: List of prediction JSON objects. :param helper: Instance of PredictHelper that wraps the nuScenes val set. :param config: Config file. :return: Metrics. Nested dictionary where keys are metric names and value is a dictionary mapping the Aggregator name to the results. """ n_preds = len(predictions) containers = { metric.name: np.zeros((n_preds, metric.shape)) for metric in config.metrics } containers['imginfo'] = [] containers['tokens'] = [] BACK_SAMPLES = 5 for i, prediction_str in enumerate(tqdm(predictions)): prediction = Prediction.deserialize(prediction_str) ground_truth = helper.get_future_for_agent(prediction.instance, prediction.sample, config.seconds, in_agent_frame=False) agent_past = helper.get_past_for_agent(prediction.instance, prediction.sample, BACK_SAMPLES * 0.5, in_agent_frame=False, just_xy=False) if len(agent_past) < BACK_SAMPLES: print('Sample {} didnt have enough history {}'.format( i, len(agent_past))) continue cam_names = [ sensor['channel'] for sensor in nusc.sensor if 'CAM' in sensor['channel'] ] containers['tokens'].append(prediction_str) cameras = [] containers['imginfo'].append(cameras) # Append current timestep d = {} cameras.append(d) token = prediction.instance + '_' + prediction.sample for cam_name in cam_names: d[cam_name] = get_im_and_box( nusc, token, cam_name=cam_name, imgAsName=True) # impath, box, camera_intrinsics # Append earlier timesteps for t in range(BACK_SAMPLES): d = {} cameras.append(d) token = agent_past[t]['instance_token'] + '_' + agent_past[t][ 'sample_token'] for cam_name in cam_names: d[cam_name] = get_im_and_box( nusc, token, cam_name=cam_name, imgAsName=True) # impath, box, camera_intrinsics for metric in config.metrics: containers[metric.name][i] = metric(ground_truth, prediction) aggregations: Dict[str, Dict[str, List[float]]] = defaultdict(dict) for metric in config.metrics: for agg in metric.aggregators: aggregations[metric.name][agg.name] = agg(containers[metric.name]) return aggregations, containers
class NuScenesFormatTransformer: def __init__(self, DATAROOT='./data/sets/nuscenes', dataset_version='v1.0-mini'): self.DATAROOT = DATAROOT self.dataset_version = dataset_version self.nuscenes = NuScenes(dataset_version, dataroot=self.DATAROOT) self.helper = PredictHelper(self.nuscenes) def get_format_mha_jam(self, samples_agents, out_file="./transformer_format.txt"): instance_token_to_id_dict = {} sample_token_to_id_dict = {} scene_token_dict = {} sample_id = 0 instance_id = 0 for current_sample in tqdm(samples_agents): instance_token, sample_token = current_sample.split("_") scene_token = self.nuscenes.get('sample', sample_token)["scene_token"] if scene_token in scene_token_dict: continue # get the first sample in this sequence scene_token_dict[scene_token] = True first_sample_token = self.nuscenes.get( "scene", scene_token)["first_sample_token"] current_sample = self.nuscenes.get('sample', first_sample_token) while True: if current_sample['token'] not in sample_token_to_id_dict: sample_token_to_id_dict[ current_sample['token']] = sample_id sample_token_to_id_dict[sample_id] = current_sample[ 'token'] sample_id += 1 else: print("should not happen?") instances_in_sample = self.helper.get_annotations_for_sample( current_sample['token']) for sample_instance in instances_in_sample: if sample_instance[ 'instance_token'] not in instance_token_to_id_dict: instance_token_to_id_dict[ sample_instance['instance_token']] = instance_id instance_token_to_id_dict[ instance_id] = sample_instance['instance_token'] instance_id += 1 if current_sample['next'] == "": break current_sample = self.nuscenes.get('sample', current_sample['next']) mode = "train" if out_file.find("_train") != -1 else "val" mini = "mini" if out_file.find("mini") != -1 else "main" with open( "dicts_sample_and_instances_id2token_" + mode + "_" + mini + ".json", 'w') as fw: json.dump([instance_token_to_id_dict, sample_token_to_id_dict], fw) ############# # Converting to the transformer network format # frame_id, agent_id, pos_x, pos_y # todo: # loop on all the agents, if agent not taken: # 1- add it to takens agents (do not retake the agent) # 2- get the number of appearance of this agent # 3- skip this agent if the number is less than 10s (4 + 6) # 4- get the middle agent's token # 5- get the past and future agent's locations relative to its location samples_new_format = [] taken_instances = {} ds_size = 0 # max_past_traj_len = -1 for current_sample in samples_agents: instance_token, sample_token = current_sample.split("_") instance_id = instance_token_to_id_dict[instance_token] if instance_id in taken_instances: continue taken_instances[instance_id] = True # trajectory_full_instances = self.get_trajectory_around_sample(instance_token, sample_token, # just_xy=False) # ////////////////////// future_samples = self.helper.get_future_for_agent( instance_token, sample_token, 6, True, False) past_samples = self.helper.get_past_for_agent( instance_token, sample_token, 1000, True, False)[:MAX_TRAJ_LEN - 1][::-1] current_sample = self.helper.get_sample_annotation( instance_token, sample_token) assert len(past_samples) >= 1 assert len(future_samples) == 12 # assert len(past_samples) < 7 # if len(past_samples) > max_past_traj_len: # max_past_traj_len = len(past_samples) # past_samples = np.append(past_samples, [current_sample], axis=0) ds_size += 1 # get_trajectory at this position center_pos = len(past_samples) future_samples_local = self.helper.get_future_for_agent( instance_token, sample_token, 6, True, True) past_samples_local = self.helper.get_past_for_agent( instance_token, sample_token, 1000, True, True)[:MAX_TRAJ_LEN - 1][::-1] # current_sample = self.helper.get_sample_annotation(instance_token, sample_token) assert len(future_samples_local) == 12 # if len(past_samples) > 7: # past_samples = past_samples[len(past_samples)-7:] # past_samples_local = past_samples_local[past_samples_local.shape[0]-7:] trajectory = np.append(past_samples_local, np.append([[0, 0]], future_samples_local, axis=0), axis=0) past_samples = [ sample_token_to_id_dict[p['sample_token']] for p in past_samples ] future_samples = [ sample_token_to_id_dict[p['sample_token']] for p in future_samples ] trajectory_tokens = np.append( past_samples, np.append([sample_token_to_id_dict[sample_token]], future_samples, axis=0), axis=0) trajectory_ = np.zeros((trajectory.shape[0], 6)) trajectory_[:, 0] = trajectory_tokens[:] trajectory_[:, 1:3] = trajectory trajectory = trajectory_ len_future_samples = len(future_samples) del trajectory_, trajectory_tokens, past_samples, future_samples, past_samples_local, future_samples_local curr_sample = self.helper.get_past_for_agent( instance_token, sample_token, 1000, False, False)[:MAX_TRAJ_LEN][-1] for i in range(trajectory.shape[0]): # instance_id, sample_id, x, y, velocity, acc, yaw velocity = self.helper.get_velocity_for_agent( instance_token, curr_sample["sample_token"]) acceleration = self.helper.get_acceleration_for_agent( instance_token, curr_sample["sample_token"]) heading_change_rate = self.helper.get_heading_change_rate_for_agent( instance_token, curr_sample["sample_token"]) if math.isnan(velocity): velocity = 0 if math.isnan(acceleration): acceleration = 0 if math.isnan(heading_change_rate): heading_change_rate = 0 # need to check paper for relative velocity? same for acc and yaw trajectory[i][3:] = [ velocity, acceleration, heading_change_rate ] # if curr_sample['next'] == '': # import pdb # pdb.set_trace() # No need to get next sample token in case this is last element in the series # prevents bug if i < trajectory.shape[0] - 1: next_sample_token = self.nuscenes.get( 'sample_annotation', curr_sample['next'])['sample_token'] curr_sample = self.helper.get_sample_annotation( instance_token, next_sample_token) s = str(instance_id) + "," # assert (MAX_TRAJ_LEN+len_future_samples) >= trajectory.shape[0] repeat = (MAX_TRAJ_LEN + len_future_samples) - trajectory.shape[0] leading_arr = np.array( repeat * [-1, -64, -64, -64, -64, -64]).reshape((repeat, 6)) trajectory = np.append(leading_arr, trajectory, axis=0) # print("Built In!") # self.nuim.render_trajectory(sample_token, rotation_yaw=0, center_key_pose=True) # print("Bassel's!") # visualize_traffic(trajectory[(trajectory != [-1, -64, -64, -64, -64, -64]).all(axis=1), 1:3].copy()) for i in range(trajectory.shape[0]): sample_id, x, y, velocity, acceleration, heading_change_rate = trajectory[ i] s += str(sample_id) + "," + str(x) + "," + str(y) + "," + str(velocity) + "," \ + str(acceleration) + "," + str(heading_change_rate) if i != trajectory.shape[0] - 1: s += "," else: s += "\n" samples_new_format.append(s) # print("max past trajectory len:",max_past_traj_len) # samples_new_format.sort(key=lambda x: int(x.split(",")[0])) with open(out_file, 'w') as fw: fw.writelines(samples_new_format) print(out_file + "size " + str(ds_size)) def get_format_mha_jam_context(self, states_filepath, out_file): with open(states_filepath) as fr: agents_states = fr.readlines() # format # agent_id, 20x(frame_id, x, y, v, a, yaw_rate)] agents_states = [[float(x.rstrip()) for x in s.split(',')] for s in agents_states] mode = "train" if out_file.find("_train") != -1 else "val" mini = "mini" if out_file.find("mini") != -1 else "main" with open("dicts_sample_and_instances_id2token_" + mode + "_" + mini + ".json") as fr: instance_dict_id_token, sample_dict_id_token = json.load(fr) # Get Context for each sample in states context = [] agent_ind = 0 for agent in tqdm(agents_states): instance_token = instance_dict_id_token[str(int(agent[0]))] mid_frame_id = int(agent[1 + 6 * (MAX_TRAJ_LEN - 1)]) sample_token = sample_dict_id_token[str(mid_frame_id)] frame_annotations = self.helper.get_annotations_for_sample( sample_token) surroundings_agents_coords = [] surroundings_agents_instance_token = [] for ann in frame_annotations: if ann['category_name'].find("vehicle") == -1: continue if ann['instance_token'] == instance_token: agent_ann = ann else: surroundings_agents_coords.append(ann["translation"][:2]) surroundings_agents_instance_token.append( ann["instance_token"]) if len(surroundings_agents_coords) != 0: surroundings_agents_coords = convert_global_coords_to_local( surroundings_agents_coords, agent_ann["translation"], agent_ann["rotation"]) # for i in range(len(surroundings_agents_coords)): # if surroundings_agents_coords[i][0] < -25 or surroundings_agents_coords[i][0] > 25 \ # or surroundings_agents_coords[i][1] < -10 or surroundings_agents_coords[i][1] > 40: # surroundings_agents_coords[i] = None # surroundings_agents_instance_token[i] = None total_area_side = 50 cell_size = 1.5625 map_side_size = int(total_area_side // cell_size) map = [[[-64, -64, -64, -64, -64] for i in range(MAX_TRAJ_LEN)] for j in range(map_side_size * map_side_size)] for n in range(len(surroundings_agents_coords)): # if np.isnan(surroundings_agents_coords[n][0]): # ---> surroundings_agents_coords[n] is None # continue # search for the agent location in the map # agent_found = False # for i in range(map_side_size): # for j in range(map_side_size): # # if agent found in the cell # if surroundings_agents_coords[n][0] >= (j * cell_size) - 25\ # and surroundings_agents_coords[n][0] < (j * cell_size) - 25 + cell_size \ # and surroundings_agents_coords[n][1] < 40 - (i * cell_size) \ # and surroundings_agents_coords[n][1] > 40 - (i * cell_size + cell_size): # found_i, found_j = i, j # break # get the agent location in the map! alpha_y = (surroundings_agents_coords[n][1] - (-10)) / (40 - (-10)) i = (map_side_size - 1) - int(alpha_y * map_side_size + 0) alpha_x = (surroundings_agents_coords[n][0] - (-25)) / (25 - (-25)) j = int(alpha_x * map_side_size + 0) # Confirmation the 2 methods yield the same results # if not(found_i == i and found_j == j): # raise Exception("Calculations error") # prevent out of bound cases (which shall never happen if none is set for out of bound (line 240) if not (i >= 0 and i < map_side_size and j >= 0 and j < map_side_size): # raise Exception("Calculations error") continue pos = i * map_side_size + j past_trajectory = self.get_current_past_trajectory( surroundings_agents_instance_token[n], sample_token, num_seconds=1000)[:MAX_TRAJ_LEN] assert len(past_trajectory) <= MAX_TRAJ_LEN retrieved_trajectory_len = len(past_trajectory) if map[pos][-1][0] != -64: skip_traj = False # Save the trajectory with greater length for ind, map_pos in enumerate(map[pos]): if map_pos[0] != 64: if MAX_TRAJ_LEN - ind > retrieved_trajectory_len: skip_traj = True if skip_traj: agent_found = True break else: # print("new longer agent trajectory in cell") pass past_trajectory = convert_global_coords_to_local( past_trajectory, agent_ann["translation"], agent_ann["rotation"]) if retrieved_trajectory_len != MAX_TRAJ_LEN: past_trajectory = np.concatenate([ np.array([[-64, -64] for _ in range(MAX_TRAJ_LEN - past_trajectory.shape[0])]), past_trajectory ], axis=0) neighbour_agent_features = [] skip_traj = False for k in range(0, MAX_TRAJ_LEN): if retrieved_trajectory_len > k: if k == 0: sample_token_i = sample_dict_id_token[str( mid_frame_id)] else: sample_token_i = self.helper.get_sample_annotation( surroundings_agents_instance_token[n], sample_token_i)["prev"] sample_token_i = self.nuscenes.get( 'sample_annotation', sample_token_i)['sample_token'] try: velocity = self.helper.get_velocity_for_agent( surroundings_agents_instance_token[n], sample_token_i) except: skip_traj = True # print("error") break acceleration = self.helper.get_acceleration_for_agent( surroundings_agents_instance_token[n], sample_token_i) heading_change_rate = self.helper.get_heading_change_rate_for_agent( surroundings_agents_instance_token[n], sample_token_i) if math.isnan(velocity): velocity = 0 if math.isnan(acceleration): acceleration = 0 if math.isnan(heading_change_rate): heading_change_rate = 0 neighbour_agent_features.append( [velocity, acceleration, heading_change_rate]) else: neighbour_agent_features.append([-64, -64, -64]) if skip_traj: print("skipping agent because it has missing data") agent_found = True break past_trajectory = np.concatenate( [past_trajectory, neighbour_agent_features], axis=1) map[pos] = past_trajectory.tolist() # agent_found = True # break # if agent_found: # break map = np.array(map).astype(np.float16) if VISUALIZE_DATA: visualize_traffic_neighbours(map, map_side_size * map_side_size) # context.append(map) if not os.path.exists(os.path.dirname(out_file)): os.makedirs(os.path.dirname(out_file)) np.save(out_file.replace("_.txt", "__" + str(agent_ind) + ".txt"), map) agent_ind += 1 # with open(out_file, 'ab') as fw: # pickle.dump(map, fw) # continue # fw.write(map) def get_current_past_trajectory(self, instance_token, sample_token, num_seconds, just_xy=True, in_agent_frame=False): past_samples = self.helper.get_past_for_agent( instance_token, sample_token, num_seconds, in_agent_frame, just_xy)[::-1] #[0:7][::-1] current_sample = self.helper.get_sample_annotation( instance_token, sample_token) if just_xy: current_sample = current_sample["translation"][:2] if past_samples.shape[0] == 0: trajectory = np.array([current_sample]) else: trajectory = np.append(past_samples, [current_sample], axis=0) else: trajectory = np.append(past_samples, [current_sample], axis=0) return trajectory def get_format_mha_jam_maps(self, states_filepath, out_file): with open(states_filepath) as fr: agents_states = fr.readlines() # format # agen t_id, 20x(frame_id, x, y, v, a, yaw_rate)] agents_states = [[float(x.rstrip()) for x in s.split(',')] for s in agents_states] mode = "train" if out_file.find("_train") != -1 else "val" mini = "mini" if out_file.find("mini") != -1 else "main" with open("dicts_sample_and_instances_id2token_" + mode + "_" + mini + ".json") as fr: instance_dict_id_token, sample_dict_id_token = json.load(fr) # Get map for each sample in states agent_ind = 0 static_layer_rasterizer = StaticLayerRasterizer(self.helper) agent_rasterizer = AgentBoxesWithFadedHistory(self.helper, seconds_of_history=1) mtp_input_representation = InputRepresentation(static_layer_rasterizer, agent_rasterizer, Rasterizer()) if not os.path.exists(os.path.dirname(out_file)): os.makedirs(os.path.dirname(out_file)) for agent in tqdm(agents_states): instance_token = instance_dict_id_token[str(int(agent[0]))] mid_frame_id = int(agent[1 + 6 * (MAX_TRAJ_LEN)]) sample_token = sample_dict_id_token[str(mid_frame_id)] img = mtp_input_representation.make_input_representation( instance_token, sample_token) # img = cv2.resize(img, (1024, 1024)) cv2.imwrite( out_file.replace("_.jpg", "__" + str(agent_ind) + ".jpg"), img) agent_ind += 1 def run(self, out_dir): if self.dataset_version.find("mini") != -1: train_agents = get_prediction_challenge_split( "mini_train", dataroot=self.DATAROOT) val_agents = get_prediction_challenge_split("mini_val", dataroot=self.DATAROOT) else: train_agents = get_prediction_challenge_split( "train", dataroot=self.DATAROOT) train_agents.extend( get_prediction_challenge_split("train_val", dataroot=self.DATAROOT)) val_agents = get_prediction_challenge_split("val", dataroot=self.DATAROOT) ## Statistics # mx =-1 # for in train_agents: # instance_token, sample_token = current_sample.split("_") # past_samples_local = self.helper.get_past_for_agent(instance_token, sample_token, 100, True, True)[::-1] # if len(past_samples_local) > mx: # mx = len(past_samples_local) # print("max length of the past sequences for trainval is:",mx) # for instance_token, sample_token in train_agents: # past_samples_local = self.helper.get_past_for_agent(instance_token, sample_token, 100, True, True)[::-1] # if len(past_samples_local) > mx: # mx = len(past_samples_local) # print("max length of the past sequence for val is:",mx) # return self.get_format_mha_jam( train_agents, os.path.join(out_dir, "states_train_" + self.dataset_version + ".txt")) self.get_format_mha_jam_context( os.path.join(out_dir, "states_train_" + self.dataset_version + ".txt"), os.path.join(out_dir, "context_train_" + self.dataset_version, "context_train_.txt")) self.get_format_mha_jam_maps( os.path.join(out_dir, "states_train_" + self.dataset_version + ".txt"), os.path.join(out_dir, "maps_train_" + self.dataset_version, "maps_train_.jpg")) # 25 self.get_format_mha_jam( val_agents, os.path.join(out_dir, "states_val_" + self.dataset_version + ".txt")) self.get_format_mha_jam_context( os.path.join(out_dir, "states_val_" + self.dataset_version + ".txt"), os.path.join(out_dir, "context_val_" + self.dataset_version, "context_val_.txt")) self.get_format_mha_jam_maps( os.path.join(out_dir, "states_val_" + self.dataset_version + ".txt"), os.path.join(out_dir, "maps_val_" + self.dataset_version, "maps_val_.jpg"))
class NS(Dataset): def __init__(self, dataroot: str, split: str, t_h: float = 2, t_f: float = 6, grid_dim: int = 25, img_size: int = 200, horizon: int = 40, grid_extent: Tuple[int, int, int, int] = (-25, 25, -10, 40), num_actions: int = 4, image_extraction_mode: bool = False): """ Initializes dataset class for nuScenes prediction :param dataroot: Path to tables and data :param split: Dataset split for prediction benchmark ('train'/'train_val'/'val') :param t_h: Track history in seconds :param t_f: Prediction horizon in seconds :param grid_dim: Size of grid, default: 25x25 :param img_size: Size of raster map image in pixels, default: 200x200 :param horizon: MDP horizon :param grid_extent: Map extents in meters, (-left, right, -behind, front) :param num_actions: Number of actions for each state (4: [D,R,U,L] or 8: [D, R, U, L, DR, UR, DL, UL]) :param image_extraction_mode: Whether dataset class is being used for image extraction """ # Nuscenes dataset and predict helper self.dataroot = dataroot self.ns = NuScenes('v1.0-trainval', dataroot=dataroot) self.helper = PredictHelper(self.ns) self.token_list = get_prediction_challenge_split(split, dataroot=dataroot) # Useful parameters self.grid_dim = grid_dim self.grid_extent = grid_extent self.img_size = img_size self.t_f = t_f self.t_h = t_h self.horizon = horizon self.num_actions = num_actions # Map row, column and velocity states to actual values grid_size_m = self.grid_extent[1] - self.grid_extent[0] self.row_centers = np.linspace( self.grid_extent[3] - grid_size_m / (self.grid_dim * 2), self.grid_extent[2] + grid_size_m / (self.grid_dim * 2), self.grid_dim) self.col_centers = np.linspace( self.grid_extent[0] + grid_size_m / (self.grid_dim * 2), self.grid_extent[1] - grid_size_m / (self.grid_dim * 2), self.grid_dim) # Surrounding agent input representation: populate grid with velocity, acc, yaw-rate self.agent_ip = AgentMotionStatesOnGrid(self.helper, resolution=grid_size_m / img_size, meters_ahead=grid_extent[3], meters_behind=-grid_extent[2], meters_left=-grid_extent[0], meters_right=grid_extent[1]) # Image extraction mode is used for extracting map images offline prior to training self.image_extraction_mode = image_extraction_mode if self.image_extraction_mode: # Raster map representation self.map_ip = StaticLayerRasterizer(self.helper, resolution=grid_size_m / img_size, meters_ahead=grid_extent[3], meters_behind=-grid_extent[2], meters_left=-grid_extent[0], meters_right=grid_extent[1]) # Raster map with agent boxes. Only used for visualization static_layer_rasterizer = StaticLayerRasterizer( self.helper, resolution=grid_size_m / img_size, meters_ahead=grid_extent[3], meters_behind=-grid_extent[2], meters_left=-grid_extent[0], meters_right=grid_extent[1]) agent_rasterizer = AgentBoxesWithFadedHistory( self.helper, seconds_of_history=1, resolution=grid_size_m / img_size, meters_ahead=grid_extent[3], meters_behind=-grid_extent[2], meters_left=-grid_extent[0], meters_right=grid_extent[1]) self.map_ip_agents = InputRepresentation(static_layer_rasterizer, agent_rasterizer, Rasterizer()) def __len__(self): return len(self.token_list) def __getitem__(self, idx): """ Returns inputs, ground truth values and other utilities for data point at given index :return hist: snippet of track history, default 2s at 0.5 Hz sampling frequency :return fut: ground truth future trajectory, default 6s at 0.5 Hz sampling frequency :return img: Imagenet normalized bird's eye view map around the target vehicle :return svf_e: Goal and path state visitation frequencies for expert demonstration, ie. path from train set :return motion_feats: motion and position features used for reward model :return waypts_e: (x,y) BEV co-ordinates corresponding to grid cells of svf_e :return agents: tensor of surrounding agent states populated in grid around target agent :return grid_idcs: grid co-ordinates of svf_e :return bc_targets: ground truth actions for training behavior cloning model :return img_agents: image with agent boxes for visualization / debugging :return instance_token: nuScenes instance token for prediction instance :return sample_token: nuScenes sample token for prediction instance :return idx: instance id (mainly for debugging) """ # Nuscenes instance and sample token for prediction data point instance_token, sample_token = self.token_list[idx].split("_") # If dataset is being used for image extraction grid_size_m = self.grid_extent[1] - self.grid_extent[0] if self.image_extraction_mode: # Make directory to store raster map images img_dir = os.path.join( self.dataroot, 'prediction_raster_maps', 'images' + str(self.img_size) + "_" + str(int(grid_size_m)) + 'm') if not os.path.isdir(img_dir): os.mkdir(img_dir) # Generate and save raster map image with just static elements img = self.map_ip.make_representation(instance_token, sample_token) img_save = Image.fromarray(img) img_save.save( os.path.join(img_dir, instance_token + "_" + sample_token + '.png')) # Generate and save raster map image with static elements and agent boxes (for visualization only) img_agents = self.map_ip_agents.make_input_representation( instance_token, sample_token) img_agents_save = Image.fromarray(img_agents) img_agents_save.save( os.path.join( img_dir, instance_token + "_" + sample_token + 'agents.png')) # Return dummy values return 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 # If dataset is being used for training/validation/testing else: # Get track history for agent: hist = self.get_hist(instance_token, sample_token) hist = torch.from_numpy(hist) # Get ground truth future for agent: fut = self.helper.get_future_for_agent(instance_token, sample_token, seconds=self.t_f, in_agent_frame=True) fut = torch.from_numpy(fut) # Get indefinite future for computing expert State visitation frequencies (SVF): fut_indefinite = self.helper.get_future_for_agent( instance_token, sample_token, seconds=300, in_agent_frame=True) # Up sample indefinite future by a factor of 10 fut_interpolated = np.zeros((fut_indefinite.shape[0] * 10 + 1, 2)) param_query = np.linspace(0, fut_indefinite.shape[0], fut_indefinite.shape[0] * 10 + 1) param_given = np.linspace(0, fut_indefinite.shape[0], fut_indefinite.shape[0] + 1) val_given_x = np.concatenate(([0], fut_indefinite[:, 0])) val_given_y = np.concatenate(([0], fut_indefinite[:, 1])) fut_interpolated[:, 0] = np.interp(param_query, param_given, val_given_x) fut_interpolated[:, 1] = np.interp(param_query, param_given, val_given_y) # Read pre-extracted raster map image img_dir = os.path.join( self.dataroot, 'prediction_raster_maps', 'images' + str(self.img_size) + "_" + str(int(grid_size_m)) + 'm') img = cv2.imread( os.path.join(img_dir, instance_token + "_" + sample_token + '.png')) # Pre-process image img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = torch.from_numpy(img) img = img.permute((2, 0, 1)).float() / 255 # Normalize using Imagenet stats img = normalize_imagenet(img) # Read pre-extracted raster map with agent boxes (for visualization + debugging) img_agents = cv2.imread( os.path.join( img_dir, instance_token + "_" + sample_token + 'agents.png')) # Pre-process image img_agents = cv2.cvtColor(img_agents, cv2.COLOR_BGR2RGB) img_agents = torch.from_numpy(img_agents) img_agents = img_agents.permute((2, 0, 1)).float() / 255 # Get surrounding agent states agents = torch.from_numpy( self.agent_ip.make_representation(instance_token, sample_token)) agents = agents.permute((2, 0, 1)).float() # Sum pool states to down-sample to grid dimensions agents = f.avg_pool2d(agents[None, :, :, :], self.img_size // self.grid_dim) agents = agents.squeeze(dim=0) * ( (self.img_size // self.grid_dim)**2) # Get expert SVF: svf_e, waypts_e, grid_idcs = self.get_expert_waypoints( fut_interpolated) svf_e = torch.from_numpy(svf_e) waypts_e = torch.from_numpy(waypts_e) grid_idcs = torch.from_numpy(grid_idcs) # Get motion and position feats: motion_feats = self.get_motion_feats(instance_token, sample_token) motion_feats = torch.from_numpy(motion_feats) # Targets for behavior cloning model: bc_targets = self.get_bc_targets(fut_interpolated) bc_targets = torch.from_numpy(bc_targets) return hist, fut, img, svf_e, motion_feats, waypts_e, agents, grid_idcs, bc_targets, img_agents, \ instance_token, sample_token, idx def get_hist(self, instance_token: str, sample_token: str): """ Function to get track history of agent :param instance_token: nuScenes instance token for datapoint :param sample_token nuScenes sample token for datapoint """ # x, y co-ordinates in agent's frame of reference xy = self.helper.get_past_for_agent(instance_token, sample_token, seconds=self.t_h, in_agent_frame=True) # Get all history records for obtaining velocity, acceleration and turn rate values hist_records = self.helper.get_past_for_agent(instance_token, sample_token, seconds=self.t_h, in_agent_frame=True, just_xy=False) if xy.shape[0] > self.t_h * 2: xy = xy[0:int(self.t_h) * 2] if len(hist_records) > self.t_h * 2: hist_records = hist_records[0:int(self.t_h) * 2] # Initialize hist tensor and set x and y co-ordinates returned by prediction helper hist = np.zeros((xy.shape[0], 5)) hist[:, 0:2] = xy # Instance and sample tokens from history records i_tokens = [ hist_records[i]['instance_token'] for i in range(len(hist_records)) ] i_tokens.insert(0, instance_token) s_tokens = [ hist_records[i]['sample_token'] for i in range(len(hist_records)) ] s_tokens.insert(0, sample_token) # Set velocity, acc and turn rate values for hist for k in range(hist.shape[0]): i_t = i_tokens[k] s_t = s_tokens[k] v = self.helper.get_velocity_for_agent(i_t, s_t) a = self.helper.get_acceleration_for_agent(i_t, s_t) theta = self.helper.get_heading_change_rate_for_agent(i_t, s_t) # If function returns nan values due to short tracks, set corresponding value to 0 if np.isnan(v): v = 0 if np.isnan(a): a = 0 if np.isnan(theta): theta = 0 hist[k, 2] = v hist[k, 3] = a hist[k, 4] = theta # Zero pad for track histories shorter than t_h hist_zeropadded = np.zeros((int(self.t_h) * 2, 5)) # Flip to have correct order of timestamps hist = np.flip(hist, 0) hist_zeropadded[-hist.shape[0]:] = hist return hist_zeropadded def get_expert_waypoints(self, fut: np.ndarray): """ Function to get the expert's state visitation frequencies based on their trajectory :param fut: numpy array with future trajectory of for all available future timestamps, up-sampled by 10 """ # Expert state visitation frequencies for training reward model, waypoints in meters and grid indices svf_e = np.zeros((2, self.grid_dim, self.grid_dim)) waypts_e = np.zeros((self.horizon, 2)) grid_idcs = np.zeros((self.horizon, 2)) count = 0 row_prev = np.nan column_prev = np.nan for k in range(fut.shape[0]): # Convert trajectory (x,y) co-ordinates to grid locations: column = np.argmin(np.absolute(fut[k, 0] - self.col_centers)) row = np.argmin(np.absolute(fut[k, 1] - self.row_centers)) # Demonstration ends when expert leaves the image crop corresponding to the grid: if self.grid_extent[0] <= fut[k, 0] <= self.grid_extent[1] and \ self.grid_extent[2] <= fut[k, 1] <= self.grid_extent[3]: # Check if cell location has changed if row != row_prev or column != column_prev: # Add cell location to path states of expert svf_e[0, row.astype(int), column.astype(int)] = 1 if count < self.horizon: # Get BEV coordinates corresponding to cell locations waypts_e[count, 0] = self.row_centers[row] waypts_e[count, 1] = self.col_centers[column] grid_idcs[count, 0] = row grid_idcs[count, 1] = column count += 1 else: break column_prev = column row_prev = row # Last cell location where demonstration terminates is the goal state: svf_e[1, row_prev.astype(int), column_prev.astype(int)] = 1 return svf_e, waypts_e, grid_idcs def get_motion_feats(self, instance_token: str, sample_token: str): """ Function to get motion and position features over grid for reward model :param instance_token: NuScenes instance token for datapoint :param sample_token: NuScenes sample token for datapoint """ feats = np.zeros((3, self.grid_dim, self.grid_dim)) # X and Y co-ordinates over grid grid_size_m = self.grid_extent[1] - self.grid_extent[0] y = (np.linspace( self.grid_extent[3] - grid_size_m / (self.grid_dim * 2), self.grid_extent[2] + grid_size_m / (self.grid_dim * 2), self.grid_dim)).reshape(-1, 1).repeat(self.grid_dim, axis=1) x = (np.linspace( self.grid_extent[0] + grid_size_m / (self.grid_dim * 2), self.grid_extent[1] - grid_size_m / (self.grid_dim * 2), self.grid_dim)).reshape(-1, 1).repeat(self.grid_dim, axis=1).transpose() # Velocity of agent v = self.helper.get_velocity_for_agent(instance_token, sample_token) if np.isnan(v): v = 0 # Normalize X and Y co-ordinates over grid feats[0] = v feats[1] = x / grid_size_m feats[2] = y / grid_size_m return feats def get_bc_targets(self, fut: np.ndarray): """ Function to get targets for behavior cloning model :param fut: numpy array with future trajectory of for all available future timestamps, up-sampled by 10 """ bc_targets = np.zeros( (self.num_actions + 1, self.grid_dim, self.grid_dim)) column_prev = np.argmin(np.absolute(fut[0, 0] - self.col_centers)) row_prev = np.argmin(np.absolute(fut[0, 1] - self.row_centers)) for k in range(fut.shape[0]): # Convert trajectory (x,y) co-ordinates to grid locations: column = np.argmin(np.absolute(fut[k, 0] - self.col_centers)) row = np.argmin(np.absolute(fut[k, 1] - self.row_centers)) # Demonstration ends when expert leaves the image crop corresponding to the grid: if self.grid_extent[0] <= fut[k, 0] <= self.grid_extent[1] and self.grid_extent[2] <= fut[k, 1] <= \ self.grid_extent[3]: # Check if cell location has changed if row != row_prev or column != column_prev: bc_targets[:, int(row_prev), int(column_prev)] = 0 d_x = column - column_prev d_y = row - row_prev theta = np.arctan2(d_y, d_x) # Assign ground truth actions for expert demonstration if self.num_actions == 4: # [D,R,U,L,end] if np.pi / 4 <= theta < 3 * np.pi / 4: bc_targets[0, int(row_prev), int(column_prev)] = 1 elif -np.pi / 4 <= theta < np.pi / 4: bc_targets[1, int(row_prev), int(column_prev)] = 1 elif -3 * np.pi / 4 <= theta < -np.pi / 4: bc_targets[2, int(row_prev), int(column_prev)] = 1 else: bc_targets[3, int(row_prev), int(column_prev)] = 1 else: # [D, R, U, L, DR, UR, DL, UL, end] if 3 * np.pi / 8 <= theta < 5 * np.pi / 8: bc_targets[0, int(row_prev), int(column_prev)] = 1 elif -np.pi / 8 <= theta < np.pi / 8: bc_targets[1, int(row_prev), int(column_prev)] = 1 elif -5 * np.pi / 8 <= theta < -3 * np.pi / 8: bc_targets[2, int(row_prev), int(column_prev)] = 1 elif np.pi / 8 <= theta < 3 * np.pi / 8: bc_targets[4, int(row_prev), int(column_prev)] = 1 elif -3 * np.pi / 8 <= theta < -np.pi / 8: bc_targets[5, int(row_prev), int(column_prev)] = 1 elif 5 * np.pi / 8 <= theta < 7 * np.pi / 8: bc_targets[6, int(row_prev), int(column_prev)] = 1 elif -7 * np.pi / 8 <= theta < -5 * np.pi / 8: bc_targets[7, int(row_prev), int(column_prev)] = 1 else: bc_targets[3, int(row_prev), int(column_prev)] = 1 else: break column_prev = column row_prev = row # Final action is the end action to transition to the goal state: bc_targets[self.num_actions, int(row_prev), int(column_prev)] = 1 return bc_targets
class NuScenesFormatTransformer: def __init__(self, DATAROOT='./data/sets/nuscenes', dataset_version='v1.0-mini'): self.DATAROOT = DATAROOT self.dataset_version = dataset_version self.nuscenes = NuScenes(dataset_version, dataroot=self.DATAROOT) self.helper = PredictHelper(self.nuscenes) # ['vehicle.car', 'vehicle.truck', 'vehicle.bus.rigid', 'vehicle.bus.bendy', 'vehicle.construction'] self.category_token_to_id = { "fd69059b62a3469fbaef25340c0eab7f": 1, # 'vehicle.car' "6021b5187b924d64be64a702e5570edf": 1, # 'vehicle.truck' "fedb11688db84088883945752e480c2c": 2, # 'vehicle.bus.rigid' "003edbfb9ca849ee8a7496e9af3025d4": 2, # 'vehicle.bus.bendy' "5b3cd6f2bca64b83aa3d0008df87d0e4": 3, # 'vehicle.construction' "7b2ff083a64e4d53809ae5d9be563504": 1 } # vehicle.emergency.police def get_new_format(self, samples_agents, format_for_model, out_file=("./transformer_format.txt"), num_seconds=None): # for current_sample in samples_agents: # instance_token, sample_token = current_sample.split("_") # traj = self.helper.get_future_for_agent(instance_token, sample_token, 6, True) # past_traj = self.helper.get_past_for_agent(instance_token, sample_token, 6, True) # # if len(past_traj) + len(traj) + 1 < 20: # print(len(past_traj) + len(traj) + 1) # # exit() #################### # Sample Token (frame) to a sequential id # for each sample (agent_frame), get the scene it belongs to, and then get the first sample (frame) # loop on all samples from the first sample till the end # set to dictionary the sequential id for each sample splitting_format = '\t' if format_for_model.value == FORMAT_FOR_MODEL.TRAFFIC_PREDICT.value: splitting_format = " " instance_token_to_id_dict = {} sample_token_to_id_dict = {} scene_token_dict = {} sample_id = 0 instance_id = 0 for current_sample in samples_agents: instance_token, sample_token = current_sample.split("_") scene_token = self.nuscenes.get('sample', sample_token)["scene_token"] if scene_token in scene_token_dict: continue # get the first sample in this sequence scene_token_dict[scene_token] = True first_sample_token = self.nuscenes.get( "scene", scene_token)["first_sample_token"] current_sample = self.nuscenes.get('sample', first_sample_token) while True: if current_sample['token'] not in sample_token_to_id_dict: sample_token_to_id_dict[ current_sample['token']] = sample_id sample_id += 1 else: print("should not happen?") instances_in_sample = self.helper.get_annotations_for_sample( current_sample['token']) for sample_instance in instances_in_sample: if sample_instance[ 'instance_token'] not in instance_token_to_id_dict: instance_token_to_id_dict[ sample_instance['instance_token']] = instance_id instance_id += 1 if current_sample['next'] == "": break current_sample = self.nuscenes.get('sample', current_sample['next']) ############# # Converting to the transformer network format # frame_id, agent_id, pos_x, pos_y # todo: # loop on all the agents, if agent not taken: # 1- add it to takens agents (do not retake the agent) # 2- get the number of appearance of this agent # 3- skip this agent if the number is less than 10s (4 + 6) # 4- get the middle agent's token # 5- get the past and future agent's locations relative to its location samples_new_format = [] taken_instances = {} ds_size = 0 for current_sample in samples_agents: instance_token, sample_token = current_sample.split("_") instance_id, sample_id = instance_token_to_id_dict[ instance_token], sample_token_to_id_dict[sample_token] if instance_id in taken_instances: continue taken_instances[instance_id] = True trajectory = self.get_trajectory_around_sample( instance_token, sample_token) trajectory_full_instances = self.get_trajectory_around_sample( instance_token, sample_token, just_xy=False) # traj_samples_token = [instance['sample_token'] for instance in trajectory_full_instances] if len(trajectory) < 20: print("length is less than 20 samples, trajectory length is: ", len(trajectory)) continue ds_size += 1 if num_seconds is not None: start, end = len(trajectory) // 2 - 9, len( trajectory) // 2 + 11, starting_frame = (start + end) // 2 middle_sample_token = trajectory_full_instances[ starting_frame]["sample_token"] trajectory = self.get_trajectory_around_sample( instance_token, middle_sample_token, just_xy=True, num_seconds=num_seconds, in_agent_frame=True) trajectory_full_instances = self.get_trajectory_around_sample( instance_token, middle_sample_token, just_xy=False, num_seconds=num_seconds, in_agent_frame=True) # traj_samples_token = [instance['sample_token'] for instance in trajectory_full_instances] # get_trajectory at this position for i in range(trajectory.shape[0]): traj_sample, sample_token = trajectory[ i], trajectory_full_instances[i]["sample_token"] sample_id = sample_token_to_id_dict[sample_token] if format_for_model.value == FORMAT_FOR_MODEL.TRANSFORMER_NET.value: yaw = quaternion_yaw( Quaternion(trajectory_full_instances[i]["rotation"])) # samples_new_format.append(str(sample_id) + splitting_format + str(instance_id)\ # + splitting_format + str(traj_sample[0]) + splitting_format \ # + str(traj_sample[1]) + splitting_format + str(yaw) + "\n") x, y, z = trajectory_full_instances[i]["translation"] w, l, h = trajectory_full_instances[i]["size"] samples_new_format.append( str(sample_id) + splitting_format + str(instance_id ) #+ splitting_format + str(object_type)\ + splitting_format + str(x) + splitting_format + str(y) #+ splitting_format + str(z) # + splitting_format + str(l) + splitting_format + str(w) + splitting_format + str(h) + splitting_format + str(yaw) + "\n") elif format_for_model.value == FORMAT_FOR_MODEL.TRAFFIC_PREDICT.value: # raise Exception("not implemented yet") category_token = self.nuscenes.get( "instance", instance_token)["category_token"] object_type = self.category_token_to_id[category_token] # frame_id, object_id, object_type, # position_x, position_y, position_z, # object_length, object_width, object_height, # heading x, y, z = trajectory_full_instances[i]["translation"] w, l, h = trajectory_full_instances[i]["size"] # yaw = angle_of_rotation(quaternion_yaw(Quaternion(trajectory_full_instances[i]["rotation"]))) yaw = quaternion_yaw( Quaternion(trajectory_full_instances[i]["rotation"])) samples_new_format.append(str(sample_id) + splitting_format + str(instance_id) + splitting_format + str(object_type)\ + splitting_format + str(x) + splitting_format + str(y) + splitting_format + str(z) + splitting_format + splitting_format + str(l) + splitting_format + str(w) + splitting_format + str(h) + splitting_format + str(yaw) + "\n") # annotations = helper.get_annotations_for_sample(sample_token) # for ann in annotations: # # q = ann['rotation'] # # yaw = math.atan2(2.0 * (q[3] * q[0] + q[1] * q[2]), - 1.0 + 2.0 * (q[0] * q[0] + q[1] * q[1]))*180/math.pi # # if yaw < 0: # # yaw += 360 # # selected_sample_data = [sample_id, instance_id] + ann['translation'] + [yaw] + ann['size'] # selected_sample_data = str(sample_id) + " " + str(instance_token_to_id_dict[ann['instance_token']])\ # + " " + str(ann['translation'][0]) + " " + str(ann['translation'][2]) + "\n" # samples_new_format.append(selected_sample_data) # no need for sorting as it occurs in the TransformationNet data loader # left it for similarity samples_new_format.sort( key=lambda x: int(x.split(splitting_format)[0])) with open(out_file, 'w') as fw: fw.writelines(samples_new_format) print(out_file + "size " + str(ds_size)) def run(self, format_for_model): if self.dataset_version.find("mini") != -1: train_agents = get_prediction_challenge_split( "mini_train", dataroot=self.DATAROOT) val_agents = get_prediction_challenge_split("mini_val", dataroot=self.DATAROOT) else: train_agents = get_prediction_challenge_split( "train", dataroot=self.DATAROOT) train_agents.extend( get_prediction_challenge_split("train_val", dataroot=self.DATAROOT)) val_agents = get_prediction_challenge_split("val", dataroot=self.DATAROOT) self.get_new_format( train_agents, format_for_model, "/home/bassel/PycharmProjects/Trajectory-Transformer/datasets/nuscenes/bkup/transformer_train_" + self.dataset_version + ".txt") self.get_new_format( val_agents, format_for_model, "/home/bassel/PycharmProjects/Trajectory-Transformer/datasets/nuscenes/bkup/transformer_val_" + self.dataset_version + ".txt") # shutil.copy("/home/bassel/PycharmProjects/Trajectory-Transformer/datasets/nuscenes/val/transformer_val_"+self.dataset_version+".txt", # "/home/bassel/PycharmProjects/Trajectory-Transformer/datasets/nuscenes/test/transformer_val_"+self.dataset_version+".txt") def get_trajectory_around_sample(self, instance_token, sample_token, just_xy=True, num_seconds=1000, in_agent_frame=False): future_samples = self.helper.get_future_for_agent( instance_token, sample_token, num_seconds, in_agent_frame, just_xy) past_samples = self.helper.get_past_for_agent(instance_token, sample_token, num_seconds, in_agent_frame, just_xy)[::-1] current_sample = self.helper.get_sample_annotation( instance_token, sample_token) if num_seconds == 5: if len(past_samples) > 9: past_samples = past_samples[0:9] if len(future_samples) > 10: future_samples = future_samples[0:10] if just_xy: current_sample = current_sample["translation"][:2] trajectory = np.append(past_samples, np.append([current_sample], future_samples, axis=0), axis=0) else: trajectory = np.append(past_samples, np.append([current_sample], future_samples, axis=0), axis=0) return trajectory