File size: 1,122 Bytes
55e58d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
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)
|