import gradio as gr
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from PIL import Image
import requests  

# Cargar el modelo y el extractor de características
model_name = "microsoft/swin-small-patch4-window7-224"
model = AutoModelForImageClassification.from_pretrained(model_name)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)

def predict(image):
    # Preprocesar la imagen
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    # Obtener las predicciones
    probs = torch.nn.functional.softmax(logits, dim=-1)
    top_probs, top_labels = torch.topk(probs, 3)
    top_probs = top_probs.detach().numpy().flatten()
    top_labels = top_labels.detach().numpy().flatten()
    # Convertir las etiquetas a nombres
    id2label = model.config.id2label
    labels = [id2label[label] for label in top_labels]
    return {labels[i]: float(top_probs[i]) for i in range(len(labels))}

titulo = "Mi primer demo con Hugging Face"
descripcion = "Este es un demo ejecutado durante la clase de Hugo Martinez."

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Carga una imagen aquí"),
    outputs=gr.Label(num_top_classes=3),
    title=titulo,
    description=descripcion
)

demo.launch()