import gradio as gr
from fastai.vision.all import *
from sklearn.metrics import roc_auc_score

# Model functions
def get_x(row): return Path(str(path/f"{row['rootname']}_small"/f"{row['ID']}") + ".png")    
def get_y(row): return row["LABEL"]
def auroc_score(input, target):
    input, target = input.cpu().numpy()[:,1], target.cpu().numpy()
    return roc_auc_score(target, input)

# Load model
learn = load_learner("export.pkl")   

# Labels   
labels = ["Negative", "Positive"]

# Prediction function
def predict(img):
    img = PILImage.create(img)
    pred, idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}



# Interface parameters
title = "Ethiopia TB Detection"
description = "Detect TB from chest x-rays"
examples = ['patient1.png', 'patient2.png', 'patient3.png']

# Launch interface
gr.Interface(fn=predict,
             inputs=gr.inputs.Image(shape=(512, 512)), 
             outputs=gr.outputs.Label(num_top_classes=1),
             title=title,
             description=description,
             examples=examples,
             css=custom_css).launch(inline=False)