import gradio as gr
import os
import torch

from model import create_efficientb2_model
from timeit import default_timer as timer

class_names = [
    "glioma",
    "meningioma",
    "notumor",
    "pituitary"
]

efficientb2, transforms = create_efficientb2_model(num_classes=4)

# Load the entire model directly from the file
my_model_weight = torch.load(
    f="efficientnet_mri_model.pth",
    map_location=torch.device("cpu"),
    weights_only=False
)

efficientb2.load_state_dict(my_model_weight())

def predict_img(img):

  start_time = timer()

  img = transforms(img).unsqueeze(0)

  efficientb2.eval()

  with torch.inference_mode():

    pred_probs = torch.softmax(efficientb2(img), dim=1)

    pred_labels_and_probs = {
        class_names[i] : float(pred_probs[0][i]) for i in range(len(class_names))
    }

    pred_time = round(timer() - start_time(),5)

    return pred_labels_and_probs, pred_time

title = "MRI Result Finder"
description = "Efficientnet b2 model to classify MRI images"
article = "Created at 2026"

example_list = [["examples/" + example] for example in os.listdir("examples")]

demo = gr.Interface(
    fn=predict_img,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=4,label="Predictions"),
        gr.Number(label="Prediction Time")
    ],
    examples = example_list,
    title = title,
    description = description,
    article = article
)


demo.lunch()