# Imports what we need

import os
import torch
from timeit import default_timer as timer
from typing import Tuple,Dict
import gradio as gr
from model import vit_model


# Setup Classes

with open("class_names.txt", "r") as f:
  class_names  = [food_names.strip() for food_names in f.readlines()]



# Setup models and transforms

vit, vit_transform = vit_model(num_classes = len(class_names))


# load_model

vit.load_state_dict(
    torch.load(
        f = "pretrained_vit_model.pth",
        map_location = "cpu"
    )
)

# predict function

def predict(img):

  # 1. Start timer
  start_time = timer()

  # 2. Transform the input image for use in VIT

  transformed_img = vit_transform(img).unsqueeze(0)

  # 3. Put model into eval model and make prediction

  vit.eval()
  with torch.inference_mode():
    pred_logit = vit(transformed_img)

    pred = torch.softmax(pred_logit,dim=1)

    y_pred = torch.argmax(pred,dim=1)

    class_label = class_names[y_pred]
# 4. Create pred label ane pred prob dictionary


  pred_labels_and_probs = {class_names[i]: float(pred[0][i]) for i in range(len(class_names))}



# 5. Calculate predtime
  end_time = timer()

  Total_predtime = end_time-start_time


  return (pred_labels_and_probs,Total_predtime)


# 4. Gradio app

# 4a. Creating the example list
example_list = [["examples/" + example] for example in os.listdir("examples")]
example_list

#create title, decription, article




title = "VIT MODEL ON FOOD VISION DATASET (101 CLASSES)"
description = "A basic Vit model to identify foods"
article = "Created for testing"


demo = gr.Interface(fn =predict,
                    inputs = gr.Image(type ="pil"),
                    outputs = [gr.Label(num_top_classes=3,label ="Prediction"),
                              gr.Number(label = "Prediction Time (s)")],
                    examples = example_list,
                    title= title,
                    article =article,
                    description =  description,
                    )

demo.launch(debug= False)