import os
import csv

f = open("./datasets/celeba/annotations/train.csv", "w", newline="")
wr = csv.writer(f)
wr.writerow(["_", "img_name", "mask_name"])

for i in range(23304):
    wr.writerow(
        [
            i,
            "celeba512_30k_masked/"
            + os.listdir("./datasets/celeba/images/celeba512_30k_masked")[i],
            "celeba512_30k_binary/"
            + os.listdir("./datasets/celeba/images/celeba512_30k_binary")[i],
        ]
    )

f.close()

f = open("./datasets/celeba/annotations/val.csv", "w", newline="")
wr = csv.writer(f)
wr.writerow(["_", "img_name", "mask_name"])

for i in range(23304, 29131):
    wr.writerow(
        [
            i,
            "celeba512_30k_masked/"
            + os.listdir("./datasets/celeba/images/celeba512_30k_masked")[i],
            "celeba512_30k_binary/"
            + os.listdir("./datasets/celeba/images/celeba512_30k_binary")[i],
        ]
    )

f.close()