def load_map_data(dataroot, location): # Load the NuScenes map object nusc_map = NuScenesMap(dataroot, location) map_data = OrderedDict() for layer in nusc_utils.STATIC_CLASSES: # Retrieve all data associated with the current layer records = getattr(nusc_map, layer) polygons = list() # Drivable area records can contain multiple polygons if layer == 'drivable_area': for record in records: # Convert each entry in the record into a shapely object for token in record['polygon_tokens']: poly = nusc_map.extract_polygon(token) if poly.is_valid: polygons.append(poly) else: for record in records: # Convert each entry in the record into a shapely object poly = nusc_map.extract_polygon(record['polygon_token']) if poly.is_valid: polygons.append(poly) # Store as an R-Tree for fast intersection queries map_data[layer] = STRtree(polygons) return map_data
def scenario_filter(scene_df, keep_senarios=[]): nusc_map = NuScenesMap(dataroot=full_path, map_name=scene_df.iloc[0]['scene_location']) keep_instance_token_list = [] processed_instance_token_list = [] for i, r in scene_df.iterrows(): if r.instance_token in processed_instance_token_list: continue processed_instance_token_list.append(r.instance_token) instance_road_object_traj = r.past_instance_on_road_objects + [ r.current_instance_on_road_objects ] + r.future_instance_on_road_objects for road_objects in instance_road_object_traj: # filter ados that passes through intersections # if 'intersection' in keep_senarios: if road_objects['road_segment'] is not "": ro = nusc_map.get('road_segment', road_objects['road_segment']) if ro['is_intersection']: keep_instance_token_list.append(r.instance_token) continue #filtered_df = scene_df.drop(del_idx).reset_index(drop=True) filtered_df = scene_df[scene_df.instance_token.str.contains( '|'.join(keep_instance_token_list))] return filtered_df
def convert_map(args: argparse.Namespace) -> None: """ Core function of this script. Loads each json file, calls the helper functions to extract relevant information from the json files, and finally creates xml files in the argoverse format. """ # Load json file. Loop over all the json files in directory for filename in [ "boston-seaport", "singapore-hollandvillage", "singapore-onenorth", "singapore-queenstown", ]: print(filename) with open(f"/{args.nuscenes_dir}/maps/expansion/{filename}.json") as f: data = json.load(f) nusc_map = NuScenesMap(dataroot=f"/{args.nuscenes_dir}", map_name=filename) # Create new xml ETree root = ET.Element("NuScenesMap") # Map lane token of nuscenes to node_tokens that form the centerline lane_dict = populate_lane_dict(nusc_map, root, data) # k: token, v: list of node_token poly_dict = populate_polys(data) create_lanes_xml(nusc_map, root, data, filename, args.argo_dir, lane_dict, poly_dict)
def render_scene(self, idx): sample = self.samples[idx] sample_token = sample['token'] scene = self.nus.get('scene', sample['scene_token']) log = self.nus.get('log', scene['log_token']) location = log['location'] nus_map = NuScenesMap(dataroot=self.root, map_name=location) layer_names = ['road_segment', 'lane', 'ped_crossing', 'walkway', 'stop_line', 'carpark_area'] camera_channel = 'CAM_FRONT' nus_map.render_map_in_image(self.nus, sample_token, layer_names=layer_names, camera_channel=camera_channel)
def get_nusc_maps(dataroot): nusc_maps = { map_name: NuScenesMap(dataroot=dataroot, map_name=map_name) for map_name in [ "singapore-hollandvillage", "singapore-queenstown", "boston-seaport", "singapore-onenorth", ] } return nusc_maps
def render_map(self, idx, combined=True): sample = self.samples[idx] sample_token = sample['token'] sample_data_lidar = self.nus.get('sample_data', sample['data']['LIDAR_TOP']) ego_pose = self.nus.get('ego_pose', sample_data_lidar['ego_pose_token']) ego_pose_xy = ego_pose['translation'] ego_pose_rotation = ego_pose['rotation'] timestamp = ego_pose['timestamp'] # 2. Generate Map & Agent Masks scene = self.nus.get('scene', sample['scene_token']) log = self.nus.get('log', scene['log_token']) location = log['location'] nus_map = NuScenesMap(dataroot=self.root, map_name=location) static_layer = StaticLayerRasterizer(self.helper, layer_names=self.layer_names, colors=self.colors) agent_layer = AgentBoxesWithFadedHistory( self.helper, seconds_of_history=self.seconds) map_masks, lanes, map_img = static_layer.generate_mask( ego_pose_xy, ego_pose_rotation, sample_token) agent_mask = agent_layer.generate_mask(ego_pose_xy, ego_pose_rotation, sample_token) if combined: plt.subplot(1, 2, 1) plt.title("combined map") plt.imshow(map_img) plt.subplot(1, 2, 2) plt.title("agent") plt.imshow(agent_mask) plt.show() else: num_labels = len(self.layer_names) num_rows = num_labels // 3 fig, ax = plt.subplots(num_rows, 3, figsize=(10, 3 * num_rows)) for row in range(num_rows): for col in range(3): num = 3 * row + col if num == num_labels - 1: break ax[row][col].set_title(self.layer_names[num]) ax[row][col].imshow(map_masks[num]) plt.show()
def setUp(self): """ Initialize the map for each location. """ self.nusc_maps = dict() for map_name in locations: # Load map. nusc_map = NuScenesMap(map_name=map_name, dataroot=os.environ['NUSCENES']) # Render for debugging. if self.render: nusc_map.render_layers(['lane'], figsize=1) plt.show() self.nusc_maps[map_name] = nusc_map
def load_all_maps(helper: PredictHelper, verbose: bool = False) -> Dict[str, NuScenesMap]: """ Loads all NuScenesMap instances for all available maps. :param helper: Instance of PredictHelper. :param verbose: Whether to print to stdout. :return: Mapping from map-name to the NuScenesMap api instance. """ dataroot = helper.data.dataroot maps = {} for map_name in locations: if verbose: print(f'static_layers.py - Loading Map: {map_name}') maps[map_name] = NuScenesMap(dataroot, map_name=map_name) return maps
def load_all_maps(helper: PredictHelper) -> Dict[str, NuScenesMap]: """ Loads all NuScenesMap instances for all available maps. :param helper: Instance of PredictHelper. :return: Mapping from map-name to the NuScenesMap api instance. """ dataroot = helper.data.dataroot json_files = filter(lambda f: "json" in f and "prediction_scenes" not in f, os.listdir(os.path.join(dataroot, "maps"))) maps = {} for map_file in json_files: map_name = str(map_file.split(".")[0]) print(f'static_layers.py - Loading Map: {map_name}') maps[map_name] = NuScenesMap(dataroot, map_name=map_name) return maps
def __getitem__(self, idx): # 1. past_agents_traj : (Num obv agents in batch_i X 20 X 2) # 2. past_agents_traj_len : (Num obv agents in batch_i, ) # 3. future_agents_traj : (Num pred agents in batch_i X 20 X 2) # 4. future_agents_traj_len : (Num pred agents in batch_i, ) # 5. future_agent_masks : (Num obv agents in batch_i) # 6. decode_rel_pos: (Num pred agents in batch_i X 2) # 7. decode_start_pos: (Num pred agents in batch_i X 2) # 8. map_image : (3 X 224 X 224) # 9. scene ID: (string) sample = self.samples[idx] sample_token = sample['token'] # 1. calculate ego pose sample_data_lidar = self.nus.get('sample_data', sample['data']['LIDAR_TOP']) ego_pose = self.nus.get('ego_pose', sample_data_lidar['ego_pose_token']) ego_pose_xy = ego_pose['translation'] ego_pose_rotation = ego_pose['rotation'] # 타임스탬프 timestamp = ego_pose['timestamp'] # 2. Generate Map & Agent Masks scene = self.nus.get('scene', sample['scene_token']) log = self.nus.get('log', scene['log_token']) location = log['location'] nus_map = NuScenesMap(dataroot=self.root, map_name=location) static_layer = StaticLayerRasterizer(self.helper, layer_names=self.layer_names, colors=self.colors) agent_layer = AgentBoxesWithFadedHistory(self.helper, seconds_of_history=self.seconds) map_masks, lanes, map_img = static_layer.generate_mask(ego_pose_xy, ego_pose_rotation, sample_token) agent_mask = agent_layer.generate_mask(ego_pose_xy, ego_pose_rotation, sample_token) # 3. Generate Agent Trajectory annotation_tokens = sample['anns'] num_agent = len(annotation_tokens) agents = [] for ans_token in annotation_tokens: agent_states = [] agent = self.nus.get('sample_annotation', ans_token) instance_token = agent['instance_token'] # 에이전트 주행경로 xy_global = agent['translation'] past_xy_global = self.helper.get_past_for_agent( instance_token, sample_token, seconds=self.seconds, in_agent_frame=False) future_xy_global = self.helper.get_future_for_agent( instance_token, sample_token, seconds=self.seconds, in_agent_frame=False) # 로컬 주행경로 xy_local = convert_global_coords_to_local(np.array([xy_global]), ego_pose_xy, ego_pose_rotation) past_xy_local = convert_global_coords_to_local(past_xy_global, ego_pose_xy, ego_pose_rotation) future_xy_local = convert_global_coords_to_local(future_xy_global, ego_pose_xy, ego_pose_rotation) # 에이전트 주행상태 rot = agent['rotation'] vel = self.helper.get_velocity_for_agent(instance_token, sample_token) accel = self.helper.get_acceleration_for_agent(instance_token, sample_token) yaw_rate = self.helper.get_heading_change_rate_for_agent(instance_token, sample_token) agent_states = {'present_pos': xy_global, 'past_pos': past_xy_global, 'future_pos': future_xy_global, 'rot': rot, 'vel': vel, 'accel': accel, 'yaw_rate': yaw_rate, 'present_local_xy': xy_local, 'past_local_xy': past_xy_local, 'future_local_xy': future_xy_local} agents.append(agent_states) return map_masks, agents
def plot_agent_scene( self, ego_centric: bool = False, ego_traj=None, instance_token: str = None, sample_token: str = None, map_layers=None, sensor_info=None, sensing_patch=None, predictions=None, agent_state_dict=None, plot_list=None, legend=False, text_box=False, show_axis=True, render_ego_pose_range=False, paper_ready=False, read_from_cached=False, plot_agent_trajs=True, animated_agent=False, bfig=None, # for birdseye image bax=None, sfig=None, # for camera sax=None, render_additional=None, plot_human_ego=True, patch_margin=30, sim_ego_pose=None, save_image_ax=False): if paper_ready: legend = False text_box = False show_axis = False render_ego_pose_range = False plot_agent_trajs = False animated_agent = True #show_axis = True if map_layers is None: map_layers = self.map_layers if plot_list is None: plot_list = self.plot_list sample = self.nusc.get('sample', sample_token) scene = self.nusc.get('scene', sample['scene_token']) scene_log = self.nusc.get('log', scene['log_token']) nusc_map = NuScenesMap(dataroot=self.dataroot, map_name=scene_log['location']) patch_margin = patch_margin min_diff_patch = 30 if not ego_centric: agent_future = self.helper.get_future_for_agent( instance_token, sample_token, self.na_config['pred_horizon'], in_agent_frame=False, just_xy=True) agent_past = self.helper.get_past_for_agent( instance_token, sample_token, self.na_config['obs_horizon'], in_agent_frame=False, just_xy=True) #### set plot patch #### if agent_future.shape[0] > 0: p = agent_future[0] else: p = agent_past[-1] my_patch = ( p[0] - patch_margin, p[1] - patch_margin, p[0] + patch_margin, p[1] + patch_margin, ) # min_patch = np.floor(agent_future.min(axis=0) - patch_margin) # max_patch = np.ceil(agent_future.max(axis=0) + patch_margin) # diff_patch = max_patch - min_patch # if any(diff_patch < min_diff_patch): # center_patch = (min_patch + max_patch) / 2 # diff_patch = np.maximum(diff_patch, min_diff_patch) # min_patch = center_patch - diff_patch / 2 # max_patch = center_patch + diff_patch / 2 # my_patch = (min_patch[0], min_patch[1], max_patch[0], max_patch[1]) else: sample_data = self.nusc.get('sample_data', sample['data']['CAM_FRONT']) ego_pose = self.nusc.get('ego_pose', sample_data['ego_pose_token']) my_patch = ( ego_pose['translation'][0] - patch_margin, ego_pose['translation'][1] - patch_margin, ego_pose['translation'][0] + patch_margin, ego_pose['translation'][1] + patch_margin, ) if sim_ego_pose is not None: my_patch = ( sim_ego_pose['translation'][0] - patch_margin, sim_ego_pose['translation'][1] - patch_margin, sim_ego_pose['translation'][0] + patch_margin, sim_ego_pose['translation'][1] + patch_margin, ) #### read from saved path if present #### read_img = False if read_from_cached: scene_path = os.path.join( scene_img_dir, scene['name'] + "-token-" + sample['scene_token']) p_scene = Path(scene_img_dir) saved_scene_list = [ str(f) for f in p_scene.iterdir() if f.is_dir() ] if scene_path in saved_scene_list: p_sample = Path(scene_path) for f in p_sample.iterdir(): if sample_token in str(f): ax = cloudpickle.load( open(os.path.join(scene_path, str(f)), 'rb')) fig = plt.figure(figsize=(10, 10)) fig._axstack.add('ax', ax) read_img = True if not read_img: fig, ax = nusc_map.render_map_patch( my_patch, map_layers, figsize=(10, 10), render_egoposes_range=render_ego_pose_range, render_legend=legend, fig=bfig, axes=bax) if not ego_centric: ax.set_title(scene['name'] + " instance_token: " + instance_token + ", sample_token: " + sample_token + "\n" + ", decription " + scene['description']) else: ax.set_title(scene['name'] + ", sample_token: " + sample_token + "\n" + ", decription " + scene['description']) #### label map #### if 'labeled_map' in plot_list: records_within_patch = nusc_map.get_records_in_patch( my_patch, nusc_map.non_geometric_layers, mode='within') self.label_map(ax, nusc_map, records_within_patch['stop_line'], text_box=text_box) #### Plot ego #### if 'ego' in plot_list: ego_pos, ego_quat = self.plot_ego( ax, sample, ego_traj=ego_traj, animated_agent=animated_agent, plot_ego=plot_human_ego) #### Plot other agents #### if 'pedestrian' in plot_list or 'other_cars' in plot_list: road_agents_in_patch = self.plot_road_agents( ax, instance_token, sample, plot_list, text_box, sensor_info, my_patch, plot_traj=plot_agent_trajs, animated_agent=animated_agent) ################## # Car to predict # ################## if not ego_centric: agent_pos, agent_quat = self.plot_car_to_predict( ax, agent_future, agent_past, instance_token, sample_token, text_box, predictions, agent_state_dict) #### plot all_info #### if 'map_info' in self.plot_list: if not ego_centric: self.plot_map_info(ax, agent_pos, nusc_map, text_box=text_box) else: self.plot_map_info(ax, ego_pos, nusc_map, text_box=text_box) #### plot sensor info ### if sensor_info is not None and 'sensor_info' in plot_list: self.plot_sensor_info(ax, sensor_info=sensor_info, text_box=text_box) #### render map layers on camera images #### if 'cam' in plot_list: #sfig, sax = plt.subplots(nrows=2, ncols=3, sharex=True, sharey=True, figsize=(9,16)) layer_names = [ 'road_segment', 'lane', 'ped_crossing', 'walkway', 'stop_line', 'carpark_area' ] layer_names = [] #cam_names = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', # 'CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT'] cam_names = ['CAM_FRONT'] k = 0 if sfig is None: if len(cam_names) == 6: sfig, sax = plt.subplots(nrows=2, ncols=3) for i in range(2): for j in range(3): sax[i, j].xaxis.set_visible(False) sax[i, j].yaxis.set_visible(False) cam_fig, cam_ax = nusc_map.render_map_in_image( self.nusc, sample_token, layer_names=layer_names, camera_channel=cam_names[k], ax=sax[i, j]) k += 1 elif len(cam_names) == 1: cam_fig, cam_ax = plt.subplots() cam_ax.xaxis.set_visible(False) cam_ax.yaxis.set_visible(False) nusc_map.render_map_in_image(self.nusc, sample_token, layer_names=layer_names, camera_channel=cam_names[k], ax=cam_ax) else: raise ValueError('') if sfig is not None: sfig.tight_layout(pad=0) sfig.set_figheight(7) sfig.set_figwidth(15) # for car_info in road_agents_in_patch['vehicles']: # instance_token = car_info['instance_token'] # # render annotations inside patch # ann = self.helper.get_sample_annotation(instance_token, sample_token) # ann_fig, ann_ax = self.nusc.render_annotation(ann['token']) # if sensing_patch is not None and self.in_shapely_polygon(car_info['translation'], sensing_patch): # ann_ax.set_title("Sensed") else: sfig, sax = None, None cam_fig, cam_ax = None, None #### render additional outside information #### if render_additional is not None: self.render_additional(ax, render_additional) if not show_axis: plt.axis('off') plt.grid('off') ax.grid(False) ax.set_aspect('equal') other = { 'cam_fig': cam_fig, 'cam_ax': cam_ax, 'sfig': sfig, 'sax': sax } return fig, ax, other
def main(): input_root, output_root, do_trajectories, do_maps = get_command_line_args() print("Preprocessing Script for nuScenes Dataset.") print("Trajectories: {}, Maps: {}".format("Y" if do_trajectories else "N", "Y" if do_maps else "N")) if do_trajectories: nusc = NuScenes(version='v1.0-trainval', dataroot=input_root) name2ind = {} # Maps "scene-name" to nusc.scene(list) index. for ind, member in enumerate(nusc.scene): name2ind[member['name']] = ind token2attr = {} # Maps attribute_token to attribute string. for attribute in nusc.attribute: token2attr[attribute['token']] = attribute['name'] splits = create_splits_scenes() if do_maps: from nuscenes.map_expansion.map_api import NuScenesMap city_list = ['singapore-onenorth', 'singapore-hollandvillage', 'singapore-queenstown', 'boston-seaport'] for city_name in city_list: scale_h = scale_w = 2 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) nusc_map = NuScenesMap(input_root, city_name) print("Calculating a map mask with scale_h: {:d}, scale_w: {:d}... ".format(scale_h, scale_w), end="", flush=True) map_mask = get_drivable_area_mask(nusc_map, scale_h=2, scale_w=2) print("Done.") 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, 'scale_h': 2, 'scale_w': 2}, mask_path) dump({'map': dt, 'scale_h': 2, 'scale_w': 2}, 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))) if do_trajectories: for partition in ['train', 'train_val', 'val']: print("Generating Trajectories for {:s} set.".format(partition)) if 'train' in partition: scene_list = splits['train'] if partition == "train": scene_list = scene_list[NUM_IN_TRAIN_VAL:] if partition == "train_val": scene_list = scene_list[:NUM_IN_TRAIN_VAL] else: scene_list = splits['val'] pool = Pool(cpu_count()) callback = Counting_Callback(task_name="Trajectory Imputation & Smoothing", num_data=len(scene_list)) for name in scene_list: """ Generate a raw DataFrame object for each scene_name. Filter object categories other than "human" and "vehicle". Perform Kalman Smoothing and/or rule-based Imputation. """ ind = name2ind[name] scene = nusc.scene[ind] log = nusc.get('log', scene['log_token']) location = log['location'] data_dict = annotation_dict() sample_token = scene['first_sample_token'] last_sample_token = scene['last_sample_token'] frame = 0 passed_last = False while not passed_last: sample_data = nusc.get("sample", sample_token) timestamp = sample_data["timestamp"] # Gather pose token from LiDAR since it has timestamp synced with sample_data. lidar_data = nusc.get("sample_data", sample_data["data"]["LIDAR_TOP"]) ego_pose_token = lidar_data['ego_pose_token'] ego_pose_data = nusc.get("ego_pose", ego_pose_token) ego_translation = ego_pose_data["translation"] ego_rotation = ego_pose_data["rotation"] # Append Ego Motion Data data_dict.append(location, sample_token, '00000000000000000000000000000000', 'vehicle.ego', 'ego', timestamp, frame, ego_translation, ego_rotation) for anns_token in sample_data['anns']: anns_data = nusc.get("sample_annotation", anns_token) instance_token = anns_data['instance_token'] instance_category = anns_data['category_name'] instance_attributes = [token2attr[token] for token in anns_data['attribute_tokens']] instance_attributes = ", ".join(instance_attributes) instance_translation = anns_data["translation"] instance_rotation = anns_data["rotation"] # Append Instance Motion Data data_dict.append(location, sample_token, instance_token, instance_category, instance_attributes, timestamp, frame, instance_translation, instance_rotation) # goto next sample if sample_token == last_sample_token or len(sample_data['next']) == 0: passed_last = True else: sample_token = sample_data['next'] frame += 1 df = pd.DataFrame.from_dict(data_dict) # Generate a DataFrame pool.apply_async(kalman_smoother, (df.copy(), name), callback=callback) # Perform Kalman Smoothing pool.close() pool.join() # Get Kalman Smoothed results and sort w.r.t. scene_anme. smoothed_trajectories = callback.get_results() smoothed_trajectories.sort(key=lambda x: x[1]) pool = Pool(cpu_count()) callback = Counting_Callback(task_name="Trajectory Chopping & Sample Generation", num_data=float('inf')) num_data = 0 for df, scene_name in smoothed_trajectories: """ Chop a smoothed DataFrame into multiple samples (~33 samples per scene) such that each sample spans 5 seconds where the reference frame is set at the 2 second's frame. Then, split the sample to obsv (0~2 seconds) and pred (2~5 seconds) files. """ scn_code = int(scene_name.split('-')[-1]) frames = df.FRAME.to_list() initial_frame = frames[0] last_frame = frames[-1] for ref_frame in range(initial_frame+3, last_frame-5): obsv_path = P(output_root).joinpath(partition, 'observation', '{:04d}-{:03d}.pkl'.format(scn_code, ref_frame)) obsv_path.parent.mkdir(parents=True, exist_ok=True) pred_path = P(output_root).joinpath(partition, 'prediction', '{:04d}-{:03d}.pkl'.format(scn_code, ref_frame)) pred_path.parent.mkdir(parents=True, exist_ok=True) pool.apply_async(generate_trajectories, (df.copy(), ref_frame, obsv_path, pred_path), callback=callback) # generate_trajectories(df.copy(), ref_frame, obsv_path, pred_path) num_data += 1 callback.set_num_data(num_data) pool.close() pool.join() print("Saved {:d} {:s} samples at {:s}.".format(num_data, partition, str(P(output_root).joinpath(partition)))) # 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.")
def convert_scene(self, scene_name, utils, dataroot, up_to_step, *args, **kwargs): experiment_result_dict = None if 'experiment_result_dict' in kwargs.keys(): experiment_result_dict = kwargs['experiment_result_dict'] scene_i = None for scene_i in self.nusc.scene: if scene_i['name'] == scene_name: scene = scene_i break scene_name = scene['name'] log = self.nusc.get('log', scene['log_token']) location = log['location'] print(f'Loading map "{location}"') self.nusc_map = NuScenesMap(dataroot=dataroot, map_name=location) print(f'Loading bitmap "{self.nusc_map.map_name}"') bitmap = BitMap(self.nusc_map.dataroot, self.nusc_map.map_name, 'basemap') print(f'Loaded {bitmap.image.shape} bitmap') cur_sample = self.nusc.get('sample', scene['first_sample_token']) can_parsers = [ [self.nusc_can.get_messages(scene_name, 'ms_imu'), 0, utils.get_imu_msg], [self.nusc_can.get_messages(scene_name, 'pose'), 0, utils.get_odom_msg], [self.nusc_can.get_messages(scene_name, 'steeranglefeedback'), 0, lambda x: utils.get_basic_can_msg('Steering Angle', x)], [self.nusc_can.get_messages(scene_name, 'vehicle_monitor'), 0, lambda x: utils.get_basic_can_msg('Vehicle Monitor', x)], [self.nusc_can.get_messages(scene_name, 'zoesensors'), 0, lambda x: utils.get_basic_can_msg('Zoe Sensors', x)], [self.nusc_can.get_messages(scene_name, 'zoe_veh_info'), 0, lambda x: utils.get_basic_can_msg('Zoe Vehicle Info', x)], ] bag_name = f"{self.model_name}_{self.ckpt_name}_{self.config['scene_name']}.bag" bag_path = os.path.join(self.scene_data_p, bag_name) print(f'Writing to {bag_path}') bag = rosbag.Bag(bag_path, 'w', compression='lz4') stamp = utils.get_time(self.nusc.get('ego_pose', self.nusc.get('sample_data', cur_sample['data']['LIDAR_TOP'])['ego_pose_token'])) map_msg = utils.get_scene_map(scene, self.nusc_map, bitmap, stamp) centerlines_msg = utils.get_centerline_markers(scene, self.nusc_map, stamp) bag.write('/map', map_msg, stamp) bag.write('/semantic_map', centerlines_msg, stamp) last_map_stamp = stamp #idx = list(experiment_result_dict['sim_ego_pos'].keys())[0] idx = 0 while cur_sample is not None: if up_to_step is not None and idx > up_to_step: break sample_lidar = self.nusc.get('sample_data', cur_sample['data']['LIDAR_TOP']) ego_pose = self.nusc.get('ego_pose', sample_lidar['ego_pose_token']) stamp = utils.get_time(ego_pose) #### write experiment_result to bag #### ego_marker = None if experiment_result_dict is not None: # figures # for k,v in experiment_result_dict['figures'].items(): if idx in v.keys(): png = v[idx] msg = CompressedImage() msg.header.frame_id = k msg.header.stamp = stamp msg.format = 'png' msg.data = png bag.write("/"+k+'/image_rect_compressed', msg, stamp) #print(idx) #print(experiment_result_dict['sim_ego_pos'].keys()) if idx in experiment_result_dict['sim_ego_pos'].keys(): ego_pose = { 'translation': experiment_result_dict['sim_ego_pos'][idx], 'rotation': experiment_result_dict['sim_ego_quat'][idx], 'token': ego_pose['token'], 'timestamp': ego_pose['timestamp'] } if idx == list(experiment_result_dict['sim_ego_pos'].keys())[-1]: next_idx = idx break else: next_idx = idx+1 next_sample = utils.nusc.get('sample',cur_sample['next']) sample_lidar = utils.nusc.get('sample_data', next_sample['data']['LIDAR_TOP']) real_next_ego_pose = utils.nusc.get('ego_pose', sample_lidar['ego_pose_token']) next_ego_pose = { 'translation': experiment_result_dict['sim_ego_pos'][next_idx], 'rotation': experiment_result_dict['sim_ego_quat'][next_idx], 'token': real_next_ego_pose['token'], 'timestamp': real_next_ego_pose['timestamp'] } # publish /tf tf_array = utils.get_tfmessage(cur_sample, current_pose=ego_pose, next_pose=next_ego_pose) bag.write('/tf', tf_array, stamp) else: # publish /tf tf_array = utils.get_tfmessage(cur_sample) bag.write('/tf', tf_array, stamp) else: # publish /tf tf_array = utils.get_tfmessage(cur_sample) bag.write('/tf', tf_array, stamp) # write map topics every two seconds if stamp - rospy.Duration(2.0) >= last_map_stamp: map_msg.header.stamp = stamp for marker in centerlines_msg.markers: marker.header.stamp = stamp bag.write('/map', map_msg, stamp) bag.write('/semantic_map', centerlines_msg, stamp) last_map_stamp = stamp # write CAN messages to /pose, /odom, and /diagnostics can_msg_events = [] for i in range(len(can_parsers)): (can_msgs, index, msg_func) = can_parsers[i] while index < len(can_msgs) and utils.get_utime(can_msgs[index]) < stamp: can_msg_events.append(msg_func(can_msgs[index])) index += 1 can_parsers[i][1] = index can_msg_events.sort(key = lambda x: x[0]) for (msg_stamp, topic, msg) in can_msg_events: bag.write(topic, msg, stamp) # /driveable_area occupancy grid utils.write_occupancy_grid(bag, self.nusc_map, ego_pose, stamp) # iterate sensors for (sensor_id, sample_token) in cur_sample['data'].items(): sample_data = self.nusc.get('sample_data', sample_token) topic = '/' + sensor_id # write the sensor data # if sample_data['sensor_modality'] == 'radar': # msg = utils.get_radar(sample_data, sensor_id) # bag.write(topic, msg, stamp) # elif sample_data['sensor_modality'] == 'lidar': # msg = utils.get_lidar(sample_data, sensor_id) # bag.write(topic, msg, stamp) # if sample_data['sensor_modality'] == 'camera': # msg = utils.get_camera(sample_data, sensor_id, dataroot) # bag.write(topic + '/image_rect_compressed', msg, stamp) # msg = utils.get_camera_info(sample_data, sensor_id) # bag.write(topic + '/camera_info', msg, stamp) # if sample_data['sensor_modality'] == 'camera': # msg = utils.get_lidar_imagemarkers(sample_lidar, sample_data, sensor_id) # bag.write(topic + '/image_markers_lidar', msg, stamp) # utils.write_boxes_imagemarkers(bag, cur_sample['anns'], sample_data, sensor_id, topic, stamp) # publish /markers/annotations marker_array = MarkerArray() ego_marker = Marker() ego_marker.header.frame_id = 'base_link' ego_marker.header.stamp = stamp ego_marker.id = 10000 ego_marker.text = 'ego' ego_marker.type = Marker.CUBE #ego_marker.pose = utils.get_pose(next_ego_pose) ego_marker.frame_locked = True ego_marker.scale.x = 4 ego_marker.scale.y = 1 ego_marker.scale.z = 2 ego_marker.color = utils.make_color((255,61,99), 0.5) marker_array.markers.append(ego_marker) # publish /pose pose_stamped = PoseStamped() pose_stamped.header.frame_id = 'base_link' pose_stamped.header.stamp = stamp pose_stamped.pose.orientation.w = 1 bag.write('/pose', pose_stamped, stamp) # publish /gps coordinates = utils.derive_latlon(location, ego_pose) gps = NavSatFix() gps.header.frame_id = 'base_link' gps.header.stamp = stamp gps.status.status = 1 gps.status.service = 1 gps.latitude = coordinates['latitude'] gps.longitude = coordinates['longitude'] gps.altitude = utils.get_transform(ego_pose).translation.z bag.write('/gps', gps, stamp) for annotation_id in cur_sample['anns']: ann = self.nusc.get('sample_annotation', annotation_id) marker_id = int(ann['instance_token'][:4], 16) c = np.array(self.nusc.explorer.get_color(ann['category_name'])) / 255.0 marker = Marker() marker.header.frame_id = 'map' marker.header.stamp = stamp marker.id = marker_id marker.text = ann['instance_token'][:4] marker.type = Marker.CUBE marker.pose = utils.get_pose(ann) marker.frame_locked = True marker.scale.x = ann['size'][1] marker.scale.y = ann['size'][0] marker.scale.z = ann['size'][2] marker.color = utils.make_color(c, 0.5) marker_array.markers.append(marker) # ego marker # # marker = Marker() # marker.header.frame_id = 'map' # marker.header.stamp = stamp # marker.id = marker_id # marker.text = 'ego' # marker.type = Marker.CUBE # marker.pose = utils.get_pose(ego_pose) # marker.frame_locked = True # marker.scale.x = 4.3 # marker.scale.y = 1.8 # marker.scale.z = 1.6 # c = [0.1,0.1,0.6] # marker.color = utils.make_color(c, 0.9) # marker_array.markers.append(marker) bag.write('/markers/annotations', marker_array, stamp) # collect all sensor frames after this sample but before the next sample non_keyframe_sensor_msgs = [] for (sensor_id, sample_token) in cur_sample['data'].items(): topic = '/' + sensor_id next_sample_token = self.nusc.get('sample_data', sample_token)['next'] while next_sample_token != '': next_sample_data = self.nusc.get('sample_data', next_sample_token) # if next_sample_data['is_key_frame'] or get_time(next_sample_data).to_nsec() > next_stamp.to_nsec(): # break if next_sample_data['is_key_frame']: break # if next_sample_data['sensor_modality'] == 'radar': # msg = utils.get_radar(next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic, msg)) # elif next_sample_data['sensor_modality'] == 'lidar': # msg = utils.get_lidar(next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic, msg)) # if next_sample_data['sensor_modality'] == 'camera': # msg = utils.get_camera(next_sample_data, sensor_id, dataroot) # camera_stamp_nsec = msg.header.stamp.to_nsec() # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/image_rect_compressed', msg)) # msg = utils.get_camera_info(next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/camera_info', msg)) # closest_lidar = utils.find_closest_lidar(cur_sample['data']['LIDAR_TOP'], camera_stamp_nsec) # if closest_lidar is not None: # msg = utils.get_lidar_imagemarkers(closest_lidar, next_sample_data, sensor_id) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic + '/image_markers_lidar', msg)) # else: # msg = utils.get_remove_imagemarkers(sensor_id, 'LIDAR_TOP', msg.header.stamp) # non_keyframe_sensor_msgs.append((msg.header.stamp.to_nsec(), topic + '/image_markers_lidar', msg)) # Delete all image markers on non-keyframe camera images # msg = get_remove_imagemarkers(sensor_id, 'LIDAR_TOP', msg.header.stamp) # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/image_markers_lidar', msg)) # msg = get_remove_imagemarkers(sensor_id, 'annotations', msg.header.stamp) # non_keyframe_sensor_msgs.append((camera_stamp_nsec, topic + '/image_markers_annotations', msg)) next_sample_token = next_sample_data['next'] # sort and publish the non-keyframe sensor msgs non_keyframe_sensor_msgs.sort(key=lambda x: x[0]) for (_, topic, msg) in non_keyframe_sensor_msgs: bag.write(topic, msg, msg.header.stamp) # move to the next sample cur_sample = self.nusc.get('sample', cur_sample['next']) if cur_sample.get('next') != '' else None idx += 1 bag.close() print(f'Finished writing {bag_name}')
def __init__(self, config: dict = {}, helper: PredictHelper = None, py_logger=None, tb_logger=None): self.na_config = Configuration({ 'debug': False, 'pred_horizon': 6, 'obs_horizon': 2, 'freq': 2, 'load_dataset': False, 'version': 'v1.0-mini', 'debug': False, 'py_logging_path': None, 'tb_logging_path': None }) self.na_config.update(config) self.name = None self.py_logger = py_logger self.tb_logger = tb_logger self.dataroot = None if 'mini' in self.na_config['version']: self.dataroot = mini_path else: self.dataroot = full_path if self.py_logger is None and self.na_config[ 'py_logging_path'] is not None: print(f"py logging path: {self.na_config['py_logging_path']}") self.py_logger = logger self.py_logger.add(self.na_config['py_logging_path'] + "/log.txt") # self.py_logger = logging.getLogger(self.name) # print(f"py logging path: {self.na_config['py_logging_path']}") # self.py_logger.addHandler(logging.FileHandler(os.path.join(self.na_config['py_logging_path'], 'log.txt'), mode='w')) # if self.py_logger is not None: # self.py_logger.propagate = False if self.tb_logger is None and self.na_config[ 'tb_logging_path'] is not None: self.tb_logger = SummaryWriter( log_dir=os.path.join(self.na_config['tb_logging_path'])) self.helper = helper self.nusc = None if self.helper is not None: self.nusc = self.helper.data else: if self.dataroot is not None and self.na_config[ 'version'] is not None and self.na_config[ 'load_dataset'] and self.helper is None: self.nusc = NuScenes(dataroot=self.dataroot, version=self.na_config['version'], verbose=True) self.helper = PredictHelper(self.nusc) #### Initialize Map #### self.nusc_map_dict = { 'boston-seaport': NuScenesMap(dataroot=self.dataroot, map_name='boston-seaport'), 'singapore-hollandvillage': NuScenesMap(dataroot=self.dataroot, map_name='singapore-hollandvillage'), 'singapore-onenorth': NuScenesMap(dataroot=self.dataroot, map_name='singapore-onenorth'), 'singapore-queenstown': NuScenesMap(dataroot=self.dataroot, map_name='singapore-queenstown'), } #### Initialize CAN #### self.nusc_can = NuScenesCanBus(dataroot=self.dataroot) #### self.all_info = {'config': self.na_config}
def preprocess_maps(dataroot, glob_path="./*.json"): fs = glob(glob_path) for fi, f in enumerate(fs): with open(f, "r") as reader: data = json.load(reader) nusc_map = NuScenesMap(dataroot=dataroot, map_name=data["map_name"]) dataset = dict() center, h, w = data["center"], data["height"], data["width"] patch = [ center[0] - w / 2, center[1] - h / 2, center[0] + w / 2, center[1] + h / 2, ] dataset["patch"] = patch dataset["center"] = np.array([center]) dataset["height"] = h dataset["width"] = w dataset["map_name"] = data["map_name"] dataset["dx"] = np.array(data["dx"]) dataset["bx"] = np.array(data["bx"]) dataset["road_img"] = np.array(data["road_img"]) # Needed for lidar sensors pt1, pt2 = nuscenes_map_to_line_representation(nusc_map, patch, False) dataset["edges"] = (pt1, pt2) drivable_area, xs, ys = get_drivable_area_matrix(data, patch, res=500) dataset["plotting_utils"] = ( drivable_area.numpy().flatten(), xs.numpy().flatten(), ys.numpy().flatten(), [(0.5, 0.5, 0.5) if row else (1, 1, 1) for row in drivable_area.numpy().flatten()], ) dataset["splines"] = dict() signal_loc = torch.as_tensor(data["signal_locations"]) signal_color = data["mapping"] dataset["signal_locations"] = signal_loc dataset["color_mapping"] = signal_color dataset["starts_to_signal"] = data["starts_to_signal"] dataset["signal_loc"] = [] dataset["signal_color"] = [] for starti, (key, paths) in enumerate(data["all_paths"].items()): idx = data["starts_to_signal"][starti] loc = signal_loc[idx] col = signal_color[idx] dataset["signal_loc"].append(loc) dataset["signal_color"].append(col) dataset["splines"][starti] = dict() for pathi, path in enumerate(paths): path = np.array(path) if path.shape[0] < 75: print( "Skipping spline as it contains very few control points" ) continue dataset["splines"][starti][pathi] = [] for i in range(0, 50, 10): cps = path[ np.linspace(i, path.shape[0] - 15, 12, dtype=np.int), :2, ] diff = cps[0] - cps[1] theta = np.arctan2(diff[1], diff[0]) start_orientation = (angle_normalize( torch.as_tensor(math.pi + theta)).float().reshape( 1, 1)) if i == 0: extra_pt1 = np.array([[ cps[0, 0] + np.cos(theta) * 30.0, cps[0, 1] + np.sin(theta) * 30.0, ]]) else: extra_pt1 = path[0:1, :2] diff = cps[-1] - cps[-2] theta = np.arctan2(diff[1], diff[0]) dest_orientation = (angle_normalize( torch.as_tensor(theta)).float().reshape(1, 1)) extra_pt2 = np.array([[ cps[-1, 0] + np.cos(theta) * 30.0, cps[-1, 1] + np.sin(theta) * 30.0, ]]) cps = torch.cat([ torch.from_numpy(cps), torch.from_numpy(extra_pt2), torch.from_numpy(extra_pt1), ])[None, :, :].float() start_position = cps[:, 0, :] destination = cps[:, -3, :] dataset["splines"][starti][pathi].append(( start_position, destination, start_orientation, dest_orientation, cps, )) outname = f"env{data['map_name']}_{data['center'][0]}_{data['center'][1]}.pth" print("saving", outname) torch.save(dataset, outname)
data.sort_values('frame_id', inplace=True) max_timesteps = data['frame_id'].max() x_min = np.round(data['x'].min() - 50) x_max = np.round(data['x'].max() + 50) y_min = np.round(data['y'].min() - 50) y_max = np.round(data['y'].max() + 50) data['x'] = data['x'] - x_min data['y'] = data['y'] - y_min scene = Scene(timesteps=max_timesteps + 1, dt=0.5, name=str(scene_id)) # Generate Maps map_name = nusc.get('log', nuscene['log_token'])['location'] nusc_map = NuScenesMap(dataroot=data_path, map_name=map_name) type_map = dict() x_size = x_max - x_min y_size = y_max - y_min patch_box = (x_min + 0.5 * (x_max - x_min), y_min + 0.5 * (y_max - y_min), y_size, x_size) patch_angle = 0 # Default orientation where North is up canvas_size = (np.round(3 * y_size).astype(int), np.round(3 * x_size).astype(int)) homography = np.array([[3., 0., 0.], [0., 3., 0.], [0., 0., 3.]]) layer_names = ['lane', 'road_segment', 'drivable_area', 'road_divider', 'lane_divider', 'stop_line', 'ped_crossing', 'stop_line', 'ped_crossing', 'walkway'] map_mask = (nusc_map.get_map_mask(patch_box, patch_angle, layer_names, canvas_size) * 255.0).astype( np.uint8) map_mask = np.swapaxes(map_mask, 1, 2) # x axis comes first # PEDESTRIANS map_mask_pedestrian = np.stack((map_mask[9], map_mask[8], np.max(map_mask[:3], axis=0)), axis=2)
def process_scene(sample_token, processed_sample_tokens, env, nusc, helper, data_path, data_class, half_dt, dynamic): data = pd.DataFrame(columns=[ 'frame_id', 'type', 'node_id', 'robot', 'x', 'y', 'z', 'length', 'width', 'height', 'heading' ]) samples = aggregate_samples(nusc, start=sample_token, data_class=data_class) attribute_dict = defaultdict(set) frame_id = 0 for sample in samples: annotation_tokens = sample['anns'] for annotation_token in annotation_tokens: annotation = nusc.get('sample_annotation', annotation_token) category = annotation['category_name'] if len(annotation['attribute_tokens']): attribute = nusc.get('attribute', annotation['attribute_tokens'][0])['name'] else: if 'vehicle' in category: attribute = '' else: continue if 'pedestrian' in category and 'stroller' not in category and 'wheelchair' not in category: our_category = env.NodeType.PEDESTRIAN elif 'vehicle' in category and 'bicycle' not in category and 'motorcycle' not in category: # and 'parked' not in attribute: our_category = env.NodeType.VEHICLE else: continue attribute_dict[annotation['instance_token']].add(attribute) data_point = pd.Series({ 'frame_id': frame_id, 'type': our_category, 'node_id': annotation['instance_token'], 'robot': False, 'x': annotation['translation'][0], 'y': annotation['translation'][1], 'z': annotation['translation'][2], 'length': annotation['size'][0], 'width': annotation['size'][1], 'height': annotation['size'][2], 'heading': Quaternion(annotation['rotation']).yaw_pitch_roll[0] }) data = data.append(data_point, ignore_index=True) # Ego Vehicle our_category = env.NodeType.VEHICLE sample_data = nusc.get('sample_data', sample['data']['CAM_FRONT']) annotation = nusc.get('ego_pose', sample_data['ego_pose_token']) data_point = pd.Series({ 'frame_id': frame_id, 'type': our_category, 'node_id': 'ego', 'robot': True, 'x': annotation['translation'][0], 'y': annotation['translation'][1], 'z': annotation['translation'][2], 'length': 4, 'width': 1.7, 'height': 1.5, 'heading': Quaternion(annotation['rotation']).yaw_pitch_roll[0], 'orientation': None }) data = data.append(data_point, ignore_index=True) processed_sample_tokens.add(sample['token']) frame_id += 1 if len(data.index) == 0: return None data.sort_values('frame_id', inplace=True) max_timesteps = data['frame_id'].max() x_min = np.round(data['x'].min() - 50) x_max = np.round(data['x'].max() + 50) y_min = np.round(data['y'].min() - 50) y_max = np.round(data['y'].max() + 50) data['x'] = data['x'] - x_min data['y'] = data['y'] - y_min scene = Scene(timesteps=max_timesteps + 1, dt=dt, name=sample_token, aug_func=augment, sample_tokens=[sample['token'] for sample in samples], x_min=x_min, y_min=y_min) # Generate Maps map_name = helper.get_map_name_from_sample_token(sample_token) nusc_map = NuScenesMap(dataroot=data_path, map_name=map_name) type_map = dict() x_size = x_max - x_min y_size = y_max - y_min patch_box = (x_min + 0.5 * (x_max - x_min), y_min + 0.5 * (y_max - y_min), y_size, x_size) patch_angle = 0 # Default orientation where North is up canvas_size = (np.round(3 * y_size).astype(int), np.round(3 * x_size).astype(int)) homography = np.array([[3., 0., 0.], [0., 3., 0.], [0., 0., 3.]]) layer_names = [ 'lane', 'road_segment', 'drivable_area', 'road_divider', 'lane_divider', 'stop_line', 'ped_crossing', 'stop_line', 'ped_crossing', 'walkway' ] map_mask = (nusc_map.get_map_mask(patch_box, patch_angle, layer_names, canvas_size) * 255.0).astype(np.uint8) map_mask = np.swapaxes(map_mask, 1, 2) # x axis comes first # PEDESTRIANS map_mask_pedestrian = np.stack( (map_mask[9], map_mask[8], np.max(map_mask[:3], axis=0)), axis=0) type_map['PEDESTRIAN'] = GeometricMap(data=map_mask_pedestrian, homography=homography, description=', '.join(layer_names)) # VEHICLES map_mask_vehicle = np.stack( (np.max(map_mask[:3], axis=0), map_mask[3], map_mask[4]), axis=0) type_map['VEHICLE'] = GeometricMap(data=map_mask_vehicle, homography=homography, description=', '.join(layer_names)) map_mask_plot = np.stack( ((np.max(map_mask[:3], axis=0) - (map_mask[3] + 0.5 * map_mask[4]).clip(max=255)).clip(min=0).astype( np.uint8), map_mask[8], map_mask[9]), axis=0) type_map['VISUALIZATION'] = GeometricMap( data=map_mask_plot, homography=homography, description=', '.join(layer_names)) scene.map = type_map del map_mask del map_mask_pedestrian del map_mask_vehicle del map_mask_plot for node_id in pd.unique(data['node_id']): node_frequency_multiplier = 1 node_df = data[data['node_id'] == node_id] if dynamic: if 'vehicle.parked' in attribute_dict[ node_id] or 'vehicle.stopped' in attribute_dict[node_id]: continue elif len(attribute_dict[node_id] ) == 1 and 'vehicle.parked' in attribute_dict[node_id]: # Catching instances of vehicles that were parked and then moving, but not allowing # only parked vehicles through. continue if node_df['x'].shape[0] < 2: continue if not np.all(np.diff(node_df['frame_id']) == 1): min_index = node_df['frame_id'].min() max_index = node_df['frame_id'].max() node_df = node_df.set_index('frame_id').reindex( range(min_index, max_index + 1)).interpolate().reset_index() node_df['type'] = node_df['type'].mode()[0] node_df['node_id'] = node_id node_df['robot'] = False # print('Occlusion') # continue # TODO Make better node_values = node_df[['x', 'y']].values x = node_values[:, 0] y = node_values[:, 1] heading = node_df['heading'].values if node_df.iloc[0][ 'type'] == env.NodeType.VEHICLE and not node_id == 'ego': # Kalman filter Agent vx = derivative_of(x, scene.dt) vy = derivative_of(y, scene.dt) velocity = np.linalg.norm(np.stack((vx, vy), axis=-1), axis=-1) filter_veh = NonlinearKinematicBicycle(dt=scene.dt, sMeasurement=1.0) P_matrix = None for i in range(len(x)): if i == 0: # initalize KF # initial P_matrix P_matrix = np.identity(4) elif i < len(x): # assign new est values x[i] = x_vec_est_new[0][0] y[i] = x_vec_est_new[1][0] heading[i] = x_vec_est_new[2][0] velocity[i] = x_vec_est_new[3][0] if i < len(x) - 1: # no action on last data # filtering x_vec_est = np.array([[x[i]], [y[i]], [heading[i]], [velocity[i]]]) z_new = np.array([[x[i + 1]], [y[i + 1]], [heading[i + 1]], [velocity[i + 1]]]) x_vec_est_new, P_matrix_new = filter_veh.predict_and_update( x_vec_est=x_vec_est, u_vec=np.array([[0.], [0.]]), P_matrix=P_matrix, z_new=z_new) P_matrix = P_matrix_new curvature, pl, _ = trajectory_curvature(np.stack((x, y), axis=-1)) if pl < 1.0: # vehicle is "not" moving x = x[0].repeat(max_timesteps + 1) y = y[0].repeat(max_timesteps + 1) heading = heading[0].repeat(max_timesteps + 1) global total global curv_0_2 global curv_0_1 total += 1 if pl > 1.0: if curvature > .2: curv_0_2 += 1 node_frequency_multiplier = 3 * int( np.floor(total / curv_0_2)) elif curvature > .1: curv_0_1 += 1 node_frequency_multiplier = 3 * int( np.floor(total / curv_0_1)) if half_dt: t_old = np.linspace(0, 1, x.shape[0]) t_new = np.linspace(0, 1, 2 * x.shape[0]) x = np.interp(t_new, t_old, x) y = np.interp(t_new, t_old, y) heading = np.interp( t_new, t_old, heading + np.pi, period=2 * np.pi) - np.pi vx = derivative_of(x, scene.dt / 2) vy = derivative_of(y, scene.dt / 2) ax = derivative_of(vx, scene.dt / 2) ay = derivative_of(vy, scene.dt / 2) else: vx = derivative_of(x, scene.dt) vy = derivative_of(y, scene.dt) ax = derivative_of(vx, scene.dt) ay = derivative_of(vy, scene.dt) if node_df.iloc[0]['type'] == env.NodeType.VEHICLE: v = np.stack((vx, vy), axis=-1) v_norm = np.linalg.norm(np.stack((vx, vy), axis=-1), axis=-1, keepdims=True) heading_v = np.divide(v, v_norm, out=np.zeros_like(v), where=(v_norm > 1.)) heading_x = heading_v[:, 0] heading_y = heading_v[:, 1] a_norm = np.divide(ax * vx + ay * vy, v_norm[..., 0], out=np.zeros_like(ax), where=(v_norm[..., 0] > 1.)) if half_dt: d_heading = derivative_of(heading, scene.dt / 2, radian=True) else: d_heading = derivative_of(heading, scene.dt, radian=True) data_dict = { ('position', 'x'): x, ('position', 'y'): y, ('velocity', 'x'): vx, ('velocity', 'y'): vy, ('velocity', 'norm'): np.linalg.norm(np.stack((vx, vy), axis=-1), axis=-1), ('acceleration', 'x'): ax, ('acceleration', 'y'): ay, ('acceleration', 'norm'): a_norm, ('heading', 'x'): heading_x, ('heading', 'y'): heading_y, ('heading', '°'): heading, ('heading', 'd°'): d_heading } node_data = pd.DataFrame(data_dict, columns=data_columns_vehicle) else: data_dict = { ('position', 'x'): x, ('position', 'y'): y, ('velocity', 'x'): vx, ('velocity', 'y'): vy, ('acceleration', 'x'): ax, ('acceleration', 'y'): ay } node_data = pd.DataFrame(data_dict, columns=data_columns_pedestrian) node = Node(node_type=node_df.iloc[0]['type'], node_id=0 if node_id == 'ego' else abs(hash(node_id)), data=node_data, frequency_multiplier=node_frequency_multiplier, description=node_id) node.first_timestep = node_df['frame_id'].iloc[0] if half_dt: node.first_timestep *= 2 if node_df.iloc[0]['robot'] == True: node.is_robot = True scene.robot = node scene.nodes.append(node) if half_dt: # We are interpolating to make the overall dt half of what it was. scene.dt /= 2.0 scene.timesteps *= 2 return scene
def get_info(self, sample_token: str, ego_pos: np.ndarray = None, ego_quat: np.ndarray = None, instance_based=False): sample = self.nusc.get('sample', sample_token) scene = self.nusc.get('scene', sample['scene_token']) scene_log = self.nusc.get('log', scene['log_token']) nusc_map = NuScenesMap(dataroot=self.dataroot, map_name=scene_log['location']) sample_data = self.nusc.get('sample_data', sample['data']['CAM_FRONT']) if ego_pos is None: ego_pose = self.nusc.get('ego_pose', sample_data['ego_pose_token']) ego_pos = np.array(ego_pose['translation']) if ego_quat is None: ego_pose = self.nusc.get('ego_pose', sample_data['ego_pose_token']) ego_quat = np.array(ego_pose['rotation']) ego_yaw = Quaternion(ego_quat) ego_yaw = quaternion_yaw(ego_yaw) ego_yaw_rad = angle_of_rotation(ego_yaw) ego_yaw_degrees = np.rad2deg(ego_yaw_rad) ego_on_road_objects = None if self.agent_road_objects: ego_on_road_objects = self.get_road_objects(ego_pos, nusc_map) ego_info = { 'translation': ego_pos, 'rotation_deg': ego_yaw_degrees, 'rotation_quat': ego_quat, 'velocity': 0, 'acceleration': 0, 'heading_change_rate': 0, 'road_objects': ego_on_road_objects } #### define patch #### sensing_patch_width = self.config['sensing_patch_size'][0] sensing_patch_length = self.config['sensing_patch_size'][1] # patch_center_before_rotation = np.array([ego_pos[0], # ego_pos[1] + sensing_patch_length/2]) # lower_left_before_rotation = np.array([ego_pos[0] - sensing_patch_width/2, # ego_pos[1] - sensing_patch_length/2]) # sensing_patch = self.get_patch_coord(patch_box=(patch_center_before_rotation[0], # patch_center_before_rotation[1], # sensing_patch_length, # sensing_patch_width), # rotate_center=(ego_pos[0], ego_pos[1]), # patch_angle=-ego_yaw_degrees) patch_center_before_rotation = np.array([ego_pos[0], ego_pos[1]]) lower_left_before_rotation = np.array([ ego_pos[0] - sensing_patch_width / 2, ego_pos[1] - sensing_patch_length / 2 ]) sensing_patch_coord_before_rotation = [ ego_pos[0] - sensing_patch_width / 2, ego_pos[1] - sensing_patch_length / 2, ego_pos[0] + sensing_patch_width / 2, ego_pos[1] + sensing_patch_length / 2 ] ## generate sensing patch mesh x = np.arange(sensing_patch_coord_before_rotation[0], sensing_patch_coord_before_rotation[2], 0.2) y = np.arange(sensing_patch_coord_before_rotation[1], sensing_patch_coord_before_rotation[3], 0.2) X, Y = np.meshgrid(x, y) ### apply rotation X, Y = rotate_mesh2D(pos=ego_pos, rot_rad=ego_yaw_rad, X=X, Y=Y, frame='current') ## generate sensing patch shapely polygon sensing_patch = self.get_patch_coord( patch_box=(patch_center_before_rotation[0], patch_center_before_rotation[1], sensing_patch_length, sensing_patch_width), rotate_center=(ego_pos[0], ego_pos[1]), patch_angle=-ego_yaw_degrees) sensing_patch_info = {'mesh': [X, Y], 'polygon': sensing_patch} # get agent_info in patch agent_info = self.get_agent_info(sample, sensing_patch, nusc_map) # get map info in patch map_info = self.get_map_info(ego_pos, ego_yaw_degrees, nusc_map, sensing_patch) # get canbus info can_info = None if not instance_based: can_info = self.get_can_info(scene['name']) sensing_info = { 'sample_token': sample_token, 'sensing_patch': sensing_patch_info, 'ego_info': ego_info, 'agent_info': agent_info, 'map_info': map_info, 'can_info': can_info } return sensing_info
cam_path, boxes, camera_intrinsic = nusc.get_sample_data( sd_rec['token'], ) im = Image.open(cam_path) im.save( data_directory + spl + str(scene_id) + camera_channel + '/img1/' + str(count_frames) + '.jpg', 'JPEG', ) TL_found = False bb = [] patch_radius = 700 sample_record = nusc.get('sample', sample_token) log_record = nusc.get('log', scene_record['log_token']) log_location = log_record['location'] nusc_map = NuScenesMap( dataroot=opt.nuscenes_root, map_name=log_location, ) img = transforms.ToTensor()(Image.open(cam_path)) # Pad to square resolution img, _ = pad_to_square(img, 0) # Resize img = resize(img, 416) input_imgs = Variable(img.type(Tensor)).unsqueeze(0) cs_record = nusc.get( 'calibrated_sensor', cam_record['calibrated_sensor_token'], ) poserecord = nusc.get('ego_pose', cam_record['ego_pose_token']) ego_pose = poserecord['translation'] if first_time == True: first_rec = cs_record.copy()
def test_init_all_maps(self): for my_map in ['singapore-onenorth', 'singapore-hollandvillage', 'singapore-queenstown', 'boston-seaport']: nusc_map = NuScenesMap(map_name=my_map)
def process_once(scene_name_list=[], data_save_dir=None, config={}): """get raw data from dataset, everything are lists :param scene_name_list: list of scene names to process :param data_save_dir: directory to save processed raw data :param config: additional configs :returns: one pandas dataframe for each scene """ test_scene_name = config['scene_name'] if config['num_workers'] == 1: nusc = config['other_configs']['nusc'] helper = config['other_configs']['helper'] rasterizer = config['other_configs']['rasterizer'] else: nusc = ray.get(config['other_configs']['nusc']) helper = ray.get(config['other_configs']['helper']) rasterizer = ray.get(config['other_configs']['rasterizer']) dataroot = config['other_configs']['dataroot'] nusc_can = NuScenesCanBus(dataroot=dataroot) df_dict = { #### scene info #### 'scene_name': [], 'scene_token':[], 'scene_nbr_samples':[], 'scene_description':[], 'scene_log_token':[], 'scene_map_token':[], 'scene_vehicle_id': [], 'scene_date_captured':[], 'scene_location':[], #### sample info #### 'sample_idx': [], 'sample_token':[], 'time_stamp': [], #### ego info #### 'current_ego_pos': [], 'current_ego_quat':[], 'current_ego_accel':[], 'current_ego_speed':[], 'current_ego_steering': [], 'current_ego_on_road_objects':[], 'current_ego_raster_img':[], 'current_ego_raster_img_path': [], #### instance info #### 'annotation_token': [], 'instance_token': [], 'instance_category': [], 'instance_attributes': [], 'instance_nbr_annotations':[], 'current_instance_pos': [], 'current_instance_quat': [], 'current_instance_speed': [], 'current_instance_accel':[], 'current_instance_on_road_objects':[], 'current_instance_raster_img':[], 'current_instance_raster_img_path': [] } #### loop over scenes #### for scene_name in tqdm.tqdm(scene_name_list, 'processing scene'): if test_scene_name is not None: if not scene_name == test_scene_name: continue scene_df_dict = copy.deepcopy(df_dict) scene_token = nusc.field2token('scene', 'name', scene_name)[0] scene = nusc.get('scene', scene_token) scene_log = nusc.get('log', scene['log_token']) scene_map = nusc.get('map', scene_log['map_token']) nusc_map = NuScenesMap(dataroot=dataroot, map_name=scene_log['location']) raster_dir = os.path.join(data_save_dir, scene_name) if not os.path.isdir(raster_dir): os.makedirs(raster_dir) # can bus information if scene_name in NO_CANBUS_SCENES: ego_accel_traj = [None] * 50 ego_rotation_rate_traj = [None] * 50 ego_speed_traj = [None] * 50 else: pose_msg = nusc_can.get_messages(scene_name, 'pose') ego_accel_traj = [pm['accel'] for pm in pose_msg][::25] ego_rotation_rate_traj = [pm['rotation_rate'] for pm in pose_msg][::25] ego_speed_traj = [pm['vel'] for pm in pose_msg][::25] #### loop over samples #### sample_token = scene['first_sample_token'] sample = nusc.get('sample', sample_token) sample_idx = 0 t = time.time() while sample['next'] != '': # print(sample_idx) # print(t-time.time()) # t = time.time() cam_data = nusc.get('sample_data', sample['data']['CAM_FRONT']) ego_pose = nusc.get('ego_pose', cam_data['ego_pose_token']) # ego raster ego_raster = rasterizer.make_input_representation(instance_token=None, sample_token=sample_token, ego=True, ego_pose=ego_pose, include_history=False) ego_rel_path = str(sample_idx) + "_" + 'ego.png' current_ego_raster_img_path = os.path.join(raster_dir, ego_rel_path) ego_raster_png = im.fromarray(ego_raster) ego_raster_png.save(current_ego_raster_img_path) # ego road objects ego_on_road_objects = nusc_map.layers_on_point(ego_pose['translation'][0], ego_pose['translation'][1]) #### loop over annotations #### for ann_token in sample['anns']: ann = nusc.get('sample_annotation', ann_token) instance = nusc.get('instance', ann['instance_token']) if filter_instance(ann, nusc): #### populate data #### ## scene ## scene_df_dict['scene_name'].append(scene_name) scene_df_dict['scene_token'].append(scene_token) scene_df_dict['scene_nbr_samples'].append(scene['nbr_samples']) scene_df_dict['scene_description'].append(scene['description']) scene_df_dict['scene_log_token'].append(scene_log['token']) scene_df_dict['scene_map_token'].append(scene_map['token']) scene_df_dict['scene_vehicle_id'].append(scene_log['vehicle']) scene_df_dict['scene_date_captured'].append(scene_log['date_captured']) scene_df_dict['scene_location'].append(scene_log['location']) ## sample info ## scene_df_dict['sample_idx'].append(sample_idx) scene_df_dict['sample_token'].append(sample_token) scene_df_dict['time_stamp'].append(sample['timestamp']) ## ego info ## scene_df_dict['current_ego_pos'].append(ego_pose['translation']) scene_df_dict['current_ego_quat'].append(ego_pose['rotation']) idx = min(sample_idx, len(ego_speed_traj)-1) scene_df_dict['current_ego_speed'].append(ego_speed_traj[idx]) idx = min(sample_idx, len(ego_accel_traj)-1) scene_df_dict['current_ego_accel'].append(ego_accel_traj[idx]) idx = min(sample_idx, len(ego_rotation_rate_traj)-1) scene_df_dict['current_ego_steering'].append(ego_rotation_rate_traj[idx]) scene_df_dict['current_ego_on_road_objects'].append(ego_on_road_objects) scene_df_dict['current_ego_raster_img'].append(ego_raster) scene_df_dict['current_ego_raster_img_path'].append(scene_name+"/"+ego_rel_path) ## instance info ## scene_df_dict['annotation_token'].append(ann_token) scene_df_dict['instance_token'].append(ann['instance_token']) scene_df_dict['instance_category'].append(ann['category_name']) instance_attributes = [nusc.get('attribute', attribute_token)['name'] for attribute_token in ann['attribute_tokens']] scene_df_dict['instance_attributes'].append(instance_attributes) scene_df_dict['instance_nbr_annotations'].append(instance['nbr_annotations']) scene_df_dict['current_instance_pos'].append(ann['translation']) scene_df_dict['current_instance_quat'].append(ann['rotation']) instance_speed = helper.get_velocity_for_agent(ann['instance_token'], sample_token) scene_df_dict['current_instance_speed'].append(instance_speed) instance_accel = helper.get_acceleration_for_agent(ann['instance_token'], sample_token) scene_df_dict['current_instance_accel'].append(instance_accel) instance_on_road_objects = nusc_map.layers_on_point(ann['translation'][0], ann['translation'][1]) scene_df_dict['current_instance_on_road_objects'].append(instance_on_road_objects) instance_raster = rasterizer.make_input_representation(instance_token=ann['instance_token'], sample_token=sample_token, ego=False, include_history=False) scene_df_dict['current_instance_raster_img'].append(instance_raster) instance_rel_path = str(sample_idx) + "_"+ann['instance_token']+'.png' instance_raster_img_path = os.path.join(raster_dir, instance_rel_path) instance_raster_png = im.fromarray(instance_raster) instance_raster_png.save(instance_raster_img_path) scene_df_dict['current_instance_raster_img_path'].append(scene_name+"/"+instance_rel_path) sample_token = sample['next'] sample = nusc.get('sample', sample_token) sample_idx += 1 scene_df = pd.DataFrame(scene_df_dict) scene_df.to_pickle(data_save_dir+"/"+scene_name+".pkl")
def __getitem__(self, idx): sample = self.samples[idx] sample_token = sample['token'] # 1. calculate ego pose sample_data_lidar = self.nus.get('sample_data', sample['data']['LIDAR_TOP']) ego_pose = self.nus.get('ego_pose', sample_data_lidar['ego_pose_token']) ego_pose_xy = ego_pose['translation'] ego_pose_rotation = ego_pose['rotation'] # 타임스탬프 timestamp = ego_pose['timestamp'] # 2. Generate Map & Agent Masks scene = self.nus.get('scene', sample['scene_token']) log = self.nus.get('log', scene['log_token']) location = log['location'] nus_map = NuScenesMap(dataroot=self.root, map_name=location) static_layer = StaticLayerRasterizer(self.helper, layer_names=self.layer_names, colors=self.colors) agent_layer = AgentBoxesWithFadedHistory( self.helper, seconds_of_history=self.agent_seconds) map_masks, lanes, map_img = static_layer.generate_mask( ego_pose_xy, ego_pose_rotation, sample_token) agent_mask = agent_layer.generate_mask(ego_pose_xy, ego_pose_rotation, sample_token) lane_tokens = list(lanes.keys()) num_lanes = len(lane_tokens) lanes_disc = [np.array(lanes[token])[:, :2] for token in lane_tokens] lanes_arc = [] for seq in lanes_disc: seq_len = len(seq) lane_middle = seq[seq_len // 2] opt_middle = (seq[0] + seq[-1]) / 2 lane_h = np.linalg.norm(lane_middle - opt_middle) lane_w = np.linalg.norm(seq[-1] - seq[0]) curve = lane_h / 2 + lane_w**2 / (8 * lane_h) lanes_arc.append(curve) lanes_arc = np.array(lanes_arc) # 3. Generate Agent Trajectory annotation_tokens = sample['anns'] num_agent = len(annotation_tokens) agents = [] for ans_token in annotation_tokens: agent_states = [] agent = self.nus.get('sample_annotation', ans_token) instance_token = agent['instance_token'] # 에이전트 주행경로 xy_global = agent['translation'] past_xy_global = self.helper.get_past_for_agent( instance_token, sample_token, seconds=self.seconds, in_agent_frame=False) future_xy_global = self.helper.get_future_for_agent( instance_token, sample_token, seconds=self.seconds, in_agent_frame=False) # 경로 곡률 agents_arc = [] for seq in [past_xy_global, future_xy_global]: seq_len = len(seq) if seq_len < 2: continue path_middle = seq[seq_len // 2] opt_middle = (seq[0] + seq[-1]) / 2 path_h = np.linalg.norm(path_middle - opt_middle) path_w = np.linalg.norm(seq[-1] - seq[0]) if path_h == 0: path_h = 0.001 if path_w == 0: path_w = 0.001 curve = path_h / 2 + (path_w * path_w) / (8 * path_h) agents_arc.append(curve) agents_arc = np.array(agents_arc) # 로컬 주행경로 xy_local = convert_global_coords_to_local( np.array([xy_global[:2]]), ego_pose_xy, ego_pose_rotation) if len(past_xy_global) < 1: past_xy_global = np.append(past_xy_global, [0., 0.]) if len(future_xy_global) < 1: future_xy_global = np.append(future_xy_global, [0., 0.]) past_xy_local = convert_global_coords_to_local( past_xy_global, ego_pose_xy, ego_pose_rotation) future_xy_local = convert_global_coords_to_local( future_xy_global, ego_pose_xy, ego_pose_rotation) # 에이전트 주행상태 rot = agent['rotation'] vel = self.helper.get_velocity_for_agent(instance_token, sample_token) accel = self.helper.get_acceleration_for_agent( instance_token, sample_token) yaw_rate = self.helper.get_heading_change_rate_for_agent( instance_token, sample_token) agent_states = { 'present_pos': xy_global, 'past_pos': past_xy_global, 'future_pos': future_xy_global, 'rot': rot, 'vel': vel, 'accel': accel, 'yaw_rate': yaw_rate, 'present_local_xy': xy_local, 'past_local_xy': past_xy_local, 'future_local_xy': future_xy_local, 'curvature': agents_arc } agents.append(agent_states) return map_masks, agent_mask, agents, idx