import gradio as gr 
import torchvision 
from torchvision import models
from torch import nn 


import torch

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

#class names 
with open('class_names.txt', "r") as f:
    class_names  = [car.strip() for car in  f.readlines()]


#model and transforms preparation
effnetb0_weights = models.EfficientNet_B0_Weights.DEFAULT
effnetb0 = torchvision.models.efficientnet_b0(weights = effnetb0_weights)
effnetb0_transforms = effnetb0_weights.transforms()

#freeze params 
for param in effnetb0.parameters(): 
  param.requires_grad = False 
        
#change classifier 
effnetb0.classifier = nn.Sequential(
    nn.Dropout(p=.2),
    nn.Linear(in_features = 1280, 
              out_features = 196)
)

#load saved weights 
effnetb0.load_state_dict(torch.load('pretrained_effnetb0_stanford_cars_20_percent.pth', 
                          map_location=torch.device("cpu"))


#predict function 

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

  start_time = timer()

  #put model into eval mode 
  effnetb0.eval()

  with torch.inference_mode():
    pred_logits = effnetb0(img.unsqueeze(0))
    pred_probs = torch.softmax(pred_logits, dim = 1)

   # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
  pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}

  end_time = timer()

  time = round(end_time - start_time, 5) 

  return pred_labels_and_probs, time


#gradio app 

title = 'effnetb0'
description = 'Pretrained effnetb0 model on stanford cars dataset'

#create example list 
example_list = [["examples/" + example] for example in os.listdir("examples")]

# Create Gradio interface 
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
    
)

# Launch the app!
demo.launch()