# set up imports and others
import gradio as gr
import os
import torch

from model import create_effnet_b2
from timeit import default_timer as timer

# set up class names
class_names = ["pizza", "steak", "sushi"]

# create the model architecture
effnet_b2, effnet_transforms = create_effnet_b2()

# load the model with pretrained weights
effnet_b2.load_state_dict(torch.load(f = "effnet_b2.pth", map_location = torch.device("cpu"), weights_only = True)) # harcode to cpu

# create a predict function using the model
def predict(image):
  # transform the image and hardcode it to cpu
  transformed_image = effnet_transforms(image).to("cpu")

  # start the timer
  start = timer()
  # pass the image through the model
  effnet_b2.eval()
  with torch.inference_mode():
    # pass
    logits = effnet_b2(transformed_image.unsqueeze(dim = 0))
    pred_probs = torch.softmax(logits, dim = 1)

    # store them in dict like {'pizza': 0.9785208702087402, 'steak': 0.01169557310640812, 'sushi': 0.009783552028238773}
    pred_label_prob = {class_names[i]: pred_probs.squeeze()[i].item() for i  in range(len(class_names))}

    # end timer
    end = timer()

    # time taken
    time_taken = end - start

    return pred_label_prob, round(time_taken, 4)

# create the gradio app

# some setups
title = "Food Vision Mini 🍕🥩🍣"
description = "An effnet b2 feature extractor model to classify three food classes - pizza, steak and sushi. THIS FOR TESTING PURPOSES!"

article = "Created by [Abdiaziz Muse](https://www.abdiazizmuse.com/)"

# create an example dir
example_list = [["examples/" + example]for example in os.listdir("examples") if example.endswith(".jpg")]

# create the demo
demo = gr.Interface(fn = predict,
                    inputs = gr.Image(type = "pil"),
                    outputs = [gr.Label(num_top_classes = 3, label = "Predictions"),
                               gr.Number(label = "Prediction time (seconds)")
                               ],
                    examples = example_list,
                    description = description,
                    article = article
                    )

# output the demo
demo.launch()