import argparse
import os
import shutil

parser = argparse.ArgumentParser()
parser.add_argument("-d", "--dir", help="dir with the images", required=True)
parser.add_argument("-l", "--labels", help="file with image name to class label mapping", required=True)

args = parser.parse_args()

processed_classes = set()

with open(args.labels, "r") as file:
    # skip header
    next(file)
    for line in file:
        img_name, labels = line.split(",")
        class_name = labels.split(" ")[0]
        # create a dir for this classname
        if class_name not in processed_classes:
            dir_path = args.dir + "/" + class_name
            if not os.path.exists(dir_path):
                os.mkdir(dir_path)
        shutil.move(args.dir + "/" + img_name + ".JPEG", args.dir + "/" + class_name+ "/" + img_name + ".JPEG")