def vis_scene_data(nuscenes_path=None, nuscenes_version='v1.0-trainval', which_scene=0, max_seq_num=10, begin_frame=0, frame_skip=3, trained_model_path=None, img_save_dir=None, which_model='MotionNet', use_adj_frame_pred=True, use_motion_state_pred_masking=True, disp=True): """ Visualize the scene data. nuscenes_path: the path to the nuScenes dataset nuscenes_version: the dataset version ['v1.0-trainval'/'v1.0-mini'] which_scene: for which we want to visualize max_seq_num: how many frames want to visualize begin_frame: for this scene, from which frame we want to start our prediction frame_skip: how many future frames we want to skip. This is used for generated preprocessed bev data. trained_model_path: the path to the pretrained model img_save_dir: the directory for saving the predicted image which_model: which network ['MotionNet'/'MotionNetMGDA'] use_adj_frame_pred: whether to predict the relative offsets between two adjacent frames use_motion_state_pred_masking: whether to threshold the prediction with motion state estimation results disp: whether to immediately show the predicted results """ if nuscenes_path is None: raise ValueError("Should specify the nuScenes data path.") nusc = NuScenes(version=nuscenes_version, dataroot=nuscenes_path, verbose=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") nsweeps_back = 20 nsweeps_forward = 20 num_frame_skipped = 0 voxel_size = (0.25, 0.25, 0.4) area_extents = np.array([[-32., 32.], [-32., 32.], [-3, 2]]) sample_cnt = 1 class_map = { 'vehicle.car': 1, 'vehicle.bus.rigid': 1, 'vehicle.bus.bendy': 1, 'human.pedestrian': 2, 'vehicle.bicycle': 3 } # background: 0, other: 4 curr_scene = nusc.scene[which_scene] first_sample_token = curr_scene['first_sample_token'] last_sample_token = curr_scene['last_sample_token'] curr_sample = nusc.get('sample', first_sample_token) curr_sample_data = nusc.get('sample_data', curr_sample['data']['LIDAR_TOP']) has_reached_last_keyframe = False seq_num = 0 # Load pre-trained network weights loaded_models = list() if which_model == "MotionNet": model = MotionNet(out_seq_len=20, motion_category_num=2, height_feat_size=13) model = nn.DataParallel(model) checkpoint = torch.load(trained_model_path) model.load_state_dict(checkpoint['model_state_dict']) model = model.to(device) loaded_models = [model] else: model_encoder = FeatEncoder() model_head = MotionNetMGDA(out_seq_len=20, motion_category_num=2) model_encoder = nn.DataParallel(model_encoder) model_head = nn.DataParallel(model_head) checkpoint = torch.load(trained_model_path) model_encoder.load_state_dict(checkpoint['encoder_state_dict']) model_head.load_state_dict(checkpoint['head_state_dict']) model_encoder = model_encoder.to(device) model_head = model_head.to(device) loaded_models = [model_encoder, model_head] print("Loaded pretrained model {}".format(which_model)) while curr_sample_data['next'] != '': if has_reached_last_keyframe: break # has reached the final keyframe if curr_sample_data['token'] == last_sample_token: has_reached_last_keyframe = True # Skip current keyframe if possible if num_frame_skipped > 0 and sample_cnt % (num_frame_skipped + 1) == 0: sample_cnt += 1 curr_sample_data = nusc.get('sample_data', curr_sample_data['next']) continue # Get the synchronized point clouds all_pc, all_times = LidarPointCloud.from_file_multisweep_bf_sample_data( nusc, curr_sample_data, nsweeps_back=nsweeps_back, nsweeps_forward=nsweeps_forward) # Store point cloud of each sweep pc = all_pc.points _, sort_idx = np.unique(all_times, return_index=True) unique_times = all_times[np.sort( sort_idx)] # Preserve the item order in unique_times num_sweeps = len(unique_times) # Make sure we have sufficient past and future sweeps if num_sweeps != (nsweeps_back + nsweeps_forward): sample_cnt += 1 curr_sample_data = nusc.get('sample_data', curr_sample_data['next']) continue # Prepare data dictionary for visualization save_data_dict = dict() for tid in range(num_sweeps): _time = unique_times[tid] points_idx = np.where(all_times == _time)[0] _pc = pc[:, points_idx] save_data_dict['pc_' + str(tid)] = _pc save_data_dict['times'] = unique_times save_data_dict['num_sweeps'] = num_sweeps # Get the synchronized bounding boxes # First, we need to iterate all the instances, and then retrieve their corresponding bounding boxes num_instances = 0 # The number of instances within this sample corresponding_sample_token = curr_sample_data['sample_token'] corresponding_sample_rec = nusc.get('sample', corresponding_sample_token) for ann_token in corresponding_sample_rec['anns']: ann_rec = nusc.get('sample_annotation', ann_token) category_name = ann_rec['category_name'] flag = False for c, v in class_map.items(): if category_name.startswith(c): save_data_dict['category_' + str(num_instances)] = v flag = True break if not flag: save_data_dict['category_' + str(num_instances)] = 4 # Other category instance_token = ann_rec['instance_token'] instance_boxes, instance_all_times, _, _ = LidarPointCloud. \ get_instance_boxes_multisweep_sample_data(nusc, curr_sample_data, instance_token, nsweeps_back=nsweeps_back, nsweeps_forward=nsweeps_forward) assert np.array_equal( unique_times, instance_all_times ), "The sweep and instance times are not consistent!" assert num_sweeps == len( instance_boxes ), "The number of instance boxes does not match that of sweeps!" # Each row corresponds to a box annotation; the column consists of box center, box size, and quaternion box_data = np.zeros((len(instance_boxes), 3 + 3 + 4), dtype=np.float32) box_data.fill(np.nan) for r, box in enumerate(instance_boxes): if box is not None: row = np.concatenate( [box.center, box.wlh, box.orientation.elements]) box_data[r] = row[:] # Save the box data for current instance save_data_dict['instance_boxes_' + str(num_instances)] = box_data num_instances += 1 save_data_dict['num_instances'] = num_instances if seq_num < begin_frame: seq_num += 1 sample_cnt += 1 print("Finish loading sequence sample {}".format(seq_num)) continue # ------------------------------------ Visualization ------------------------------------ # -- The following code is simply borrowed from gen_data.py and currently not optimized num_sweeps = save_data_dict['num_sweeps'] times = save_data_dict['times'] num_past_sweeps = len(np.where(times >= 0)[0]) num_future_sweeps = len(np.where(times < 0)[0]) assert num_past_sweeps + num_future_sweeps == num_sweeps, "The number of sweeps is incorrect!" # Load point cloud pc_list = [] for i in range(num_sweeps): pc = save_data_dict['pc_' + str(i)] pc_list.append(pc.T) # Reorder the pc, and skip sample frames if wanted tmp_pc_list_1 = pc_list[0:num_past_sweeps:(frame_skip + 1)] tmp_pc_list_1 = tmp_pc_list_1[::-1] tmp_pc_list_2 = pc_list[(num_past_sweeps + frame_skip)::(frame_skip + 1)] pc_list = tmp_pc_list_1 + tmp_pc_list_2 num_past_pcs = len(tmp_pc_list_1) num_future_pcs = len(tmp_pc_list_2) # Voxelize the input point clouds, and compute the ground truth displacement vectors padded_voxel_points_list = list( ) # This contains the compact representation of voxelization, as in the paper for i in range(num_past_pcs): res = voxelize_occupy(pc_list[i], voxel_size=voxel_size, extents=area_extents) padded_voxel_points_list.append(res) # Compile the batch of voxels, so that they can be fed into the network padded_voxel_points = torch.from_numpy( np.stack(padded_voxel_points_list, axis=0)) # Finally, generate the ground-truth displacement field all_disp_field_gt, all_valid_pixel_maps, non_empty_map, pixel_cat_map \ = gen_2d_grid_gt_for_visualization(save_data_dict, grid_size=voxel_size[0:2], reordered=True, extents=area_extents, frame_skip=frame_skip) bev_input_data = (padded_voxel_points, all_disp_field_gt, all_valid_pixel_maps, non_empty_map, pixel_cat_map, num_past_pcs, num_future_pcs) vis_model_per_sample_data( bev_input_data, save_data_dict, frame_skip=frame_skip, loaded_models=loaded_models, voxel_size=voxel_size, which_model=which_model, model_path=trained_model_path, img_save_dir=img_save_dir, use_adj_frame_pred=use_adj_frame_pred, disp=disp, use_motion_state_pred_masking=use_motion_state_pred_masking, frame_idx=seq_num) seq_num += 1 print("Finish loading sequence sample {}".format(seq_num)) sample_cnt += 1 curr_sample_data = nusc.get('sample_data', curr_sample_data['next']) if seq_num - begin_frame >= max_seq_num: break
def gen_data(): res_scenes = list() for s in scenes: s_id = s.split('_')[1] res_scenes.append(int(s_id)) for scene_idx in res_scenes: curr_scene = nusc.scene[scene_idx] first_sample_token = curr_scene['first_sample_token'] curr_sample = nusc.get('sample', first_sample_token) curr_sample_data = nusc.get('sample_data', curr_sample['data']['LIDAR_TOP']) save_data_dict_list = list( ) # for storing consecutive sequences; the data consists of timestamps, points, etc save_box_dict_list = list( ) # for storing box annotations in consecutive sequences save_instance_token_list = list() adj_seq_cnt = 0 save_seq_cnt = 0 # only used for save data file name # Iterate each sample data print("Processing scene {} ...".format(scene_idx)) while curr_sample_data['next'] != '': # Get the synchronized point clouds all_pc, all_times, trans_matrices = \ LidarPointCloud.from_file_multisweep_bf_sample_data(nusc, curr_sample_data, return_trans_matrix=True, nsweeps_back=nsweeps_back, nsweeps_forward=nsweeps_forward) # Store point cloud of each sweep pc = all_pc.points _, sort_idx = np.unique(all_times, return_index=True) unique_times = all_times[np.sort( sort_idx)] # Preserve the item order in unique_times num_sweeps = len(unique_times) # Make sure we have sufficient past and future sweeps if num_sweeps != (nsweeps_back + nsweeps_forward): # Skip some keyframes if necessary flag = False for _ in range(num_keyframe_skipped + 1): if curr_sample['next'] != '': curr_sample = nusc.get('sample', curr_sample['next']) else: flag = True break if flag: # No more keyframes break else: curr_sample_data = nusc.get( 'sample_data', curr_sample['data']['LIDAR_TOP']) # Reset adj_seq_cnt = 0 save_data_dict_list = list() save_box_dict_list = list() save_instance_token_list = list() continue # Prepare data dictionary for the next step (ie, generating BEV maps) save_data_dict = dict() box_data_dict = dict( ) # for remapping the instance ids, according to class_map curr_token_list = list() for tid in range(num_sweeps): _time = unique_times[tid] points_idx = np.where(all_times == _time)[0] _pc = pc[:, points_idx] save_data_dict['pc_' + str(tid)] = _pc save_data_dict['times'] = unique_times save_data_dict['num_sweeps'] = num_sweeps save_data_dict['trans_matrices'] = trans_matrices # Get the synchronized bounding boxes # First, we need to iterate all the instances, and then retrieve their corresponding bounding boxes num_instances = 0 # The number of instances within this sample corresponding_sample_token = curr_sample_data['sample_token'] corresponding_sample_rec = nusc.get('sample', corresponding_sample_token) for ann_token in corresponding_sample_rec['anns']: ann_rec = nusc.get('sample_annotation', ann_token) category_name = ann_rec['category_name'] instance_token = ann_rec['instance_token'] flag = False for c, v in class_map.items(): if category_name.startswith(c): box_data_dict['category_' + instance_token] = v flag = True break if not flag: box_data_dict['category_' + instance_token] = 4 # Other category instance_boxes, instance_all_times, _, _ = LidarPointCloud. \ get_instance_boxes_multisweep_sample_data(nusc, curr_sample_data, instance_token, nsweeps_back=nsweeps_back, nsweeps_forward=nsweeps_forward) assert np.array_equal( unique_times, instance_all_times ), "The sweep and instance times are inconsistent!" assert num_sweeps == len( instance_boxes ), "The number of instance boxes does not match that of sweeps!" # Each row corresponds to a box annotation; the column consists of box center, box size, and quaternion box_data = np.zeros((len(instance_boxes), 3 + 3 + 4), dtype=np.float32) box_data.fill(np.nan) for r, box in enumerate(instance_boxes): if box is not None: row = np.concatenate( [box.center, box.wlh, box.orientation.elements]) box_data[r] = row[:] # Save the box data for current instance box_data_dict['instance_boxes_' + instance_token] = box_data num_instances += 1 curr_token_list.append(instance_token) save_data_dict['num_instances'] = num_instances save_data_dict_list.append(save_data_dict) save_box_dict_list.append(box_data_dict) save_instance_token_list.append(curr_token_list) # Update the counter and save the data if desired (But here we do not want to # save the data to disk since it would cost about 2TB space) adj_seq_cnt += 1 if adj_seq_cnt == num_adj_seqs: # First, we need to reorganize the instance tokens (ids) num_instance_token_list = len(save_instance_token_list) if num_instance_token_list > 1: common_tokens = set( save_instance_token_list[0]).intersection( save_instance_token_list[1]) for l in range(2, num_instance_token_list): common_tokens = common_tokens.intersection( save_instance_token_list[l]) for l in range(num_instance_token_list): exclusive_tokens = set( save_instance_token_list[l]).difference( common_tokens) # we store the common instances first, then store the remaining instances curr_save_data_dict = save_data_dict_list[l] curr_save_box_dict = save_box_dict_list[l] counter = 0 for token in common_tokens: box_info = curr_save_box_dict['instance_boxes_' + token] box_cat = curr_save_box_dict['category_' + token] curr_save_data_dict['instance_boxes_' + str(counter)] = box_info curr_save_data_dict['category_' + str(counter)] = box_cat counter += 1 for token in exclusive_tokens: box_info = curr_save_box_dict['instance_boxes_' + token] box_cat = curr_save_box_dict['category_' + token] curr_save_data_dict['instance_boxes_' + str(counter)] = box_info curr_save_data_dict['category_' + str(counter)] = box_cat counter += 1 assert counter == curr_save_data_dict[ 'num_instances'], "The number of instances is inconsistent." save_data_dict_list[l] = curr_save_data_dict else: curr_save_box_dict = save_box_dict_list[0] curr_save_data_dict = save_data_dict_list[0] for index, token in enumerate(save_instance_token_list[0]): box_info = curr_save_box_dict['instance_boxes_' + token] box_cat = curr_save_box_dict['category_' + token] curr_save_data_dict['instance_boxes_' + str(index)] = box_info curr_save_data_dict['category_' + str(index)] = box_cat save_data_dict_list[0] = curr_save_data_dict # ------------------------ Now we generate dense BEV maps ------------------------ for seq_idx, seq_data_dict in enumerate(save_data_dict_list): dense_bev_data = convert_to_dense_bev(seq_data_dict) sparse_bev_data = convert_to_sparse_bev(dense_bev_data) # save the data save_directory = os.path.join( args.savepath, str(scene_idx) + '_' + str(save_seq_cnt)) os.makedirs(save_directory, exist_ok=True) save_file_name = os.path.join(save_directory, str(seq_idx) + '.npy') np.save(save_file_name, arr=sparse_bev_data) print(" >> Finish sample: {}, sequence {}".format( save_seq_cnt, seq_idx)) # -------------------------------------------------------------------------------- save_seq_cnt += 1 adj_seq_cnt = 0 save_data_dict_list = list() save_box_dict_list = list() save_instance_token_list = list() # Skip some keyframes if necessary flag = False for _ in range(num_keyframe_skipped + 1): if curr_sample['next'] != '': curr_sample = nusc.get('sample', curr_sample['next']) else: flag = True break if flag: # No more keyframes break else: curr_sample_data = nusc.get( 'sample_data', curr_sample['data']['LIDAR_TOP']) else: flag = False for _ in range(skip_frame + 1): if curr_sample_data['next'] != '': curr_sample_data = nusc.get('sample_data', curr_sample_data['next']) else: flag = True break if flag: # No more sample frames break