├── LICENSE ├── NUSC-PR ├── configs │ ├── config_self_supervised_scheme.yml │ └── config_supervised_scheme.yml ├── self_supervised │ ├── generate_basic_infos.py │ └── split_dataset.py └── supervised │ ├── generate_basic_infos.py │ ├── generate_selected_indicies.py │ ├── select_pos_neg_samples_by_dis.py │ └── split_dataset.py ├── README.md ├── map_pointcloud_to_image.py └── requirements.txt /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 BIT-XJY 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /NUSC-PR/configs/config_self_supervised_scheme.yml: -------------------------------------------------------------------------------- 1 | locations: ['boston-seaport', 'singapore-onenorth', 'singapore-queenstown', 'singapore-hollandvillage'] 2 | locations_abbr: ['bs', 'son', 'sq', 'shv'] 3 | 4 | num_pos: 2 5 | num_neg: 4 6 | positive_distance_threshold: 9 7 | negative_time_threshold: 6 # 6 frames -> 3 s 8 | use_date_date_threshold: False 9 | date_threshold: 105 # 105 for bs and son 10 | 11 | 12 | data_root: 13 | # nuscenes path 14 | nusc_root: '/media/xjy/T1/nuscenes/raw_data' 15 | # path to basic infos 16 | save_root_basic_infos: '/home/xjy/dev/EINet/NUSC-PR/self_supervised_data/generate_basic_infos' 17 | # path to saved datasets 18 | save_root_splitted_data: '/home/xjy/dev/EINet/NUSC-PR/self_supervised_data/split_dataset' -------------------------------------------------------------------------------- /NUSC-PR/configs/config_supervised_scheme.yml: -------------------------------------------------------------------------------- 1 | # Dataset configs in supervised-learning scheme in NUSC-PR 2 | 3 | locations: ['boston-seaport', 'singapore-onenorth', 'singapore-queenstown', 'singapore-hollandvillage'] 4 | locations_abbr: ['bs', 'son', 'sq', 'shv'] 5 | distance_interval: 1.0 6 | date_threshold_bs: 105 7 | date_threshold_son: 50 8 | date_threshold_sq: 105 9 | date_threshold_shv: 50 10 | val_ratio: 0.25 11 | positive_distance_threshold: 9 12 | negative_distance_threshold: 18 13 | nbr_positive_samples: 2 14 | nbr_negative_samples: 4 15 | 16 | data_root: 17 | # nuscenes path 18 | nusc_root: '/media/xjy/T1/nuscenes/raw_data' 19 | # path to basic infos 20 | save_root_basic_infos: '/home/xjy/dev/EINet/NUSC-PR/supervised_data/generate_basic_infos' 21 | # path to splitted dataset 22 | save_root_splitted_data: '/home/xjy/dev/EINet/NUSC-PR/supervised_data/split_dataset' 23 | # path to mapping from query to reference 24 | save_root_select_by_dis: '/home/xjy/dev/EINet/NUSC-PR/supervised_data/select_pos_neg_samples_by_dis' 25 | # path to selected indices converted from save_root_select_by_dis 26 | save_root_selected_indices: '/home/xjy/dev/EINet/NUSC-PR/supervised_data/generate_selected_indicies' 27 | -------------------------------------------------------------------------------- /NUSC-PR/self_supervised/generate_basic_infos.py: -------------------------------------------------------------------------------- 1 | # Developed by Jingyi Xu, Junyi Ma, Zijie Zhou 2 | # Brief: extract basic infos from nuScenes datasets 3 | # NUSC-PR is proposed in the paper: Explicit Interaction for Fusion-Based Place Recognition 4 | 5 | import os 6 | import pickle 7 | import numpy as np 8 | from tqdm import tqdm 9 | from nuscenes.nuscenes import NuScenes 10 | import yaml 11 | 12 | 13 | def check_dir(*dirs): 14 | for dir in dirs: 15 | if not os.path.exists(dir): 16 | os.makedirs(dir) 17 | 18 | def gen_info(nusc, sample_tokens): 19 | # Collect sample infos into a dict 20 | cam_names = [ 21 | 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_RIGHT', 'CAM_BACK', 22 | 'CAM_BACK_LEFT', 'CAM_FRONT_LEFT'] 23 | lidar_names = ['LIDAR_TOP'] 24 | 25 | infos = list() 26 | for sample_token in tqdm(sample_tokens): 27 | sample = nusc.get('sample', sample_token) 28 | info = dict() 29 | cam_datas = list() 30 | lidar_datas = list() 31 | info['sample_token'] = sample_token 32 | info['timestamp'] = sample['timestamp'] 33 | info['scene_token'] = sample['scene_token'] 34 | 35 | cam_infos = dict() 36 | lidar_infos = dict() 37 | for cam_name in cam_names: 38 | cam_data = nusc.get('sample_data', sample['data'][cam_name]) 39 | cam_datas.append(cam_data) 40 | cam_info = dict() 41 | cam_info['sample_token'] = sample['data'][cam_name] 42 | cam_info['ego_pose'] = nusc.get('ego_pose', cam_data['ego_pose_token']) 43 | cam_info['timestample'] = cam_data['timestamp'] 44 | cam_info['filename'] = cam_data['filename'] 45 | cam_info['calibrated_sensor'] = nusc.get( 46 | 'calibrated_sensor', cam_data['calibrated_sensor_token']) 47 | cam_infos[cam_name] = cam_info 48 | for lidar_name in lidar_names: 49 | lidar_data = nusc.get('sample_data', 50 | sample['data'][lidar_name]) 51 | lidar_datas.append(lidar_data) 52 | lidar_info = dict() 53 | lidar_info['sample_token'] = sample['data'][lidar_name] 54 | lidar_info['ego_pose'] = nusc.get( 55 | 'ego_pose', lidar_data['ego_pose_token']) 56 | lidar_info['timestamp'] = lidar_data['timestamp'] 57 | lidar_info['filename'] = lidar_data['filename'] 58 | lidar_info['calibrated_sensor'] = nusc.get( 59 | 'calibrated_sensor', lidar_data['calibrated_sensor_token']) 60 | lidar_infos[lidar_name] = lidar_info 61 | info['cam_infos'] = cam_infos 62 | info['lidar_infos'] = lidar_infos 63 | scene = nusc.get('scene', sample['scene_token']) 64 | loc_scene = nusc.get('log', scene['log_token'])['location'] 65 | info['loc'] = loc_scene 66 | infos.append(info) 67 | 68 | return infos 69 | 70 | 71 | def get_location_sample_tokens(nusc, location): 72 | # Get the sample tokens of a specific location 73 | 74 | location_indices = get_location_indices(nusc, location) 75 | 76 | sample_token_list = [] 77 | 78 | for scene_index in location_indices: 79 | scene = nusc.scene[scene_index] 80 | sample_token = scene['first_sample_token'] 81 | 82 | while not sample_token == '': 83 | sample = nusc.get('sample', sample_token) 84 | sample_token_list.append(sample_token) 85 | sample_token = sample['next'] 86 | 87 | return sample_token_list 88 | 89 | 90 | def get_location_indices(nusc, location): 91 | # Get the indices of the specific location 92 | 93 | location_indices = [] 94 | for scene_index in range(len(nusc.scene)): 95 | scene = nusc.scene[scene_index] 96 | if nusc.get('log', scene['log_token'])['location'] != location: 97 | continue 98 | location_indices.append(scene_index) 99 | return np.array(location_indices) 100 | 101 | def get_sample_tokens(nusc): 102 | # Get all sample tokens 103 | 104 | sample_token_list = [] 105 | for scene_index in range(len(nusc.scene)): 106 | scene = nusc.scene[scene_index] 107 | sample_token = scene['first_sample_token'] 108 | 109 | while not sample_token == '': 110 | sample = nusc.get('sample', sample_token) 111 | sample_token_list.append(sample_token) 112 | sample_token = sample['next'] 113 | return sample_token_list 114 | 115 | def main(config): 116 | 117 | nusc_root = config['data_root']['nusc_root'] 118 | save_root = config['data_root']['save_root_basic_infos'] 119 | check_dir(save_root) 120 | 121 | nusc_trainval = NuScenes(version='v1.0-trainval', dataroot=nusc_root, verbose=True) 122 | 123 | locations = config['locations'] 124 | locations_ = config['locations_abbr'] 125 | 126 | # ====================generate location infos==================== 127 | for idx, location in enumerate(locations): 128 | sample_tokens = get_location_sample_tokens(nusc_trainval, location=location) 129 | location_infos = gen_info(nusc_trainval, sample_tokens) 130 | with open(os.path.join(save_root, 'nuscenes_infos-'+locations_[idx]+'.pkl'), 'wb') as f: 131 | pickle.dump(location_infos, f) 132 | 133 | # ====================generate all infos==================== 134 | sample_tokens_trainval = get_sample_tokens(nusc_trainval) 135 | infos = gen_info(nusc_trainval, sample_tokens_trainval) 136 | with open(os.path.join(save_root, 'nuscenes_infos.pkl'), 'wb') as f: 137 | pickle.dump(infos, f) 138 | 139 | print("done!") 140 | 141 | 142 | if __name__ == '__main__': 143 | # load config ================================================================ 144 | config_filename = '../configs/config_self_supervised_scheme.yml' 145 | config = yaml.safe_load(open(config_filename)) 146 | # ============================================================================ 147 | main(config) -------------------------------------------------------------------------------- /NUSC-PR/self_supervised/split_dataset.py: -------------------------------------------------------------------------------- 1 | # Developed by Jingyi Xu, Junyi Ma, Zijie Zhou 2 | # Brief: split query and database for NUSC-PR 3 | # NUSC-PR is proposed in the paper: Explicit Interaction for Fusion-Based Place Recognition 4 | 5 | 6 | import numpy as np 7 | from sklearn.neighbors import NearestNeighbors 8 | import os 9 | import random 10 | import pickle 11 | from nuscenes import NuScenes 12 | import yaml 13 | 14 | def check_dir(*dirs): 15 | for dir in dirs: 16 | if not os.path.exists(dir): 17 | os.makedirs(dir) 18 | 19 | def scene_index_to_sample_index(scene_index, scene_nbr_list): 20 | sample_index_start = 0 21 | for i in range(scene_index): 22 | sample_index_start += scene_nbr_list[i] 23 | sample_index_end = sample_index_start + scene_nbr_list[scene_index] - 1 24 | return sample_index_start, sample_index_end 25 | 26 | def load_txt_index(txt_path): 27 | index_list = [] 28 | with open(txt_path, 'r') as f: 29 | index_lines = f.readlines() 30 | for i in range(len(index_lines)): 31 | index_list.append(int(index_lines[i].split('\n')[0])) 32 | return index_list 33 | 34 | def translation_to_info_index(t_list, translation_index, scene_nbr_list): 35 | # t_list为train_list或test_list 36 | index_list = [] 37 | index_all = 0 38 | for i in t_list: 39 | index_all += scene_nbr_list[i] 40 | index_list.append(index_all - 1) 41 | for i in range(len(index_list)): 42 | if translation_index <= index_list[0]: 43 | t_scene_index = t_list[0] 44 | rest_test = translation_index 45 | else: 46 | if (translation_index <= index_list[i]) & (translation_index > index_list[i-1]): 47 | t_scene_index = t_list[i] 48 | rest_test = translation_index - index_list[i-1] - 1 49 | 50 | info_index = 0 51 | for i in range(t_scene_index): 52 | info_index += scene_nbr_list[i] 53 | info_index = info_index + rest_test 54 | return info_index 55 | 56 | def split_test_for_different_loc(infos, database_list_loc, test_list_loc, scene_nbr_list, save_root, loc, positive_distance_threshold): 57 | database_indicies = [] 58 | database_lidar_translation = [] 59 | for i in range(len(database_list_loc)): 60 | database_scene_index = database_list_loc[i] 61 | sample_index_start, sample_index_end = scene_index_to_sample_index(database_scene_index, scene_nbr_list) 62 | 63 | for j in range(scene_nbr_list[database_scene_index]): 64 | sample_index = sample_index_start + j 65 | database_indicies.append(sample_index) 66 | sample_info = infos[sample_index] 67 | lidar_translation = sample_info['lidar_infos']['LIDAR_TOP']['ego_pose']['translation'] 68 | database_lidar_translation.append(lidar_translation) 69 | 70 | np.save(os.path.join(save_root, loc+'_db_index_in_infos.npy'), database_indicies) 71 | 72 | test_lidar_translation = [] 73 | for i in range(len(test_list_loc)): 74 | test_scene_index = test_list_loc[i] 75 | sample_index_start, sample_index_end = scene_index_to_sample_index(test_scene_index, scene_nbr_list) 76 | 77 | for j in range(scene_nbr_list[test_scene_index]): 78 | sample_index = sample_index_start + j 79 | sample_info = infos[sample_index] 80 | lidar_translation = sample_info['lidar_infos']['LIDAR_TOP']['ego_pose']['translation'] 81 | test_lidar_translation.append(lidar_translation) 82 | 83 | knn = NearestNeighbors(n_jobs=-1) 84 | database_lidar_translation = np.array(database_lidar_translation) 85 | test_lidar_translation = np.array(test_lidar_translation) 86 | database = np.ascontiguousarray(database_lidar_translation) 87 | queries = np.ascontiguousarray(test_lidar_translation) 88 | knn.fit(database) 89 | dis, gt_samples = knn.radius_neighbors(queries, radius=positive_distance_threshold, return_distance=True) 90 | 91 | # convert to index in infos 92 | gt_mapping = dict() 93 | for i in range(len(gt_samples)): 94 | gt_index_list = [] 95 | test_query = translation_to_info_index(test_list_loc, i, scene_nbr_list) 96 | for j in range(len(gt_samples[i])): 97 | translation_index = gt_samples[i][j] 98 | gt_index_list.append(translation_to_info_index(database_list_loc, translation_index, scene_nbr_list)) 99 | gt_mapping[test_query] = gt_index_list 100 | 101 | with open(os.path.join(save_root, loc+'_test_query_gt_index_in_infos.pkl'), 'wb') as f: 102 | pickle.dump(gt_mapping, f) 103 | 104 | 105 | def main(config): 106 | 107 | nusc_root = config["data_root"]["nusc_root"] 108 | save_root = config["data_root"]["save_root_splitted_data"] 109 | check_dir(save_root) 110 | infos_root = config['data_root']['save_root_basic_infos'] 111 | num_pos = config["num_pos"] 112 | num_neg = config["num_neg"] 113 | positive_distance_threshold = config["positive_distance_threshold"] 114 | negative_time_threshold = config["negative_time_threshold"] 115 | use_date_date_threshold = config["use_date_date_threshold"] 116 | date_threshold = config["date_threshold"] 117 | locations_ = config['locations_abbr'] 118 | 119 | infos_root = os.path.join(infos_root, 'nuscenes_infos.pkl') 120 | with open(infos_root, 'rb') as f: 121 | infos = pickle.load(f) 122 | 123 | nusc_trainval = NuScenes(version='v1.0-trainval', dataroot=nusc_root, verbose=True) 124 | 125 | scene_nbr_list = [] 126 | old_scenes = [] 127 | new_scenes = [] 128 | timestamps = [] 129 | 130 | for i in range(len(nusc_trainval.scene)): 131 | scene_nbr_list.append(nusc_trainval.scene[i]['nbr_samples']) 132 | current_scene = nusc_trainval.scene[i] 133 | first_sample_token = current_scene['first_sample_token'] 134 | first_sample = nusc_trainval.get('sample', first_sample_token) 135 | timestamp = first_sample['timestamp'] 136 | timestamps.append(timestamp) 137 | timestamps = np.array(timestamps, dtype=np.float32) 138 | timestamps = np.array(timestamps - min(timestamps)) / (3600 * 24 * 1e6) 139 | 140 | if use_date_date_threshold: 141 | for i in range(len(nusc_trainval.scene)): 142 | if timestamps[i] < date_threshold: 143 | old_scenes.append(i) 144 | else: 145 | new_scenes.append(i) 146 | else: 147 | scene_indicies = np.arange(0, len(nusc_trainval.scene),1).tolist() 148 | new_scenes = random.sample(scene_indicies, 150) 149 | old_scenes = list(set(scene_indicies) - set(new_scenes)) 150 | 151 | 152 | print('======================') 153 | print('the number of scenes: ', len(nusc_trainval.scene)) 154 | print('the number of scenes as train queries and database: ', len(old_scenes)) 155 | print('the number of scenes as test queries: ', len(new_scenes)) 156 | print('======================') 157 | 158 | # ====================generate training tuples==================== 159 | print('splitting train data ...') 160 | all_query_dict = dict() 161 | bs_query_dict = dict() 162 | son_query_dict = dict() 163 | sq_query_dict = dict() 164 | shv_query_dict = dict() 165 | 166 | for train_scene_index in old_scenes: 167 | sample_index_start, sample_index_end = scene_index_to_sample_index(train_scene_index, scene_nbr_list) 168 | list_sample = range(sample_index_start, sample_index_end+1) 169 | scene = nusc_trainval.scene[train_scene_index] 170 | loc_scene = nusc_trainval.get('log', scene['log_token'])['location'] 171 | for j in range(scene_nbr_list[train_scene_index]): 172 | query_index = list_sample[j] 173 | 174 | if j < num_pos+num_neg+negative_time_threshold: 175 | continue 176 | else: 177 | pos_neg_dict = dict() 178 | pos_index = [] 179 | for pos in range(num_pos): 180 | pos_index.append(list_sample[j-pos-1]) 181 | pos_neg_dict['pos'] = pos_index 182 | 183 | potential_neg_list = [] 184 | for i in range(0, j-num_pos-negative_time_threshold): 185 | potential_neg_list.append(list_sample[i]) 186 | 187 | neg_index = sorted(random.sample(potential_neg_list, num_neg)) 188 | pos_neg_dict['neg'] = neg_index 189 | pos_neg_dict['loc'] = loc_scene 190 | all_query_dict[query_index] = pos_neg_dict 191 | if loc_scene == 'singapore-onenorth': 192 | son_query_dict[query_index] = pos_neg_dict 193 | elif loc_scene == 'singapore-hollandvillage': 194 | shv_query_dict[query_index] = pos_neg_dict 195 | elif loc_scene == 'singapore-queenstown': 196 | sq_query_dict[query_index] = pos_neg_dict 197 | elif loc_scene == 'boston-seaport': 198 | bs_query_dict[query_index] = pos_neg_dict 199 | 200 | with open(os.path.join(save_root, 'all_train_query_pos_neg_index_in_infos.pkl'), 'wb') as f: 201 | pickle.dump(all_query_dict, f) 202 | with open(os.path.join(save_root, 'son_train_query_pos_neg_index_in_infos.pkl'), 'wb') as f: 203 | pickle.dump(son_query_dict, f) 204 | with open(os.path.join(save_root, 'shv_train_query_pos_neg_index_in_infos.pkl'), 'wb') as f: 205 | pickle.dump(shv_query_dict, f) 206 | with open(os.path.join(save_root, 'sq_train_query_pos_neg_index_in_infos.pkl'), 'wb') as f: 207 | pickle.dump(sq_query_dict, f) 208 | with open(os.path.join(save_root, 'bs_train_query_pos_neg_index_in_infos.pkl'), 'wb') as f: 209 | pickle.dump(bs_query_dict, f) 210 | print('======================') 211 | 212 | # ====================generate test tuples==================== 213 | # We use all the samples in the old scenes as database 214 | # and use the samples in the new scenes as query 215 | print('splitting test data ...') 216 | 217 | son_test_list = [] 218 | shv_test_list = [] 219 | sq_test_list = [] 220 | bs_test_list = [] 221 | for i in new_scenes: 222 | scene = nusc_trainval.scene[i] 223 | loc_scene = nusc_trainval.get('log', scene['log_token'])['location'] 224 | if loc_scene == 'singapore-onenorth': 225 | son_test_list.append(i) 226 | elif loc_scene == 'singapore-hollandvillage': 227 | shv_test_list.append(i) 228 | elif loc_scene == 'singapore-queenstown': 229 | sq_test_list.append(i) 230 | elif loc_scene == 'boston-seaport': 231 | bs_test_list.append(i) 232 | 233 | son_database_list = [] 234 | shv_database_list = [] 235 | sq_database_list = [] 236 | bs_database_list = [] 237 | for i in old_scenes: 238 | scene = nusc_trainval.scene[i] 239 | loc_scene = nusc_trainval.get('log', scene['log_token'])['location'] 240 | if loc_scene == 'singapore-onenorth': 241 | son_database_list.append(i) 242 | elif loc_scene == 'singapore-hollandvillage': 243 | shv_database_list.append(i) 244 | elif loc_scene == 'singapore-queenstown': 245 | sq_database_list.append(i) 246 | elif loc_scene == 'boston-seaport': 247 | bs_database_list.append(i) 248 | 249 | # print("son") 250 | # print(len(son_test_list)) 251 | # print(len(son_database_list)) 252 | # print("shv") 253 | # print(len(shv_test_list)) 254 | # print(len(shv_database_list)) 255 | # print("sq") 256 | # print(len(sq_test_list)) 257 | # print(len(sq_database_list)) 258 | # print("bs") 259 | # print(len(bs_test_list)) 260 | # print(len(bs_database_list)) 261 | 262 | split_test_for_different_loc(infos, son_database_list, son_test_list, scene_nbr_list, save_root, 'son', positive_distance_threshold) 263 | split_test_for_different_loc(infos, shv_database_list, shv_test_list, scene_nbr_list, save_root, 'shv', positive_distance_threshold) 264 | split_test_for_different_loc(infos, sq_database_list, sq_test_list, scene_nbr_list, save_root, 'sq', positive_distance_threshold) 265 | split_test_for_different_loc(infos, bs_database_list, bs_test_list, scene_nbr_list, save_root, 'bs', positive_distance_threshold) 266 | 267 | print("done!") 268 | 269 | if __name__ == '__main__': 270 | # load config ================================================================ 271 | config_filename = '../configs/config_self_supervised_scheme.yml' 272 | config = yaml.safe_load(open(config_filename)) 273 | # ============================================================================ 274 | main(config) -------------------------------------------------------------------------------- /NUSC-PR/supervised/generate_basic_infos.py: -------------------------------------------------------------------------------- 1 | # Developed by Jingyi Xu, Junyi Ma, Zijie Zhou 2 | # Brief: extract basic infos from nuScenes datasets 3 | # NUSC-PR is proposed in the paper: Explicit Interaction for Fusion-Based Place Recognition 4 | 5 | import os 6 | import pickle 7 | import numpy as np 8 | from tqdm import tqdm 9 | from nuscenes.nuscenes import NuScenes 10 | import yaml 11 | 12 | 13 | def check_dir(*dirs): 14 | for dir in dirs: 15 | if not os.path.exists(dir): 16 | os.makedirs(dir) 17 | 18 | def gen_info(nusc, sample_tokens): 19 | # Collect sample infos into a dict 20 | cam_names = [ 21 | 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_RIGHT', 'CAM_BACK', 22 | 'CAM_BACK_LEFT', 'CAM_FRONT_LEFT'] 23 | lidar_names = ['LIDAR_TOP'] 24 | 25 | infos = list() 26 | for sample_token in tqdm(sample_tokens): 27 | sample = nusc.get('sample', sample_token) 28 | info = dict() 29 | cam_datas = list() 30 | lidar_datas = list() 31 | info['sample_token'] = sample_token 32 | info['timestamp'] = sample['timestamp'] 33 | info['scene_token'] = sample['scene_token'] 34 | 35 | cam_infos = dict() 36 | lidar_infos = dict() 37 | for cam_name in cam_names: 38 | cam_data = nusc.get('sample_data', sample['data'][cam_name]) 39 | cam_datas.append(cam_data) 40 | cam_info = dict() 41 | cam_info['sample_token'] = sample['data'][cam_name] 42 | cam_info['ego_pose'] = nusc.get('ego_pose', cam_data['ego_pose_token']) 43 | cam_info['timestample'] = cam_data['timestamp'] 44 | cam_info['filename'] = cam_data['filename'] 45 | cam_info['calibrated_sensor'] = nusc.get( 46 | 'calibrated_sensor', cam_data['calibrated_sensor_token']) 47 | cam_infos[cam_name] = cam_info 48 | for lidar_name in lidar_names: 49 | lidar_data = nusc.get('sample_data', 50 | sample['data'][lidar_name]) 51 | lidar_datas.append(lidar_data) 52 | lidar_info = dict() 53 | lidar_info['sample_token'] = sample['data'][lidar_name] 54 | lidar_info['ego_pose'] = nusc.get( 55 | 'ego_pose', lidar_data['ego_pose_token']) 56 | lidar_info['timestamp'] = lidar_data['timestamp'] 57 | lidar_info['filename'] = lidar_data['filename'] 58 | lidar_info['calibrated_sensor'] = nusc.get( 59 | 'calibrated_sensor', lidar_data['calibrated_sensor_token']) 60 | lidar_infos[lidar_name] = lidar_info 61 | info['cam_infos'] = cam_infos 62 | info['lidar_infos'] = lidar_infos 63 | scene = nusc.get('scene', sample['scene_token']) 64 | loc_scene = nusc.get('log', scene['log_token'])['location'] 65 | info['loc'] = loc_scene 66 | infos.append(info) 67 | 68 | return infos 69 | 70 | 71 | def get_location_sample_tokens(nusc, location): 72 | # Get the sample tokens of a specific location 73 | 74 | location_indices = get_location_indices(nusc, location) 75 | 76 | sample_token_list = [] 77 | 78 | for scene_index in location_indices: 79 | scene = nusc.scene[scene_index] 80 | sample_token = scene['first_sample_token'] 81 | 82 | while not sample_token == '': 83 | sample = nusc.get('sample', sample_token) 84 | sample_token_list.append(sample_token) 85 | sample_token = sample['next'] 86 | 87 | return sample_token_list 88 | 89 | 90 | def get_location_indices(nusc, location): 91 | # Get the indices of the specific location 92 | 93 | location_indices = [] 94 | for scene_index in range(len(nusc.scene)): 95 | scene = nusc.scene[scene_index] 96 | if nusc.get('log', scene['log_token'])['location'] != location: 97 | continue 98 | location_indices.append(scene_index) 99 | return np.array(location_indices) 100 | 101 | def get_sample_tokens(nusc): 102 | # Get all sample tokens 103 | 104 | sample_token_list = [] 105 | for scene_index in range(len(nusc.scene)): 106 | scene = nusc.scene[scene_index] 107 | sample_token = scene['first_sample_token'] 108 | 109 | while not sample_token == '': 110 | sample = nusc.get('sample', sample_token) 111 | sample_token_list.append(sample_token) 112 | sample_token = sample['next'] 113 | return sample_token_list 114 | 115 | def main(config): 116 | 117 | nusc_root = config['data_root']['nusc_root'] 118 | save_root = config['data_root']['save_root_basic_infos'] 119 | check_dir(save_root) 120 | 121 | nusc_trainval = NuScenes(version='v1.0-trainval', dataroot=nusc_root, verbose=True) 122 | 123 | locations = config['locations'] 124 | locations_ = config['locations_abbr'] 125 | 126 | # ====================generate location infos==================== 127 | for idx, location in enumerate(locations): 128 | sample_tokens = get_location_sample_tokens(nusc_trainval, location=location) 129 | location_infos = gen_info(nusc_trainval, sample_tokens) 130 | with open(os.path.join(save_root, 'nuscenes_infos-'+locations_[idx]+'.pkl'), 'wb') as f: 131 | pickle.dump(location_infos, f) 132 | 133 | # ====================generate all infos==================== 134 | sample_tokens_trainval = get_sample_tokens(nusc_trainval) 135 | infos = gen_info(nusc_trainval, sample_tokens_trainval) 136 | with open(os.path.join(save_root, 'nuscenes_infos.pkl'), 'wb') as f: 137 | pickle.dump(infos, f) 138 | 139 | print("done!") 140 | 141 | 142 | if __name__ == '__main__': 143 | # load config ================================================================ 144 | config_filename = '../configs/config_supervised_scheme.yml' 145 | config = yaml.safe_load(open(config_filename)) 146 | # ============================================================================ 147 | main(config) -------------------------------------------------------------------------------- /NUSC-PR/supervised/generate_selected_indicies.py: -------------------------------------------------------------------------------- 1 | # Developed by Jingyi Xu, Junyi Ma, Zijie Zhou 2 | # Brief: convert tokens to indices in basic infos of NUSC-PR for possible use 3 | # NUSC-PR is proposed in the paper: Explicit Interaction for Fusion-Based Place Recognition 4 | 5 | import pickle 6 | import numpy as np 7 | import os 8 | import yaml 9 | 10 | def check_dir(*dirs): 11 | for dir in dirs: 12 | if not os.path.exists(dir): 13 | os.makedirs(dir) 14 | 15 | def main(config): 16 | prev_root0 = config['data_root']['save_root_basic_infos'] 17 | with open(os.path.join(prev_root0, 'nuscenes_infos.pkl'), 'rb') as f: 18 | infos = pickle.load(f) 19 | save_root = config['data_root']['save_root_selected_indices'] 20 | check_dir(save_root) 21 | prev_root1 = config['data_root']['save_root_splitted_data'] 22 | prev_root2 = config['data_root']['save_root_select_by_dis'] 23 | locations = config['locations_abbr'] 24 | 25 | 26 | for location in locations: 27 | with open(os.path.join(prev_root2, location+'_train_query_pos_neg_tokens.pkl'), 'rb') as f: 28 | query_pos_neg_dict = pickle.load(f) 29 | with open(os.path.join(prev_root2, location+'_test_query_gt_tokens.pkl'), 'rb') as f: 30 | gt_mapping = pickle.load(f) 31 | 32 | # sample_token -> id 33 | infos_new = {} 34 | for i in range(len(infos)): 35 | sample_token = infos[i]['sample_token'] 36 | infos_new[sample_token] = int(i) 37 | 38 | ############################################################ 39 | print("generating database indices in " + location) 40 | db_index_in_infos = [] 41 | sample_tokens_db = np.load(os.path.join(prev_root1, location+'_db_sample_token.npy')) 42 | for i in range(len(sample_tokens_db)): 43 | db_token_ = sample_tokens_db[i] 44 | db_index = infos_new[db_token_] 45 | db_index_in_infos.append(db_index) 46 | 47 | np.save(os.path.join(save_root, location+'_db_index_in_infos.npy'), db_index_in_infos) 48 | 49 | ############################################################ 50 | print("generating test-query-gt mapping indices in " + location) 51 | test_dict = {} 52 | for key, value in gt_mapping.items(): 53 | query_sample_token = key 54 | test_query_index = infos_new[query_sample_token] 55 | 56 | gt_index_list = [] 57 | for j in range(len(value['gt'])): 58 | gt_sample_token = value['gt'][j] 59 | gt_idex = infos_new[gt_sample_token] 60 | gt_index_list.append(gt_idex) 61 | test_dict[test_query_index] = gt_index_list 62 | 63 | with open(os.path.join(save_root, location+'_test_query_gt_index_in_infos.pkl'), 'wb') as f: 64 | pickle.dump(test_dict, f) 65 | 66 | ############################################################ 67 | print("generating train-query-pos-neg mapping indices in " + location) 68 | train_dict = {} 69 | for key, value in query_pos_neg_dict.items(): 70 | query_sample_token = key 71 | train_query_index = infos_new[query_sample_token] 72 | 73 | pos_index_list = [] 74 | neg_index_list = [] 75 | pos_neg_dict = {} 76 | for j in range(len(value['pos'])): 77 | pos_sample_token = value['pos'][j] 78 | pos_idex = infos_new[pos_sample_token] 79 | pos_index_list.append(pos_idex) 80 | pos_neg_dict['pos'] = pos_index_list 81 | 82 | for j in range(len(value['neg'])): 83 | neg_sample_token = value['neg'][j] 84 | neg_idex = infos_new[neg_sample_token] 85 | neg_index_list.append(neg_idex) 86 | pos_neg_dict['neg'] = neg_index_list 87 | 88 | train_dict[train_query_index] = pos_neg_dict 89 | 90 | with open(os.path.join(save_root, location+'_train_query_pos_neg_index_in_infos.pkl'), 'wb') as f: 91 | pickle.dump(train_dict, f) 92 | 93 | print("=========================================") 94 | 95 | print("done!") 96 | 97 | if __name__ == '__main__': 98 | # load config ================================================================ 99 | config_filename = '../configs/config_supervised_scheme.yml' 100 | config = yaml.safe_load(open(config_filename)) 101 | # ============================================================================ 102 | main(config) -------------------------------------------------------------------------------- /NUSC-PR/supervised/select_pos_neg_samples_by_dis.py: -------------------------------------------------------------------------------- 1 | # Developed by Jingyi Xu, Junyi Ma, Zijie Zhou 2 | # Brief: select positve and negtive samples for queries in the supervised-learning scheme of NUSC-PR 3 | # NUSC-PR is proposed in the paper: Explicit Interaction for Fusion-Based Place Recognition 4 | 5 | import os 6 | # import h5py 7 | import pickle 8 | import numpy as np 9 | from sklearn.neighbors import NearestNeighbors 10 | import yaml 11 | 12 | def check_dir(*dirs): 13 | for dir in dirs: 14 | if not os.path.exists(dir): 15 | os.makedirs(dir) 16 | 17 | def main(config): 18 | 19 | prev_root = config['data_root']['save_root_splitted_data'] 20 | save_root = config['data_root']['save_root_select_by_dis'] 21 | check_dir(save_root) 22 | locations = config['locations_abbr'] 23 | 24 | for location in locations: 25 | 26 | poses_db = np.load(os.path.join(prev_root, location+'_db.npy')) 27 | sample_tokens = np.load(os.path.join(prev_root, location+'_sample_token.npy')) 28 | 29 | poses_test_query = np.load(os.path.join(prev_root, location+'_test_query.npy')) 30 | sample_tokens_db = np.load(os.path.join(prev_root, location+'_db_sample_token.npy')) 31 | tokens_test_query = np.load(os.path.join(prev_root, location+'_test_query_sample_token.npy')) 32 | 33 | poses_train_query = np.load(os.path.join(prev_root, location+'_train_query.npy')) 34 | tokens_train_query = np.load(os.path.join(prev_root, location+'_train_query_sample_token.npy')) 35 | 36 | ############################################################ 37 | print("generating gt samples for test queries in " + location) 38 | positive_distance_threshold = config['positive_distance_threshold'] 39 | 40 | knn = NearestNeighbors(n_jobs=-1) 41 | knn.fit(poses_db[:, 1:]) 42 | positives_for_test = list(knn.radius_neighbors(poses_test_query[:, 1:], radius=positive_distance_threshold, return_distance=False)) 43 | 44 | gt_mapping = {} 45 | for i, posi in enumerate(positives_for_test): 46 | positive_sample_tokens = [] 47 | 48 | for t in posi: 49 | positive_index_in_db = int(t) 50 | positive_pose = poses_db[positive_index_in_db] 51 | positive_sample_token = sample_tokens[int(positive_pose[0])] 52 | positive_sample_tokens.append(positive_sample_token) 53 | 54 | gt_dict = {} 55 | query_key = tokens_test_query[i] 56 | gt_dict['gt'] = positive_sample_tokens 57 | gt_mapping[query_key] = gt_dict 58 | 59 | with open(os.path.join(save_root, location+'_test_query_gt_tokens.pkl'), 'wb') as f: 60 | pickle.dump(gt_mapping, f) 61 | 62 | ############################################################ 63 | print("generating positive and negative samples for train queries in " + location) 64 | negative_distance_threshold = config['negative_distance_threshold'] 65 | nbr_positive_samples = config['nbr_positive_samples'] 66 | nbr_negative_samples = config['nbr_negative_samples'] 67 | 68 | knn = NearestNeighbors() 69 | knn.fit(poses_db[:, 1:]) 70 | positives_for_train = list(knn.radius_neighbors(poses_train_query[:, 1:], 71 | radius=positive_distance_threshold, 72 | return_distance=False)) 73 | query_ref_mapping = {} 74 | for i, posi in enumerate(positives_for_train): 75 | positive_sample_tokens = [] 76 | selected_positives = np.random.choice(positives_for_train[i], nbr_positive_samples) 77 | 78 | for t in selected_positives: 79 | positive_index_in_db = int(t) 80 | positive_pose = poses_db[positive_index_in_db] 81 | positive_sample_token = sample_tokens[int(positive_pose[0])] 82 | positive_sample_tokens.append(positive_sample_token) 83 | 84 | mapping_dict = {} 85 | query_key = tokens_train_query[i] 86 | mapping_dict['pos'] = positive_sample_tokens 87 | query_ref_mapping[query_key] = mapping_dict 88 | 89 | selected_no_negtives = list(knn.radius_neighbors(poses_train_query[:, 1:], radius=negative_distance_threshold, 90 | return_distance=False)) 91 | for i, posi in enumerate(selected_no_negtives): 92 | 93 | negative_sample_tokens = [] 94 | potential_negatives = np.setdiff1d(np.arange(poses_db.shape[0]), posi, 95 | assume_unique=True) 96 | selected_negtives = np.random.choice(potential_negatives, nbr_negative_samples) 97 | for t in selected_negtives: 98 | negtive_index_in_db = int(t) 99 | negtive_pose = poses_db[negtive_index_in_db] 100 | negtive_sample_token = sample_tokens[int(negtive_pose[0])] 101 | negative_sample_tokens.append(negtive_sample_token) 102 | 103 | query_key = tokens_train_query[i] 104 | mapping_dict = query_ref_mapping[query_key] 105 | mapping_dict['neg'] = negative_sample_tokens 106 | query_ref_mapping[query_key] = mapping_dict 107 | 108 | with open(os.path.join(save_root, location+'_train_query_pos_neg_tokens.pkl'), 'wb') as f: 109 | pickle.dump(query_ref_mapping, f) 110 | 111 | print("=========================================") 112 | 113 | print("done!") 114 | 115 | if __name__ == '__main__': 116 | # load config ================================================================ 117 | config_filename = '../configs/config_supervised_scheme.yml' 118 | config = yaml.safe_load(open(config_filename)) 119 | # ============================================================================ 120 | main(config) -------------------------------------------------------------------------------- /NUSC-PR/supervised/split_dataset.py: -------------------------------------------------------------------------------- 1 | # Developed by Jingyi Xu, Junyi Ma, Zijie Zhou 2 | # Brief: split query and database for NUSC-PR 3 | # NUSC-PR is proposed in the paper: Explicit Interaction for Fusion-Based Place Recognition 4 | 5 | import numpy as np 6 | from sklearn.neighbors import NearestNeighbors 7 | import os 8 | import random 9 | import pickle 10 | import yaml 11 | 12 | 13 | def check_dir(*dirs): 14 | for dir in dirs: 15 | if not os.path.exists(dir): 16 | os.makedirs(dir) 17 | 18 | def main(config): 19 | 20 | random.seed(1) 21 | prev_root = config['data_root']['save_root_basic_infos'] 22 | save_root = config['data_root']['save_root_splitted_data'] 23 | 24 | check_dir(save_root) 25 | locations = config['locations_abbr'] 26 | 27 | for location in locations: 28 | 29 | infos_path = os.path.join(prev_root, 'nuscenes_infos-'+location+'.pkl') 30 | 31 | with open(infos_path, 'rb') as f: 32 | infos = pickle.load(f) 33 | 34 | poses = [] 35 | timestamps = [] 36 | sample_tokens = [] 37 | sample_tokens_db = [] 38 | for i, info in enumerate(infos): 39 | pose = info['lidar_infos']['LIDAR_TOP']['ego_pose']['translation'] 40 | poses.append(pose[:2]) 41 | timestamp = info['timestamp'] 42 | timestamps.append(timestamp) 43 | sample_token = info['sample_token'] 44 | sample_tokens.append(sample_token) 45 | poses = np.array(poses, dtype=np.float32) 46 | timestamps = np.array(timestamps, dtype=np.float32).reshape(-1, 1) 47 | print('total frames for '+location+':', i) 48 | 49 | ############################################################ 50 | print('==> generating database for '+location) 51 | distance_interval = config['distance_interval'] 52 | 53 | poses = np.concatenate( 54 | (np.arange(len(poses), dtype=np.int32).reshape(-1, 1), np.array(poses)), 55 | axis=1).astype(np.float32) 56 | 57 | poses_db = poses[0, :].reshape(1, -1) 58 | sample_tokens_db.append(sample_tokens[0]) 59 | for i in range(1, poses.shape[0]): 60 | knn = NearestNeighbors(n_neighbors=1) 61 | knn.fit(poses_db[:, 1:3]) 62 | dis, index = knn.kneighbors(poses[i, 1:3].reshape(1, -1), 1, return_distance=True) 63 | if dis > distance_interval: 64 | poses_db = np.concatenate((poses_db, poses[i, :].reshape(1, -1)), axis=0) 65 | sample_tokens_db.append(sample_tokens[i]) 66 | 67 | print('number of database frames in '+location+': ', poses_db.shape[0]) 68 | print('finding corresponding tokens: ', len(sample_tokens_db)) 69 | 70 | ############################################################ 71 | print('==> generating query indices for '+location) 72 | timestamps = np.array(timestamps - min(timestamps)) / (3600 * 24 * 1e6) 73 | 74 | SEPERATE_TH = 200 75 | if location == 'bs': 76 | SEPERATE_TH = config['date_threshold_bs'] 77 | elif location == 'sq': 78 | SEPERATE_TH = config['date_threshold_sq'] 79 | elif location == 'son': 80 | SEPERATE_TH = config['date_threshold_son'] 81 | elif location == 'shv': 82 | SEPERATE_TH = config['date_threshold_shv'] 83 | 84 | train_indices, _ = np.where(timestamps < SEPERATE_TH) 85 | testval_indices, _ = np.where(timestamps >= SEPERATE_TH) 86 | 87 | database_indices = poses_db[:, 0].astype(int) 88 | train_query_indices = list(set(train_indices) - set(database_indices)) 89 | testval_query_indices = list(set(testval_indices) - set(database_indices)) 90 | val_ratio = config['val_ratio'] 91 | val_query_indices = random.sample(testval_query_indices, int(len(testval_query_indices) * val_ratio)) 92 | test_query_indices = list(set(testval_query_indices) - set(val_query_indices)) 93 | 94 | poses_train_query = poses[train_query_indices] 95 | poses_test_query = poses[test_query_indices] 96 | poses_val_query = poses[val_query_indices] 97 | 98 | tokens_train_query = np.array(sample_tokens)[train_query_indices] 99 | tokens_test_query = np.array(sample_tokens)[test_query_indices] 100 | tokens_val_query = np.array(sample_tokens)[val_query_indices] 101 | 102 | print('the number of train query frames in '+location+': ', tokens_train_query.shape[0]) 103 | print('the number of test query frames in '+location+': ', tokens_test_query.shape[0]) 104 | print('the number of val query frames in '+location+': ', tokens_val_query.shape[0]) 105 | 106 | ############################################################ 107 | print('===> saving database and queries for '+location) 108 | np.save(os.path.join(save_root, location+'_db.npy'), poses_db) 109 | np.save(os.path.join(save_root, location+'_train_query.npy'), poses_train_query) 110 | np.save(os.path.join(save_root, location+'_val_query.npy'), poses_val_query) 111 | np.save(os.path.join(save_root, location+'_test_query.npy'), poses_test_query) 112 | np.save(os.path.join(save_root, location+'_db_sample_token.npy'), sample_tokens_db) 113 | np.save(os.path.join(save_root, location+'_train_query_sample_token.npy'), tokens_train_query) 114 | np.save(os.path.join(save_root, location+'_val_query_sample_token.npy'), tokens_val_query) 115 | np.save(os.path.join(save_root, location+'_test_query_sample_token.npy'), tokens_test_query) 116 | np.save(os.path.join(save_root, location+'_sample_token.npy'), sample_tokens) 117 | print("=========================================") 118 | 119 | print("done!") 120 | 121 | if __name__ == '__main__': 122 | # load config ================================================================ 123 | config_filename = '../configs/config_supervised_scheme.yml' 124 | config = yaml.safe_load(open(config_filename)) 125 | # ============================================================================ 126 | main(config) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # EINet 2 | The official code and benchmark for our paper: [**Explicit Interaction for Fusion-Based Place Recognition**](https://arxiv.org/abs/2402.17264). 3 | 4 | This work has been accepted by IROS 2024 :tada: 5 | 6 | [Jingyi Xu](https://github.com/BIT-XJY), [Junyi Ma](https://github.com/BIT-MJY), [Qi Wu](https://github.com/Gatsby23), [Zijie Zhou](https://github.com/ZhouZijie77), [Yue Wang](https://scholar.google.com.hk/citations?hl=zh-CN&user=N543LSoAAAAJ), [Xieyuanli Chen](https://github.com/Chen-Xieyuanli), Wenxian Yu, Ling Pei*. 7 | 8 | ![image](https://github.com/BIT-XJY/EINet/assets/83287843/c560b31c-86d0-4ec8-a1b0-c6b8da79db0f) 9 | 10 | ## Installation 11 | 12 | We follow the installation instructions of our codebase [LCPR](https://github.com/ZhouZijie77/LCPR), which are also posted here. 13 | 14 | 15 | * Create a conda virtual environment and activate it 16 | ```bash 17 | git clone git@github.com:BIT-XJY/EINet.git 18 | cd EINet 19 | conda create -n EINet python=3.8 20 | conda activate EINet 21 | ``` 22 | * Install other dependencies 23 | ```bash 24 | pip install -r requirements.txt 25 | ``` 26 | 27 | ## Data Download 28 | - Please download the offical [nuScenes dataset](https://www.nuscenes.org/nuscenes#download), and link it to the data folder. 29 | - [nuscenes_occ_infos_train.pkl](https://github.com/JeffWang987/OpenOccupancy/releases/tag/train_pkl), and [nuscenes_occ_infos_val.pkl](https://github.com/JeffWang987/OpenOccupancy/releases/tag/val_pkl) are also provided by the previous work. We also use it to train the EINet. 30 | 31 | Note that the download data structure should be like: 32 | ``` 33 | nuscenes 34 | ├─ raw_data 35 | │ ├─ maps 36 | │ │ ├─ ... 37 | │ ├─ samples 38 | │ │ ├─ CAM_BACK 39 | │ │ ├─ CAM_BACK_LEFT 40 | │ │ ├─ CAM_BACK_RIGHT 41 | │ │ ├─ CAM_FRONT 42 | │ │ ├─ CAM_FRONT_LEFT 43 | │ │ ├─ CAM_FRONT_RIGHT 44 | │ │ ├─ LIDAR_TOP 45 | │ │ ├─ RADAR_BACK_LEFT 46 | │ │ ├─ RADAR_BACK_RIGHT 47 | │ │ ├─ RADAR_FRONT 48 | │ │ ├─ RADAR_FRONT_LEFT 49 | │ │ ├─ RADAR_FRONT_RIGHT 50 | │ ├─ sweeps 51 | │ │ ├─ CAM_BACK 52 | │ │ ├─ CAM_BACK_LEFT 53 | │ │ ├─ CAM_BACK_RIGHT 54 | │ │ ├─ CAM_FRONT 55 | │ │ ├─ CAM_FRONT_LEFT 56 | │ │ ├─ CAM_FRONT_RIGHT 57 | │ │ ├─ LIDAR_TOP 58 | │ │ ├─ RADAR_BACK_LEFT 59 | │ │ ├─ RADAR_BACK_RIGHT 60 | │ │ ├─ RADAR_FRONT 61 | │ │ ├─ RADAR_FRONT_LEFT 62 | │ │ ├─ RADAR_FRONT_RIGHT 63 | │ ├─ v1.0-test 64 | │ │ ├─ attribute.json 65 | │ │ ├─ calibrated_sensor.json 66 | │ │ ├─ ... 67 | │ ├─ v1.0-traninval 68 | │ │ ├─ attribute.json 69 | │ │ ├─ calibrated_sensor.json 70 | │ │ ├─ ... 71 | ``` 72 | 73 | ## NUSC-PR 74 | We propose the NUSC-PR benchmark to split nuScenes datasets with self-supervised and supervised learning schemes. 75 | 76 | ### Self-supervised Data Preparation 77 | - Extract basic information from nuScenes datasets, and split query and database for NUSC-PR. 78 | ```bash 79 | cd NUSC-PR 80 | cd self_supervised 81 | python generate_basic_infos.py 82 | python split_dataset.py 83 | cd .. 84 | ``` 85 | 86 | - The data structure with a self-supervised learning scheme should be like: 87 | ``` 88 | self_supervised_data 89 | ├─ generate_basic_infos 90 | │ ├─ nuscenes_infos-bs.pkl 91 | │ ├─ nuscenes_infos-shv.pkl 92 | │ ├─ nuscenes_infos-son.pkl 93 | │ ├─ nuscenes_infos-sq.pkl 94 | │ ├─ nuscenes_infos.pkl 95 | ├─ split_dataset 96 | │ ├─ all_train_query_pos_neg_index_in_infos.pkl 97 | │ ├─ bs_db_index_in_infos.npy 98 | │ ├─ bs_test_query_gt_index_in_infos.pkl 99 | │ ├─ bs_train_query_pos_neg_index_in_infos.pkl 100 | │ ├─ shv_db_index_in_infos.npy 101 | │ ├─ shv_test_query_gt_index_in_infos.pkl 102 | │ ├─ shv_train_query_pos_neg_index_in_infos.pkl 103 | │ ├─ son_db_index_in_infos.npy 104 | │ ├─ son_test_query_gt_index_in_infos.pkl 105 | │ ├─ son_train_query_pos_neg_index_in_infos.pkl 106 | │ ├─ sq_db_index_in_infos.npy 107 | │ ├─ sq_test_query_gt_index_in_infos.pkl 108 | │ ├─ sq_train_query_pos_neg_index_in_infos.pkl 109 | ``` 110 | 111 | 112 | ### Supervised Data Preparation 113 | - Extract basic information from nuScenes datasets, and split query and database for NUSC-PR. 114 | ```bash 115 | cd supervised 116 | python generate_basic_infos.py 117 | python split_dataset.py 118 | python select_pos_neg_samples_by_dis.py 119 | python generate_selected_indicies.py 120 | cd .. 121 | cd .. 122 | ``` 123 | 124 | - The data structure with a supervised learning scheme should be like: 125 | ``` 126 | supervised_data 127 | ├─ generate_basic_infos 128 | │ ├─ nuscenes_infos-bs.pkl 129 | │ ├─ nuscenes_infos-shv.pkl 130 | │ ├─ nuscenes_infos-son.pkl 131 | │ ├─ nuscenes_infos-sq.pkl 132 | │ ├─ nuscenes_infos.pkl 133 | ├─ generate_selected_indicies 134 | │ ├─ bs_db_index_in_infos.npy 135 | │ ├─ bs_test_query_gt_index_in_infos.pkl 136 | │ ├─ bs_train_query_pos_neg_index_in_infos.pkl 137 | │ ├─ shv_db_index_in_infos.npy 138 | │ ├─ shv_test_query_gt_index_in_infos.pkl 139 | │ ├─ shv_train_query_pos_neg_index_in_infos.pkl 140 | │ ├─ son_db_index_in_infos.npy 141 | │ ├─ son_test_query_gt_index_in_infos.pkl 142 | │ ├─ son_train_query_pos_neg_index_in_infos.pkl 143 | │ ├─ sq_db_index_in_infos.npy 144 | │ ├─ sq_test_query_gt_index_in_infos.pkl 145 | │ ├─ sq_train_query_pos_neg_index_in_infos.pkl 146 | ├─ select_pos_neg_samples_by_dis 147 | │ ├─ bs_test_query_gt_tokens.pkl 148 | │ ├─ bs_train_query_pos_neg_tokens.pkl 149 | │ ├─ shv_test_query_gt_tokens.pkl 150 | │ ├─ shv_train_query_pos_neg_tokens.pkl 151 | │ ├─ son_test_query_gt_tokens.pkl 152 | │ ├─ son_train_query_pos_neg_tokens.pkl 153 | │ ├─ sq_test_query_gt_tokens.pkl 154 | │ ├─ sq_train_query_pos_neg_tokens.pkl 155 | ├─ split_dataset 156 | │ ├─ bs_db_sample_token.npy 157 | │ ├─ bs_db.npy 158 | │ ├─ bs_sample_token.npy 159 | │ ├─ bs_test_query_sample_token.npy 160 | │ ├─ bs_test_query.npy 161 | │ ├─ bs_train_query_sample_token.npy 162 | │ ├─ bs_train_query.npy 163 | │ ├─ bs_val_query_sample_token.npy 164 | │ ├─ bs_val_query.npy 165 | │ ├─ shv_db_sample_token.npy 166 | │ ├─ shv_db.npy 167 | │ ├─ shv_sample_token.npy 168 | │ ├─ shv_test_query_sample_token.npy 169 | │ ├─ shv_test_query.npy 170 | │ ├─ shv_train_query_sample_token.npy 171 | │ ├─ shv_train_query.npy 172 | │ ├─ shv_val_query_sample_token.npy 173 | │ ├─ shv_val_query.npy 174 | │ ├─ son_db_sample_token.npy 175 | │ ├─ son_db.npy 176 | │ ├─ son_sample_token.npy 177 | │ ├─ son_test_query_sample_token.npy 178 | │ ├─ son_test_query.npy 179 | │ ├─ son_train_query_sample_token.npy 180 | │ ├─ son_train_query.npy 181 | │ ├─ son_val_query_sample_token.npy 182 | │ ├─ son_val_query.npy 183 | │ ├─ sq_db_sample_token.npy 184 | │ ├─ sq_db.npy 185 | │ ├─ sq_sample_token.npy 186 | │ ├─ sq_test_query_sample_token.npy 187 | │ ├─ sq_test_query.npy 188 | │ ├─ sq_train_query_sample_token.npy 189 | │ ├─ sq_train_query.npy 190 | │ ├─ sq_val_query_sample_token.npy 191 | │ ├─ sq_val_query.npy 192 | ``` 193 | 194 | ## TODO 195 | 196 | - [X] Release the [paper](https://arxiv.org/abs/2402.17264) 197 | - [X] Release the benchmark NUSC-PR code for EINet 198 | - [ ] Release the source code for EINet 199 | - [ ] Release our pretrained baseline model 200 | 201 | 202 | ## Acknowledgement 203 | 204 | We thank the fantastic works [LCPR](https://github.com/ZhouZijie77/LCPR), [ManyDepth](https://github.com/nianticlabs/manydepth.git), and [AutoPlace](https://github.com/ramdrop/autoplace.git) for their pioneer code release, which provide codebase for this work. 205 | -------------------------------------------------------------------------------- /map_pointcloud_to_image.py: -------------------------------------------------------------------------------- 1 | def map_pointcloud_to_image(self, 2 | pointsensor_token: str, 3 | camera_token: str, 4 | min_dist: float = 1.0, 5 | render_intensity: bool = False, 6 | show_lidarseg: bool = False, 7 | filter_lidarseg_labels: List = None, 8 | lidarseg_preds_bin_path: str = None, 9 | show_panoptic: bool = False) -> Tuple: 10 | """ 11 | Given a point sensor (lidar/radar) token and camera sample_data token, load pointcloud and map it to the image 12 | plane. 13 | :param pointsensor_token: Lidar/radar sample_data token. 14 | :param camera_token: Camera sample_data token. 15 | :param min_dist: Distance from the camera below which points are discarded. 16 | :param render_intensity: Whether to render lidar intensity instead of point depth. 17 | :param show_lidarseg: Whether to render lidar intensity instead of point depth. 18 | :param filter_lidarseg_labels: Only show lidar points which belong to the given list of classes. If None 19 | or the list is empty, all classes will be displayed. 20 | :param lidarseg_preds_bin_path: A path to the .bin file which contains the user's lidar segmentation 21 | predictions for the sample. 22 | :param show_panoptic: When set to True, the lidar data is colored with the panoptic labels. When set 23 | to False, the colors of the lidar data represent the distance from the center of the ego vehicle. 24 | If show_lidarseg is True, show_panoptic will be set to False. 25 | :return (pointcloud , coloring , image ). 26 | """ 27 | 28 | cam = self.nusc.get('sample_data', camera_token) 29 | pointsensor = self.nusc.get('sample_data', pointsensor_token) 30 | pcl_path = osp.join(self.nusc.dataroot, pointsensor['filename']) 31 | if pointsensor['sensor_modality'] == 'lidar': 32 | if show_lidarseg or show_panoptic: 33 | gt_from = 'lidarseg' if show_lidarseg else 'panoptic' 34 | assert hasattr(self.nusc, gt_from), f'Error: nuScenes-{gt_from} not installed!' 35 | 36 | # Ensure that lidar pointcloud is from a keyframe. 37 | assert pointsensor['is_key_frame'], \ 38 | 'Error: Only pointclouds which are keyframes have lidar segmentation labels. Rendering aborted.' 39 | 40 | assert not render_intensity, 'Error: Invalid options selected. You can only select either ' \ 41 | 'render_intensity or show_lidarseg, not both.' 42 | 43 | pc = LidarPointCloud.from_file(pcl_path) 44 | else: 45 | pc = RadarPointCloud.from_file(pcl_path) 46 | im = Image.open(osp.join(self.nusc.dataroot, cam['filename'])) 47 | 48 | # Points live in the point sensor frame. So they need to be transformed via global to the image plane. 49 | # First step: transform the pointcloud to the ego vehicle frame for the timestamp of the sweep. 50 | cs_record = self.nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token']) 51 | pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix) 52 | pc.translate(np.array(cs_record['translation'])) 53 | 54 | # Second step: transform from ego to the global frame. 55 | poserecord = self.nusc.get('ego_pose', pointsensor['ego_pose_token']) 56 | pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix) 57 | pc.translate(np.array(poserecord['translation'])) 58 | 59 | # Third step: transform from global into the ego vehicle frame for the timestamp of the image. 60 | poserecord = self.nusc.get('ego_pose', cam['ego_pose_token']) 61 | pc.translate(-np.array(poserecord['translation'])) 62 | pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix.T) 63 | 64 | # Fourth step: transform from ego into the camera. 65 | cs_record = self.nusc.get('calibrated_sensor', cam['calibrated_sensor_token']) 66 | pc.translate(-np.array(cs_record['translation'])) 67 | pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix.T) 68 | 69 | # Fifth step: actually take a "picture" of the point cloud. 70 | # Grab the depths (camera frame z axis points away from the camera). 71 | depths = pc.points[2, :] 72 | 73 | if render_intensity: 74 | assert pointsensor['sensor_modality'] == 'lidar', 'Error: Can only render intensity for lidar, ' \ 75 | 'not %s!' % pointsensor['sensor_modality'] 76 | # Retrieve the color from the intensities. 77 | # Performs arbitary scaling to achieve more visually pleasing results. 78 | intensities = pc.points[3, :] 79 | intensities = (intensities - np.min(intensities)) / (np.max(intensities) - np.min(intensities)) 80 | intensities = intensities ** 0.1 81 | intensities = np.maximum(0, intensities - 0.5) 82 | coloring = intensities 83 | elif show_lidarseg or show_panoptic: 84 | assert pointsensor['sensor_modality'] == 'lidar', 'Error: Can only render lidarseg labels for lidar, ' \ 85 | 'not %s!' % pointsensor['sensor_modality'] 86 | 87 | gt_from = 'lidarseg' if show_lidarseg else 'panoptic' 88 | semantic_table = getattr(self.nusc, gt_from) 89 | 90 | if lidarseg_preds_bin_path: 91 | sample_token = self.nusc.get('sample_data', pointsensor_token)['sample_token'] 92 | lidarseg_labels_filename = lidarseg_preds_bin_path 93 | assert os.path.exists(lidarseg_labels_filename), \ 94 | 'Error: Unable to find {} to load the predictions for sample token {} (lidar ' \ 95 | 'sample data token {}) from.'.format(lidarseg_labels_filename, sample_token, pointsensor_token) 96 | else: 97 | if len(semantic_table) > 0: # Ensure {lidarseg/panoptic}.json is not empty (e.g. in case of v1.0-test). 98 | lidarseg_labels_filename = osp.join(self.nusc.dataroot, 99 | self.nusc.get(gt_from, pointsensor_token)['filename']) 100 | else: 101 | lidarseg_labels_filename = None 102 | 103 | if lidarseg_labels_filename: 104 | # Paint each label in the pointcloud with a RGBA value. 105 | if show_lidarseg: 106 | coloring = paint_points_label(lidarseg_labels_filename, 107 | filter_lidarseg_labels, 108 | self.nusc.lidarseg_name2idx_mapping, 109 | self.nusc.colormap) 110 | else: 111 | coloring = paint_panop_points_label(lidarseg_labels_filename, 112 | filter_lidarseg_labels, 113 | self.nusc.lidarseg_name2idx_mapping, 114 | self.nusc.colormap) 115 | 116 | else: 117 | coloring = depths 118 | print(f'Warning: There are no lidarseg labels in {self.nusc.version}. Points will be colored according ' 119 | f'to distance from the ego vehicle instead.') 120 | else: 121 | # Retrieve the color from the depth. 122 | coloring = depths 123 | 124 | # Take the actual picture (matrix multiplication with camera-matrix + renormalization). 125 | points = view_points(pc.points[:3, :], np.array(cs_record['camera_intrinsic']), normalize=True) 126 | 127 | # Remove points that are either outside or behind the camera. Leave a margin of 1 pixel for aesthetic reasons. 128 | # Also make sure points are at least 1m in front of the camera to avoid seeing the lidar points on the camera 129 | # casing for non-keyframes which are slightly out of sync. 130 | mask = np.ones(depths.shape[0], dtype=bool) 131 | mask = np.logical_and(mask, depths > min_dist) 132 | mask = np.logical_and(mask, points[0, :] > 1) 133 | mask = np.logical_and(mask, points[0, :] < im.size[0] - 1) 134 | mask = np.logical_and(mask, points[1, :] > 1) 135 | mask = np.logical_and(mask, points[1, :] < im.size[1] - 1) 136 | coloring = coloring[mask] 137 | 138 | return points, coloring, im, depths -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | faiss_gpu==1.7.2 2 | h5py==3.10.0 3 | numpy==1.24.4 4 | nuscenes_devkit==1.1.11 5 | pillow==10.1.0 6 | PyYAML==6.0.1 7 | scikit_learn==1.3.2 8 | torch==1.11.0+cu113 9 | torchvision==0.12.0+cu113 10 | tqdm==4.66.1 11 | --------------------------------------------------------------------------------