import os
from huggingface_hub import InferenceClient
import gradio as gr
from dotenv import load_dotenv


load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

if HF_TOKEN is None:
    raise ValueError("No se encontró la variable HF_TOKEN en el .env")


client = InferenceClient(provider="hf-inference", api_key=HF_TOKEN)

def clasificar_imagen(imagen):
    if imagen is None:
        return "Sube una imagen para clasificar."

    temp_path = "temp_image.jpg"
    imagen.save(temp_path)

    output = client.image_classification(temp_path, model="google/vit-base-patch16-224")

    if isinstance(output, list) and len(output) > 0:
        resultado = "\n".join([f"{item['label']}: {item['score']:.2f}" for item in output])
    else:
        resultado = "No se pudo clasificar la imagen."
    return resultado

with gr.Blocks() as demo:
    gr.Markdown("# Clasificador de Imágenes con ViT")
    gr.Markdown("Sube una imagen y se clasificará.")

    img_input = gr.Image(type="pil", label="Sube tu imagen aquí")
    boton = gr.Button("Clasificar Imagen")
    salida = gr.Textbox(label="Resultados")

    boton.click(fn=clasificar_imagen, inputs=img_input, outputs=salida)

demo.launch()