import gradio as gr
import cv2
import numpy as np
import pytesseract
from ultralytics import YOLO
import easyocr
import os
os.environ['TESSDATA_PREFIX'] = '/usr/share/tesseract-ocr/5/tessdata'

# Load YOLO model
# yolo_model = YOLO("models/NER_66_37_50_V2.pt")  # Your trained model Metz2_V2_Dataset_2_E_68_35best
yolo_model = YOLO("models/Metz2_V2_Dataset_2_E_68_35best.pt")  # Your trained model 

easyocr_readers = {}  # Cache for readers

# Store global crops
detected_crops = []

# --- Detection Function ---
def detect_objects(image):
    global detected_crops
    results = yolo_model(image)[0]
    annotated = image.copy()
    detected_crops = []

    for box in results.boxes.data.tolist():
        x1, y1, x2, y2, conf, cls = box
        x1, y1 = max(0, int(x1)), max(0, int(y1))
        x2, y2 = int(x2), int(y2)
        label = yolo_model.names[int(cls)]
        cropped = image[y1:y2, x1:x2]
        detected_crops.append((label, cropped))
        cv2.rectangle(annotated, (x1, y1), (x2, y2), (0,255,0), 2)
        cv2.putText(annotated, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255,0,0), 2)

    crops = [crop for _, crop in detected_crops]
    return annotated, crops

# --- OCR Function ---
def run_ocr(engine, lang = 'fr'):
    results = []
    for label, crop in detected_crops:
        try:
            if engine == "EasyOCR":
                if lang not in easyocr_readers:
                    easyocr_readers[lang] = easyocr.Reader([lang], gpu=False)
                reader = easyocr_readers[lang]
                text = "\n".join(reader.readtext(crop, detail=0))
            elif engine == "Tesseract":
                config = f'--oem 3 --psm 6 -l {lang}'
                text = pytesseract.image_to_string(crop, config=config)
            else:
                text = "[Unsupported Engine]"
        except Exception as e:
            text = f"[Error: {e}]"
        results.append(f"{label}:\n{text}")
    return "\n\n".join(results)

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 YOLO Detection + OCR Modular App - NER")

    with gr.Row():
        input_img = gr.Image(type="numpy", label="Upload Image")
        detected_img = gr.Image(type="numpy", label="Annotated Image")

    detect_button = gr.Button("🔍 Run Detection")
    gallery = gr.Gallery(label="Detected Crops")

    with gr.Row():
        ocr_engine = gr.Radio(["EasyOCR", "Tesseract"], value="EasyOCR", label="OCR Engine")
        lang_input = gr.Textbox("fr", label="Language Code")

    ocr_button = gr.Button("📝 Run OCR on Detected Regions")
    ocr_output = gr.Textbox(label="Extracted Text", lines=10)

    detect_button.click(fn=detect_objects, inputs=input_img, outputs=[detected_img, gallery])
    ocr_button.click(fn=run_ocr, inputs=[ocr_engine, lang_input], outputs=ocr_output)

    # Examples
    example_images = ["images/example1.jpg", "images/example2.jpg", "images/example3.jpg", "images/example4.jpg", "images/example5.jpg"]
    gr.Examples(
        examples=[[img] for img in example_images],
        inputs=[input_img],
        label="Example Images"
    )

demo.launch()