import gradio as gr
import time
import numpy as np
from PIL import Image, ImageOps
from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline

# Modelo seleccionado (EXISTENTE Y FUNCIONAL)
MODEL_ID = "bazyl/gtsrb-model"

# Cargar preprocesador y modelo
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
classifier = pipeline("image-classification",
                      model=model,
                      image_processor=processor,
                      top_k=3,
                      device=-1)

# -----------------------------
# CAPTURA DE IMAGEN DESDE WEBCAM
# -----------------------------
def capture_frame(frame):
    """Convierte la imagen capturada en PIL y corrige espejo."""
    if frame is None:
        return None

    # Convertir frame a PIL (cuando streaming=False llega como PIL correctamente)
    if isinstance(frame, np.ndarray):
        frame = Image.fromarray(frame)

    # Corregir efecto espejo típico de cámaras frontales
    frame = ImageOps.mirror(frame)
    return frame


# -----------------------------
# CLASIFICACIÓN
# -----------------------------
def predict(img):
    if img is None:
        return "⚠️ No image has been received.", "-"

    start = time.time()
    outputs = classifier(img)
    end = time.time()

    result_text = "\n".join(
        [f"{i+1}. {o['label']} ({o['score']*100:.1f}%)"
         for i, o in enumerate(outputs)]
    )

    return result_text, f"{end - start:.3f} s"


# -----------------------------
# INTERFAZ GRADIO
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## Traffic Signal Classifier (Camera + Upload images)")

    with gr.Row():
        # IMPORTANTE: Para que funcione en portátil, debe ser PIL + NO streaming
        webcam = gr.Image(
            sources=["webcam"],
            streaming=False,   # ✔ Captura funciona en PC y móvil
            type="pil",        # ✔ Los frames se reciben correctamente
            label="Camera (preview)"
        )

        captured = gr.Image(
            type="pil",
            label="Captured image",
            interactive=False
        )

    # Botón de captura
    capture_btn = gr.Button("📸 Capture image")
    capture_btn.click(fn=capture_frame, inputs=webcam, outputs=captured)

    # Subida de imagen alternativa
    upload = gr.Image(type="pil", sources=["upload"], label="Upload image from file")
    upload.change(lambda img: img, upload, captured)

    gr.Markdown("### Classification")

    pred_box = gr.Textbox(label="Prediction (top-3)", lines=5)
    time_box = gr.Textbox(label="Response time")

    classify_btn = gr.Button("Classify")
    classify_btn.click(fn=predict, inputs=captured, outputs=[pred_box, time_box])

demo.launch()