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)