def nuscenes_map_to_line_representation( nusc_map: NuScenesMap, patch: List[float], realign: bool = False) -> Tuple[torch.Tensor]: record = nusc_map.get_records_in_patch(patch, ["drivable_area"]) pt1, pt2 = [], [] for da_token in record["drivable_area"]: da = nusc_map.get("drivable_area", da_token) for poly in map(nusc_map.extract_polygon, da["polygon_tokens"]): p1, p2 = get_edges_of_polygon_in_patch(poly, patch) if len(p1) > 0 and len(p2) > 0: p1, p2 = preprocess_map_edges( torch.as_tensor(p1), torch.as_tensor(p2), passes=10, tol=0.1, ) pt1.append(p1) pt2.append(p2) pt1, pt2 = torch.cat(pt1), torch.cat(pt2) if realign: pt1, pt2 = realign_map_edges(pt1, pt2, 0.0) centers = (pt1 + pt2) / 2 centers1 = centers.unsqueeze(1) centers2 = centers.unsqueeze(0) dist = (centers2 - centers1).pow(2).sum(dim=-1).sqrt() for i in range(centers.size(0)): dist[i, i] = 1e12 very_close = (dist < 0.01).any(dim=-1) to_remove = [] for i, c in enumerate(very_close): if c: to_remove.append(i) for i, rem in enumerate(to_remove): rem = rem - i pt1 = torch.cat([pt1[:rem], pt1[rem + 1:]]) pt2 = torch.cat([pt2[:rem], pt2[rem + 1:]]) return pt1, pt2
elif scene_id >= 200 and scene_id <= 250: ego_pose[2] = -0.9 elif scene_id >= 250 and scene_id <= 300: ego_pose[2] = -0.5 else: ego_pose[2] = -0.3 box_coords = ( ego_pose[0] - patch_radius, ego_pose[1] - patch_radius, ego_pose[0] + patch_radius, ego_pose[1] + patch_radius, ) records_in_patch = nusc_map.get_records_in_patch( box_coords, layer_names, 'intersect', ) near_plane = 1e-8 # # Retrieve and render each record. taken_already = [] index = -1 for layer_name in layer_names: for token in records_in_patch[layer_name]: record = nusc_map.get(layer_name, token) line = nusc_map.extract_line(record['line_token']) if line.is_empty: # Skip lines without nodes continue xs, ys = line.xy
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