app.py

# Imports what we need


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


# Setup class names

class_names = ['pizza', 'steak', 'sushi']

# 2. Model and transforms preparation


vit, vit_transform = vit_model(num_classes =3)

# 2a. load save weights

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

# 3. Prediction 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 = "Test vision 1"
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)