| import os | |
| import pickle | |
| from tqdm.auto import tqdm | |
| from copy import deepcopy | |
| if __name__ == "__main__": | |
| pred_path = \ | |
| "/home/jovyan/users/bulat/workspace/3drec/Indoor/OKNO/data/arkitscenes/points_vggt/" | |
| out_pkl_path = \ | |
| "arkit_vggt_val_subset.pkl" | |
| gt_pkl_path = \ | |
| "/home/jovyan/users/bulat/workspace/3drec/Indoor/OKNO/data/arkitscenes/arkitscenes_offline_infos_val.pkl" | |
| with open(gt_pkl_path, 'rb') as file: | |
| gt_data = pickle.load(file) | |
| new_data = {"metainfo": gt_data["metainfo"]} | |
| data_list = [] | |
| picked_scenes = [scene for scene in os.listdir(pred_path)] | |
| num = 0 | |
| for scene in tqdm(gt_data['data_list'][:10]): | |
| scene_name = scene['lidar_points']['lidar_path'] | |
| if scene_name not in picked_scenes: | |
| print(f"Scene {scene_name} not found in {pred_path}") | |
| continue | |
| num += 1 | |
| tmp_scene = deepcopy(scene) | |
| data_list.append(tmp_scene) | |
| print(f"Number of scenes: {num}") | |
| new_data['data_list'] = data_list | |
| with open(out_pkl_path, 'wb') as f: | |
| pickle.dump(new_data, f) | |