import gradio as gr
import os
import torch
import timeit

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


#setup classnames

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

#model and transforms preparation

effnetb2,eff_b2_trans=create_effnetb2_model(num_classes=3,seed=42)


#load saved weights

effnetb2.load_state_dict(torch.load(f="eff_model_food24.pth",
                                    map_location=torch.device('cpu')))

#write predict function

def predict(img):


  #start a timer

  start_timer=timeit.default_timer()

  #turn image into tensor


  #Transform the input image for use with effnetb2

  img=eff_b2_trans(img).unsqueeze(0)

  #put model into eval mode,make prediction`

  effnetb2.eval()

  with torch.inference_mode():
    #pass transformed image thorugh the model and turn the prediction logits into probabilities
    pred_probs=torch.softmax(effnetb2(img),dim=1)

  #create a prediction label and prediction probability dictionary

  pred_labels_and_probs={'pizza':pred_probs[0][0],
                         'steak':pred_probs[0][1],
                         'sushi':pred_probs[0][2]}
  #pred_labels_and_probs=pred_probs

  #calculate pred time

  end_time=timeit.default_timer()

  pred_time=end_time-start_timer



  #return pred dictionary and pred time

  return pred_labels_and_probs,f'{pred_time:.4f}'


##gradio app

#create title,desription and article

title="Foodvision Mini"
description="An efficientB2 feature extractor"
article="created at home"

#create example list

example_list=['examples/'+i for i in os.listdir('examples')]

#create the gradio 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 (s)")],
                  examples=example_list,
                  title=title,
                  description=description,
                  article=article)

#launch the demo!
demo.launch(debug=False,share=True)