from pathlib import Path

import gradio as gr
from huggingface_hub import from_pretrained_fastai
from fastai.vision.all import *
from fastai.vision.widgets import *
from fastdownload import download_url


# Cargar modelo de clasificación de imágenes desde Hugging Face
learn = from_pretrained_fastai("AlejandroOSM1/Clasificador_de_camiones_Definitivo_Resnet18")

# Obtener etiquetas del modelo
labels = learn.dls.vocab

def predict(img):
    img = PILImage.create(img)  # Convertir imagen a formato FastAI
    pred, pred_idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Crear la interfaz con Gradio
interface = gr.Interface(
    fn=predict,  # Se usaba "classify_image", pero la función definida es "predict"
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),  # Mostrar el top 3 de clases
    title="Clasificador de Camiones",
    description="Sube una imagen y el modelo la clasificará en diferentes tipos de camiones."
)

# Lanzar la aplicación
if __name__ == "__main__":
    interface.launch()