import gradio as gr
import os
import torch

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

# Setup class names
with open("class_names.txt", "r") as f: 
  class_names = [food_name.strip() for food_name in  f.readlines()]

# Create Food101 compatible EffNetB2 instance
effnet_transforms,effnetb2_food_101 = create_effnet_b2_instance(num_classes = len(class_names))

# Load the saved model's state_dict()
effnetb2_food_101.load_state_dict(torch.load("effnetb2_food101_dict.pth",map_location = torch.device("cpu")))


def predict(img, model = effnetb2_food_101, transforms = effnet_transforms) -> Tuple[Dict,float]:
  # start a timer
  start_timer = timer()
  # transform the image to be used by the model
  prepreocpressed_image = transforms(img).unsqueeze(0)
  # turn off regularization and parameters
  model.eval()
  with torch.inference_mode():
    prediction = model(prepreocpressed_image)
    probabilities = torch.softmax(prediction,dim = 1)
  prob_dict = {class_names[i]: float(probabilities[0][i]) for i in range(len(class_names))}
  # calculate the time
  end_timer = timer()
  total_time = end_timer - start_timer
  return prob_dict,total_time


# create the gradio app
title = "FoodVision Big Classifier"
description = "An EfficientNetB2 feature extractor trained on the Food101 Dataset to classify across 101 possible classes of food."
article = "Model created using pytorch"

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

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

# Launch the app!
demo.launch()