import os
import shutil
import re

from yolov5 import detect
import easyocr

def main():
    directory = "plates"
    detect.run(
        source=directory,
        save_crop=True,
        weights="yolo-licence-plate.pt",
        conf_thres=0.25,
        iou_thres=0.45,
        agnostic_nms=False,
        max_det=1,
        project="predictions",
        exist_ok=True
    )

    reader = easyocr.Reader(["es"])
    crops_directory = "predictions/exp/crops/license_plate"
    for image in os.listdir(crops_directory):
        filepath = os.path.join(crops_directory, image)
        if os.path.isfile(filepath) and image.endswith((".jpg", ".jpeg", ".png")):
            image_path = os.path.join(crops_directory, image)
            print(f"Processing {image_path}...")
            result = reader.readtext(image_path, detail=0)
            clean_plate = re.sub(r"[^A-Z0-9]", "", result[0].upper())
            print(clean_plate)

    shutil.rmtree("predictions/exp")
if __name__ == "__main__":
    main()