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

learn = load_learner('export.pkl')

labels = learn.dls.vocab

def prever(imagem):
    img = PILImage.create(imagem)
    predicao, indice_predicao, probabilidades = learn.predict(img)
    predicao = str(predicao)
    probabilidades = [f"{labels[i]}: {prob * 100:.2f}%" for i, prob in enumerate(probabilidades)]
    
    return predicao, probabilidades

titulo = "Detecção de Câncer de Mama com Aprendizado Profundo e Transferência de Aprendizado (ResNet18)"
descricao = "<p style='text-align: center'><b>Como radiologista ou oncologista, é crucial saber o que há de errado com uma imagem de raio-X de mama.<b><br><b> Faça o upload da imagem de raio-X de mama para saber o que há de errado com o seio de uma paciente com ou sem implante<b><p>"
artigo="<p style='text-align: center'>Aplicativo web desenvolvido e gerenciado por Addai Fosberg<b></p>"
exemplos = [['img1.jpeg'], ['img2.jpeg']]
habilitar_fila = True

inputs = gr.inputs.Image(shape=(512, 512), label="Imagem de Raio-X de Mama")
outputs = [gr.outputs.Label(label="Diagnóstico"), gr.outputs.Label(label="Probabilidades")]

interface = gr.Interface(
    fn=prever,
    inputs=inputs,
    outputs=outputs,
    title=titulo,
    description=descricao,
    article=artigo,
    examples=exemplos,
    enable_queue=habilitar_fila
)

interface.launch()