import os
from timeit import default_timer as timer
from typing import Tuple, Dict

import gradio as gr
import torch
from model import create_model


title = "Food Vision Mini by Edesak"
desc = "EffitientNetB2 for recognition of Pizza,Steak,Sushi from [Zero To Mastery Course](https://www.udemy.com/course/pytorch-for-deep-learning/)"
article = "My Github page [Edesak](https://github.com/Edesak)"

class_names = ["pizza", "steak", "sushi"]
model,transform = create_model()
model.load_state_dict(torch.load(f="_Deploy_effB2.pth",map_location=torch.device('cpu')))
model.to("cpu")
example_list = [["Examples/" + example] for example in os.listdir("Examples")]

def predict(img) -> Tuple[Dict, float]:
    start_timer = timer()

    img = transform(img).unsqueeze(0)

    model.eval()
    with torch.inference_mode():
        y = torch.softmax(model(img), dim=1)
    pred_labels = {class_names[i]: float(y[0][i]) for i in range(len(class_names))}
    end_time = timer()
    pred_time = round(end_time - start_timer, 4)

    return pred_labels, pred_time


def greet(name):
    return "Hello " + name + "!"


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

demo.launch()