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

with open("class_names.txt","r") as f:
    class_names = [food_name.strip() for food_name in f.readlines()]
effnetb2,effnetb2_transforms=create_effnetb2(classes=7)
effnetb2.load_state_dict(torch.load(f="0pr_food.pth",map_location=torch.device("cpu")))
def predict(img)-> Tuple[Dict,float]:
    start_timer=timer()
    img=effnetb2_transforms(img).unsqueeze(0)
    effnetb2.eval()
    with torch.inference_mode():
        pred_probs=torch.softmax(effnetb2(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_timer,5)
    return pred_labels_and_probs,pred_time
title="FoodVision"
description="An EfficientNetB2 feature extractor computer vision model to classify images of 7 food [samosa,pizza,steak,sushi,cup cakes,french fries,omelette]"
article="Created at [Foodvision.ipynb]"
example_list=[["Ex/"+ example] for example in os.listdir("Ex")]
demo=gr.Interface(fn=predict,inputs=gr.Image(type="pil"),outputs=[gr.Label(num_top_classes=2,label="Preditions"),gr.Number(label="Predcition time(s)")],examples=example_list,title=title,description=description,article=article)
demo.launch()