import gradio as gr
import os
import torch

from model import create_effnetb0_model
from timeit import default_timer as timer
from typing import Tuple, Dict


class_names = ["pizza", "steak", "sushi"]

effnetb0, effnetb0_transforms = create_effnetb0_model()

effnetb0.load_state_dict(torch.load(
    "07_effnetb0_data_20_percent_10_epochs.pth", weights_only=True, map_location=torch.device('cpu')))


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

    pred_dict = {}

    start_time = timer()

    img = effnetb0_transforms(img).unsqueeze(0)

    effnetb0.eval()

    with torch.inference_mode():
        pred_probs = torch.softmax(effnetb0(img), dim=1)

        # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
        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, 4)

    pred_list.append(pred_dict)

    return pred_labels_and_probs, pred_time


title = "FoodVision Mini 🍕🥩🍣"
description = "An EfficientNetB0 feature extractor computer vision model to classify images of food as pizza, steak or sushi."
article = "Full Source code from scratch [deployment.ipynb](https://github.com/Victoran0/food-vision.git)."

# Create examples list from "examples/" directory
example_list = [["examples/" + example] for example in os.listdir("examples")]

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)")], examples=example_list, title=title, description=description, article=article)

demo.launch()