from fastai.vision.all import *
from fastai.vision.core import PILImage
import torch
import gradio as gr

import warnings
warnings.filterwarnings("ignore")


learn = load_learner('modelo.pkl')

classes = ['abdomen', 'antebrazo', 'cadera', 'cervical', 'clavicula', 'codo', 'col. torax', 'craneo', 'dedos', 'hombro',
           'lumbar', 'mano', 'muslo', 'muñeca', 'otros', 'pelvis', 'pierna', 'pies', 'rodilla', 'senos nasales', 'tobillo', 'torax']

# Función de predicción
def classify_image(image, model=learn, classes=classes):

    # Cargar imagen y realizar predicción
    img = PILImage.create(image)
    pred, pred_idx, probs = model.predict(img)

    # Filtrar probabilidades
    probs = torch.where(probs > 1e-2, probs, torch.tensor(0).to(probs.device))

    # Obtener top 5 resultados
    top5_probs, top5_idxs = torch.topk(probs, 5)
    top5_classes = [classes[idx] for idx in top5_idxs]

    # Crear lista de predicciones
    predictions = []
    for i in range(5):
        if top5_probs[i] > 1e-2:
            prob = round(float(top5_probs[i].numpy()), 3)
            predictions.append(f"{top5_classes[i]}: {prob}")

    return predictions


inputs = gr.inputs.Image()
outputs = gr.outputs.Textbox()

gr.Interface(fn=classify_image, inputs=inputs, outputs=outputs, title='Clasificación de Imágenes Médicas',
             description='Cargue una radiografía').launch()