data_formater.py

# Convert the data to dataloader formater.
import os


# Check root directory exists,
# If not, create it.
if not os.path.exists("dataset/root"):
    os.makedirs("dataset/root")
    




# Check if the labels.csv file exists, if it does, delete it.
if os.path.exists("dataset/root/labels.csv"):
    os.remove("dataset/root/labels.csv")

# Create a labels csv file.
print("Creating labels.csv file.")
classes_to_model_output = {"left": 0, "right": 1}
with open("dataset/root/labels.csv", "w") as file:
    # file.write("image,class\n")
    classes = ["left", "right"]
    for class_name in classes:
        image_files = os.listdir(os.path.join("dataset", class_name))
        for image in image_files:
            file.write(f"{image},{classes_to_model_output[class_name]}\n")


print("Creating uniform image dataset.")
# Create a uniform image dataset, named train


if not os.path.exists("dataset/root/train"):
    os.makedirs("dataset/root/train")

# Copy the images to the root directory.
for class_name in classes:
    image_files = os.listdir(os.path.join("dataset", class_name))
    for image in image_files:
        os.system(f"cp dataset/{class_name}/{image} dataset/root/train/{image}")