File size: 895 Bytes
663494c |
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 |
import pickle
import os
split = 'trainval'
data_dir = 'data/openscene-v1.1'
if __name__ == '__main__':
mini_infos = os.listdir(f'{data_dir}/meta_datas/{split}')
mini_infos = [os.path.join(f"{data_dir}/meta_datas/{split}", each) for each in mini_infos if each.endswith('.pkl')]
train_paths = mini_infos[:int(len(mini_infos) * 0.85)]
val_paths = mini_infos[int(len(mini_infos) * 0.85):]
train_infos = []
for file in train_paths:
with open(file, 'rb') as f:
train_infos.extend(pickle.load(f))
val_infos = []
for file in val_paths:
with open(file, 'rb') as f:
val_infos.extend(pickle.load(f))
with open(f'{data_dir}/infos/openscene_{split}_train.pkl', 'wb') as f:
pickle.dump(train_infos, f)
with open(f'{data_dir}/infos/openscene_{split}_val.pkl', 'wb') as f:
pickle.dump(val_infos, f) |