app.py

import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification, pipeline

MODEL_ID = "xpatox08/gtsrb-traffic-signs-v1"

# Cargar modelo y procesador
image_processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)

classifier = pipeline(
    "image-classification",
    model=model,
    image_processor=image_processor,
    device=-1
)

def predict(img):
    # Ejecutar modelo
    outputs = classifier(img)

    if not outputs:
        return "Señal desconocida"

    # Tomar la predicción con mayor score
    top = max(outputs, key=lambda x: x["score"])

    # El modelo YA devuelve el nombre real de la señal
    label = top["label"]

    return label  # ← devolvemos directamente el texto

title = "Clasificador de Señales de Tráfico (GTSRB)"
description = "Sube una imagen de una señal y el modelo la identificará."

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Imagen de la señal"),
    outputs=gr.Textbox(label="Predicción de la señal"),
    title=title,
    description=description,
)

if __name__ == "__main__":
    demo.launch()