File size: 749 Bytes
da2e2ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import os
import pickle
traj_root = os.getenv('NAVSIM_TRAJPDM_ROOT')
if __name__ == '__main__':
out_dir = 'vocab_expanded_8192_navtest'
os.makedirs(f'{traj_root}/{out_dir}', exist_ok=True)
ins = [f'navtest_sub{i}.pkl' for i in range(1, 6)]
out = 'navtest.pkl'
result = {}
for in_pkl in ins:
postfix = in_pkl.split('.')[0]
sub = postfix.split('_')[1]
curr_pickle = pickle.load(open(f'{traj_root}/{out_dir}_{sub}/{in_pkl}', 'rb'))
print(f'{traj_root}/{out_dir}_{sub}/{in_pkl}', len(curr_pickle))
for k, v in curr_pickle.items():
result[k] = v
print(f'Length: {len(result)}')
pickle.dump(result, open(f'{traj_root}/{out_dir}/{out}', 'wb'))
|