import gradio as gr
from ultralytics import YOLO
import cv2
import easyocr
import numpy as np
from PIL import Image

# Load YOLO model
model = YOLO('yolo11n-custom.pt')
model.fuse()

# Load EasyOCR
reader = easyocr.Reader(['en'])

def detect_license_plate(image):
    results = model.predict(image, conf=0.15, iou=0.3, classes=[0])
    plate_texts = []
    img_array = np.array(image)
    img = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    img_height, img_width, _ = img.shape

    for result in results:
        for bbox in result.boxes.xyxy:
            x1, y1, x2, y2 = map(int, bbox.tolist())
            plate = img[int(y1):int(y2), int(x1):int(x2)]
            scale = 2
            height, width = plate.shape[:2]
            plate = cv2.resize(plate, (width * scale, height * scale), interpolation=cv2.INTER_CUBIC)
            lab = cv2.cvtColor(plate, cv2.COLOR_RGB2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
            l = clahe.apply(l)
            plate = cv2.merge((l, a, b))
            plate = cv2.cvtColor(plate, cv2.COLOR_LAB2RGB)

            text = reader.readtext(plate, detail=0, allowlist="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-")
            text = " ".join(text).upper()

            text_scale = max(1, width / 250)
            thickness = max(2, width // 200)
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), thickness)
            (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, text_scale, thickness)
            text_x = x1 + (width - text_width) // 2
            text_y = y1 - 10 if y1 > 50 else y2 + text_height + 20
            text_box_y1 = text_y - text_height - 5
            text_box_y2 = text_y + 5
            cv2.rectangle(img, (text_x - 8, text_box_y1 - 3), (text_x + text_width + 8, text_box_y2 + 3), (0, 0, 0), -1)
            cv2.rectangle(img, (text_x - 5, text_box_y1), (text_x + text_width + 5, text_box_y2), (255, 255, 255), -1)
            cv2.putText(img, text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, text_scale, (0, 0, 0), thickness)

            plate_texts.append(text)

    result_img = Image.fromarray(img)
    return result_img, "\n".join(plate_texts) if plate_texts else "No license plates detected."

# Gradio UI
iface = gr.Interface(
    fn=detect_license_plate,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Detected Image"),
        gr.Textbox(label="Detected License Plates")
    ],
    title="🚘 License Plate Detector",
    description="Upload an image with license plates to detect them using YOLO and EasyOCR."
)

iface.launch()