from os import listdir
from os.path import isfile, join
import numpy as np
from shutil import copyfile
import csv

in_folder = "resized_images/"

train_folder = "resized_images_train/"
test_folder = "resized_images_test/"

onlyfiles = [f for f in listdir(in_folder) if isfile(join(in_folder, f))]

train_labels = []
test_labels = []
possible_labels = [0, 1]
for i, img_file in enumerate(onlyfiles):
    img_label = np.random.choice(possible_labels)
    if np.random.rand() >= 0.85:
        test_labels.append((img_file, str(img_label)))
        copyfile(in_folder + img_file, test_folder + img_file)
    else:
        train_labels.append((img_file, str(img_label)))
        copyfile(in_folder + img_file, train_folder + img_file)

def write_data_to_csv(data_list, f_name):
    with open(f_name, 'w', newline='') as csvfile:
        fieldnames = ["img_id", "label"]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        for img_id, label in data_list:
            writer.writerow({"img_id": img_id, "label": label})

write_data_to_csv(train_labels, "train.csv")
write_data_to_csv(test_labels, "test.csv")
