# app.py
import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO
import easyocr

# Load YOLOv8 model
model_path = "YOLOv8n.pt"  # Local path in repo
yolo_model = YOLO(model_path)

# Initialize EasyOCR
easyocr_reader = easyocr.Reader(['en'], gpu=False)

def detect_and_read_plate(image_np):
    results = yolo_model.predict(source=image_np, conf=0.25)
    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
    
    annotated_image = image_np.copy()
    detected_texts = []

    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = box[:4]
        cropped_plate = image_np[y1:y2, x1:x2]
        text_results = easyocr_reader.readtext(cropped_plate)

        plate_texts = []
        for (bbox, text, conf) in text_results:
            if conf > 0.3:
                plate_texts.append(text)

        plate_text = ' '.join(plate_texts)
        detected_texts.append(f"Plate {i+1}: {plate_text}")

        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(annotated_image, plate_text, (x1, y1 - 10), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)

    return annotated_image, '\n'.join(detected_texts)

demo = gr.Interface(
    fn=detect_and_read_plate,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=[
        gr.Image(type="numpy", label="Detected Image with Plate Text"),
        gr.Textbox(label="Detected Plate Texts")
    ],
    title="Number Plate Detection and OCR",
    description="YOLOv8 detects license plates and EasyOCR reads the number."
)

demo.launch()