import os
os.environ["YOLO_CONFIG_DIR"] = "/tmp"

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

# -------------------------------------------------------
# 🔥 Charger le modèle YOLO (détection plaques)
# -------------------------------------------------------
model = YOLO("best1.pt")   # Mets ton modèle ici

# -------------------------------------------------------
# 🔥 Charger OCR arabe + anglais
# -------------------------------------------------------
reader = easyocr.Reader(['ar', 'en'], gpu=False)

# -------------------------------------------------------
# 🔥 Fonction de détection + OCR
# -------------------------------------------------------
def recognize_license_plate(image):

    # Convertir PIL → OpenCV (BGR)
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Détection YOLO
    results = model(image_rgb, conf=0.5, verbose=False)

    output = image_rgb.copy()
    detections = []

    for r in results:
        if not hasattr(r, "boxes") or r.boxes is None:
            continue

        for box in r.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cls_id = int(box.cls[0])
            class_name = model.names.get(cls_id, "unknown")
            conf = float(box.conf[0])

            # Dessiner la boîte
            cv2.rectangle(output, (x1, y1), (x2, y2), (0, 255, 0), 2)

            # Crop de la plaque
            crop = image_rgb[y1:y2, x1:x2]

            # Vérifier si crop valide
            if crop.size == 0:
                continue

            # Préprocessing OCR
            gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
            gray = cv2.resize(gray, None, fx=2, fy=2, interpolation=cv2.INTER_LINEAR)
            gray = cv2.GaussianBlur(gray, (3, 3), 0)

            # OCR
            ocr_result = reader.readtext(gray)

            if len(ocr_result) > 0:
                text = ocr_result[0][1]
                text_conf = float(ocr_result[0][2])
            else:
                text = ""
                text_conf = 0.0

            # Ajouter le texte sur l'image
            cv2.putText(output, text, (x1, y1 - 8),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)

            # Ajouter à la liste JSON
            detections.append({
                "country": class_name,
                "bbox": [x1, y1, x2, y2],
                "plate_text": text,
                "plate_confidence": round(text_conf, 2),
                "detection_confidence": round(conf, 2)
            })

    return output, detections


# -------------------------------------------------------
# 🔥 Interface Gradio
# -------------------------------------------------------
app = gr.Interface(
    fn=recognize_license_plate,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(label="Image + OCR"),
        gr.JSON(label="Detections JSON")
    ],
    title="YOLO – Detection + EasyOCR (Arabic + English)"
)

app.launch()