| 19 | |
| 20 | |
| 21 | class EvaluationPitts(): |
| 22 | def __init__(self, config, model, device, number_neighbors=20) -> None: |
| 23 | self.config = config |
| 24 | self.number_neighbors = number_neighbors |
| 25 | self.model = model |
| 26 | self.device = device |
| 27 | if config.MODEL.TYPE == "Lidar": |
| 28 | self.trans_lidar = Augment_Point_Data(is_train=False) |
| 29 | elif config.MODEL.TYPE == "LiSPH": |
| 30 | self.trans_lidar_sph = Augment_SPH_Data(img_size=config.DATA.SPH_PROJ.IMAGE_SIZE, |
| 31 | is_train=False) |
| 32 | self.projection = LaserProjection(device=device, |
| 33 | top_size=config.DATA.BEV_PROJ.IMAGE_SIZE, |
| 34 | z_range=config.DATA.BEV_PROJ.Z_RANGE, |
| 35 | sph_size=config.DATA.SPH_PROJ.IMAGE_SIZE, |
| 36 | fov_range=config.DATA.SPH_PROJ.FOV) |
| 37 | else: |
| 38 | raise NotImplementedError('Please Try Predefined Types') |
| 39 | |
| 40 | def get_features_recall(self, traj_num, trans_IDX, rot_IDX, rand_idx=0): |
| 41 | database = [] |
| 42 | query = [] |
| 43 | running_time = 0 |
| 44 | |
| 45 | # Set global map |
| 46 | if self.config.DATA.DATASET_NAME == 'PITT': |
| 47 | global_map = o3d.io.read_point_cloud('{}/{}/DATA/Train{}/cloudGlobal.pcd'.format( |
| 48 | self.config.DATA.BENCHMARK_DIR, self.config.DATA.DATASET_NAME,traj_num)) |
| 49 | bbox_pnv = o3d.geometry.AxisAlignedBoundingBox( |
| 50 | np.array([-20, -20, 0.8]), |
| 51 | np.array([ 20, 20, 100.0])) |
| 52 | map_tree = o3d.geometry.KDTreeFlann(global_map) |
| 53 | |
| 54 | # Set filelist for test |
| 55 | file_list = sorted(glob('{}/{}/train_{}/*_sph.png'.format(self.config.DATA.BENCHMARK_DIR, self.config.DATA.DATASET_NAME, traj_num))) |
| 56 | file_list = file_list[::10] |
| 57 | |
| 58 | # load database |
| 59 | for file_name in tqdm.tqdm(file_list): |
| 60 | if self.config.MODEL.TYPE == "LiSPH": |
| 61 | queries_img = Image.open(file_name) |
| 62 | frame = [self.trans_lidar_sph(queries_img).to(self.device, dtype=torch.float)] |
| 63 | embedding, t_gpu = self.model.infer_frame(frame, t=True) |
| 64 | database.append(embedding) |
| 65 | running_time += t_gpu |
| 66 | |
| 67 | # load the test query |
| 68 | for file_name in tqdm.tqdm(file_list): |
| 69 | # preprocess generate test query |
| 70 | pose = np.load('{}_pose6d.npy'.format( |
| 71 | file_name.split('_sph.png')[0])) |
| 72 | |
| 73 | # * Add fixed transformation |
| 74 | if rand_idx == 0: |
| 75 | trans_idx = trans_IDX + np.random.random(1)[0]-0.5 # 0.5m noise |
| 76 | rot_idx = rot_IDX + 5*(np.random.random(1)[0]-0.5) # 2.5° noise |
| 77 | # * Add random transformation |
| 78 | else: |