import gradio as gr
import os
import torch
 
from model import create_effnetb2_model
from timeit import default_timer as timer
from typing import Tuple, Dict

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

effnetb2, effnetb2_transforms= create_effnetb2_model(num_classes= 3)

effnetb2.load_state_dict(torch.load(f="09_pretrained_effnetb2_feature_extractor__pizza_steak_sushi_20_percent.pth",map_location=
 torch.device("cpu")
)) 


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

  start_time= timer()
  img= effnetb2_transforms(img).unsqueeze(0)
  effnetb2.eval()
  with torch.inference_mode():
    pred_probs= torch.softmax(effnetb2(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


titile= "Foodvision Mini"
description="An efficientnetb2 feature extractor computer vision model to classify images of pizza, steak and sushi."
article= "Created at [09_Pytorch model deployment] (https://www.learnpytorch.io/09_pytorch_model_deployment/)"
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)")],
                             example_list= example_list, title= title, description= description, article= article)


demo.launch()