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

learn = load_learner("brain_tumor_prediction_model.pkl")


categories = ("No tumor detected", "A tumor was detected")


def classify_image(img):
    """
    Classify an brain mri as a tumor or not a tumor

    Parameters
    ----------
    img : PIL image
        The image to classify

    Returns
    -------
    dict
        A dictionary with the class names as keys and the probabilities as values
    """

    pred, pred_idx, probs = learn.predict(img)
    print(probs)
    return {categories[i]: float(probs[i]) for i in range(2)}


image = gr.inputs.Image(shape=(224, 224))
label = gr.outputs.Label()
examples = ["normal_brain_mri.jpeg", "images2.jpeg"]

gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples).launch(
    inline=False
)