import gradio as gr
import os
import torch

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


# Loading saved weights
model, transforms = create_model(num_classes=13)
model_load_dict = torch.load('iran_cars_model_dict.pth', map_location=torch.device('cpu'))
model.load_state_dict(model_load_dict['state_dict'])
class_names = model_load_dict['class_names']


def predict(img):
  """Transforms and performs a prediction on img 
      and returns prediction and time taken.
  """
  # starting the timer
  start_time = timer()

  # transforming the target image and adding a batch dimention
  img = transforms(img).unsqueeze(0)

  # putting model into evaluation mode
  model.eval()

  # turning on inference_mode in context manager
  with torch.inference_mode():
    pred_probs = torch.softmax(model(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



title = 'Iran Cars ComputerVision_V0 🚗'
description = 'an EfficientNetb0 CV model created by MiladAbdollahi'
article = 'github.com/Milad-Abdollahi'
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=13, label='Predictions'),
                          gr.Number(Label="Prediction time (s)")],
                  examples=example_list,
                  title=title,
                  description=description,
                  article=article
                  )

demo.launch()