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

effnetb2 ,effnetb2_transforms = create_effnetb2_model()

effnetb2.load_state_dict(torch.load(
    f="09_pretrained_effnetb2.pth" ,
    map_location=torch.device('cpu')
        )
     )
class_name = ['pizza' ,'steak' ,'sushi']
def predict(img) -> Tuple[Dict ,float]:
  start_time = timer()
  img = effnetb2_transforms(img).unsqueeze(dim=0)

  effnetb2.eval()
  with torch.inference_mode():
    logit = effnetb2(img)
    pred_probs = torch.softmax(logit ,dim=1)

  pred_labels_probs = {class_name[i] : float(pred_probs[0][i]) for i in range(len(class_name))}
  pred_time = round(timer() - start_time ,5)

  return pred_labels_probs ,pred_time

title = "FoodVision Mini 🍕🥩🍣"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of food as pizza, steak or 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)')] ,
                    examples=example_list ,
                    title=title ,
                    description=description ,
                    article=article
                    )
demo.launch()