from ultralytics import YOLO
from PIL import Image
import numpy as np
import easyocr

car_model = YOLO("yolov8n.pt")
plate_model = YOLO("license_plate_detector.pt")
reader = easyocr.Reader(["nl"])

img = Image.open("test.jpg")

results = car_model.predict(source=img)

for r in results:
    for box in r.boxes:
        cls_id = int(box.cls[0])
        cls_name = r.names[cls_id]
        if cls_name == "car":
            conf = float(box.conf[0])
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            print(f"{cls_name} ({conf:.2f}) at [{x1},{y1},{x2},{y2}]")

            # Crop Image
            cropped_img = img.crop((x1, y1, x2, y2))
            cropped_img.save("car_crop_pillow.jpg")
            results2 = plate_model.predict(source=cropped_img)
            for x in results2:
                for box2 in x.boxes:
                    cls_id = int(box2.cls[0])
                    cls_name = x.names[cls_id]
                    if cls_name == "License_Plate":
                        conf = float(box2.conf[0])
                        x1, y1, x2, y2 = map(int, box2.xyxy[0])
                        print(
                            f"{cls_name} ({conf:.2f}) at [{x1},{y1},{x2},{y2}]"
                        )
                        lp_img = cropped_img.crop((x1, y1, x2, y2))
                        lp_img.save("license_plate.jpg")
                        lp_np = np.array(lp_img)
                        result3 = reader.readtext(lp_np)
                        print(result3)