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

# ---------------------------------------------------------
# 1. Crea unos DataLoaders falsos para cargar el modelo (requerido)
# ---------------------------------------------------------
# Esto es solo para inicializar el modelo correctamente
dls = ImageDataLoaders.from_name_func(
    path='.', 
    fnames=get_image_files('.'), 
    label_func=lambda x: 'placeholder', 
    valid_pct=0.2,
    item_tfms=Resize(128),
    bs=1
)

# ---------------------------------------------------------
# 2. Crea el learner y carga el modelo desde .pth
# ---------------------------------------------------------
learn = cnn_learner(dls, resnet18, metrics=accuracy)
learn.load('resnet18_blindness') 

# Define tus clases manualmente si no están en dls.vocab
labels = ['No Blindness', 'Blindness']

# ---------------------------------------------------------
# 3. Define la función de predicción
# ---------------------------------------------------------
def predict(img):
    pred, idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# ---------------------------------------------------------
# 4. Lanza la app con Gradio
# ---------------------------------------------------------
gr.Interface(
    fn=predict,
    inputs=gr.Image(shape=(128, 128)),
    outputs=gr.Label(num_top_classes=2),
    title="Clasificador de Ceguera",
    description="Sube una imagen de retina y predice si hay ceguera o no."
).launch()