from model import create_effnetb2
from typing import Tuple, Dict
from PIL import Image

from time import time
import torch
import torchvision
import gradio as gr
import os

from pathlib import Path

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

effnetb2 , effnetb2_transforms = create_effnetb2()
# Load weights
PATH = "09_pretrained_effnetb2_feature_extractor_pizza20%_10epochs.pth"

effnetb2.load_state_dict(torch.load(f=PATH,
                                   map_location=torch.device('cpu')
                                   ))
effnetb2.eval()




def predict(img) ->Tuple[Dict, float]:
    start_time = time()

    img_tr = img
    img_tr = effnetb2_transforms(img_tr).unsqueeze(0)
    #predict
    effnetb2.eval()
    with torch.inference_mode():
     
        pred_prob = torch.softmax(effnetb2(img_tr), dim=1)
       

    pred_labesls_and_pobs ={class_names[i]:pred_prob[0][i] for i in range(len(class_names)) }
    
 


    end_time = time()
    pred_time =  round(end_time - start_time,4)
    return pred_labesls_and_pobs ,pred_time


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

# Create title, description and article strings
title = "FoodVision Classification"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of food as pizza, steak or sushi."
article = "Created at [Using pre-trained model efficientnet_b2](https://pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b2.html)."

# Create the Gradio demo
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)

# Launch the demo!
demo.launch(debug=False)