import os
import numpy as np
import pandas as pd

data_dir = 'data/cinic-10'

files = []
for root, _, filenames in os.walk(data_dir):
    for filename in filenames:
        if filename.lower().endswith(".png"):
            files.append(os.path.join(root, filename))

df = pd.DataFrame(files, columns=['path'])

df['id'] = np.arange(len(df))
df['class_name'] = df['path'].apply(lambda x: x.split('/')[-2])

all_classes = np.unique(df['class_name'])
df['label'] = df['class_name'].map({class_name: i for i, class_name in enumerate(all_classes)})

df['cifar_train'] = df['path'].str.contains("cifar10-train", na=False)
df['cifar_test'] = df['path'].str.contains("cifar10-test", na=False)
df['imagenet'] = ~(df['cifar_train'] | df['cifar_test'])

df.to_csv(os.path.join(data_dir, 'metadata.csv'), index=False)


