import gradio as gr
import os
import torch

from model import create_effnetb2_model
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_model(num_classes=101)

effnetb2.load_state_dict(
    torch.load(
        f="effnetb2_food101_large.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_prob=torch.softmax(effnetb2(img),dim=1)
    pred_label_and_prob={class_names[i]:float(pred_prob[0][i]) for i in range(len(class_names))}
    pred_time=round(timer()-start_timer,5)

    return pred_label_and_prob,pred_time

title="FoodVision Large"
description="An EfficientNet B2 Feature extractor computer vision model to classify images of food consisting of 101 classes."
article="Created on 22 Jan 2025"
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=5,label="Predictions"),
        gr.Number(label="Prediction time (s)")
    ],
    examples=example_list,
    title=title,
    description=description,
    article=article
)

demo.launch()