#%%


from pathlib import Path
import pickle 

import tensorflow_datasets as tfds

for split in ['valid', 'test', 'train']:
    dataset = tfds.load('byol_celeb_a2',
                    split=split,
                    shuffle_files=False,
                    data_dir='data',
                    #    read_config=tfds.ReadConfig(
                    #        shuffle_seed=shuffle_seed,
                    #        shuffle_reshuffle_each_iteration=reshuffle_iteration)
                    )

#     print(dataset)
#     ds_np = tfds.as_numpy(dataset)

#     print(ds_np)
#     print(next(iter(ds_np)))

#     path = 'data/byol_imagenet2012_single'
#     path = Path(path) / split
#     path.mkdir(parents=True, exist_ok=True)

#     for item in ds_np:
#         file_name = item['file_name'].split(b'.')[0]
#         with open(path / (file_name.decode("utf-8") + '.pkl'), 'wb') as f:
#             pickle.dump(item, f, pickle.HIGHEST_PROTOCOL)



# # %%

# %%
