make_index.py

import argparse
import os
import json
import multiprocessing
from glob import glob
from os.path import join
from tqdm.contrib.concurrent import process_map

parser = argparse.ArgumentParser()
parser.add_argument('--data_path', default='')
args = parser.parse_args()

dirs = [join(args.data_path, f) for f in os.listdir(args.data_path) if os.path.isdir(join(args.data_path, f))]
def fun(dir):
    files = [join(dir, f) for f in os.listdir(dir) if os.path.isfile(join(dir, f))]
    return files

#with multiprocessing.Pool(processes=os.cpu_count()) as pool:
results = process_map(fun, dirs, max_workers=os.cpu_count())

new_results = []
for d in results:
    new_results += d
    
with open("./index.json", 'w') as f:
    json.dump(new_results, f, indent=4)