#imports and class names setup

import gradio as gr
import os
import torch
import timeit
from model import create_eff_model
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.strip() for food in f.readlines()]

##model and transforms

eff_model,eff_trans=create_eff_model()

eff_model.load_state_dict(torch.load(f="eff_b2_20.pth",
                          map_location=torch.device('cpu')))


##predict function

def predict(img):


  #start a timer

  start_timer=timeit.default_timer()

  #Transform the input image for use with effnetb2

  img=eff_trans(img).unsqueeze(0)

  #put model into eval mode,make prediction`

  eff_model.eval()

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

  #create a prediction label and prediction probability dictionary

  pred_labels_and_probs={class_names[i]:float(pred_probs[0][i]) for i in range(len(class_names))}
  #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 function


#create title,description and article

title="Foodvision BIG"
description="An efficientnet B2 feature extractor"
article="Created in garage"

#create example list

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=5,label="Predictions"),
                      gr.Number(label="Prediction time (s)")],
             examples=example_list,
             title=title,
             description=description,
             article=article)
demo.launch()