import torch
import gradio as gr
from typing import Tuple,Dict
from model import create_model
import os
import time

sample_images=[["Examples/"+path]for path in os.listdir("Examples")]

Model,Eff_Net_transform=create_model(num_classes=3)

Model=Model.to("cpu")

Model.load_state_dict(torch.load("EfficientNet.pth",map_location=torch.device('cpu'),weights_only=True))

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

def predict_xyz(image)->Tuple[Dict,float]:
  image=torch.from_numpy(image)
  image=image.to("cpu")
  start_time=time.time()
  image=Eff_Net_transform(image).unsqueeze(dim=0)
  Model.eval()
  with torch.inference_mode():
    logits=Model(image)
    probability=logits.softmax(dim=1)
    class_to_prob={classes[i]:float(data.item()) for i,data in enumerate(probability[0])}
  end_time=time.time()
  return class_to_prob,end_time-start_time
    
if __name__=="__main__":
  demo=gr.Interface(fn=predict_xyz,
             inputs="image",outputs=[gr.Label(num_top_classes=3, label="Class Probabilities"),gr.Number(label="Prediction Time(s)")],
             examples=sample_images,title="Eff_Net_Prediction")
  demo.launch(debug=False,share=True)