robotinder-data/data_parser.py

import json

from os import listdir
from os.path import isfile, join, isdir
mypath = './'
onlyfolders = [f for f in listdir(mypath) if isdir(join(mypath, f))]

idx = 0
data_info = {'idx': [], 'env': [], 'name': []}
for folder in onlyfolders:
    print(folder)
    fs = [join(mypath, folder, f) for f in listdir(join(mypath, folder)) if isfile(join(mypath, folder, f))]
    # print(fs)
    for f in fs:
        if f.endswith(".gif"):
            idx += 1
            data_info['idx'].append(idx)
            data_info['env'].append(folder)
            data_info['name'].append(f.split('/')[-1])

with open('data_info.json', 'w') as f:
    json.dump(data_info, f)
    
with open('data_info.json', 'r') as f:
    data_info = json.load(f)
print(data_info)