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
from pathlib import Path

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

effnetb2, effnetb2_transforms=create_effnetb2(num_classes=3, seed=42)

effnetb2.load_state_dict(torch.load(f="effnetb2_20%_e10.pth",  map_location=torch.device("cpu"),
                                    weights_only=True))

def predict(img)->Tuple[Dict, float]:
    
    start_time=timer()
    
    img=effnetb2_transforms(img).unsqueeze(dim=0)
    
    effnetb2.eval()
    with torch.inference_mode():
        pred_probs=torch.softmax(effnetb2(img), dim=1)
        
    pred_labels={class_names[i]:round(pred_probs[0][i].item(),3) for i in range(len(class_names))}
    
    pred_time=round(timer()-start_time,3)
    
    return pred_labels, pred_time

title="BiteVision Mini 🍕 🍣 🥩"
description="Drag/upload an image out of 🍕 pizza  🍣 sushi  🥩 steak, and this BiteVision Mini will classify it accordingly.🤩" 
article="by Aakash Haldankar"

example_list=[["examples/"+i] for i in os.listdir("examples")]

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

demo.launch()