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

import pathlib

# Enlace HTML con encabezado h1
html_link = '<h1><a href="" target="_blank">Link a la V2 (seleccionando presas)</a></h1>'

# Cargar el modelo preentrenado
repo_id = "ignaciobfp/moonboard_difficulty"
learner = from_pretrained_fastai(repo_id)
labels = learner.dls.vocab

# Función para realizar predicciones
def predict(img):
    #img = PILImage.create(img)
    pred,pred_idx,probs = learner.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

image_dir = pathlib.Path('images')
examples = [[path.as_posix()] for path in sorted(image_dir.glob('*.png'))]

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(height=144, width=144, label='Input', type='numpy'),
    outputs=gr.Label(label='Output', num_top_classes=3),
    title="Análisis de dificultad de imágenes de bloques Moonboard",
    description="[Versión mejorada del modelo, seleccionando las presas](https://huggingface.co/spaces/Ignaciobfp/moonboard-2)",
    examples=examples,
)

# Agregar el enlace al principio
iface.launch(share=False)