import gradio as gr
import os
import torch
from model import create_effnetb2
from PIL import Image
from pathlib import Path
from timeit import default_timer as timer

model, transforms = create_effnetb2(101)
model = torch.load("EffnetB2_Big_Label_Smoothning_5_Epochs_50%_data.pt", map_location= torch.device('cpu'))
examples_path = [["examples/" + example] for example in os.listdir("examples")]

with open("class_names.txt", 'r') as f:
    food101_class_names_loaded = [food.strip() for food in f.readlines()]
    f.close()

def predict(img, model= model, transforms= transforms, class_names= food101_class_names_loaded, device= 'cpu'):
    pred_labels_and_probs = {}
    with torch.inference_mode():
        model.eval()
        model = model.to(device)
        start_time = timer()
        
        transformed_img = transforms(img)
        y_pred_probs = torch.softmax(model(transformed_img.unsqueeze(dim= 0)), dim= 1)
        y_pred_class = class_names[torch.argmax(y_pred_probs, dim= 1)]
        for i in range(len(class_names)):
            pred_labels_and_probs[class_names[i]] = y_pred_probs[0][i].item()
        
        end_time = timer()
        pred_time = round(end_time - start_time, 4)
    return (pred_labels_and_probs, pred_time)



title = "Food Vision Big"
description = "An EfficientNetB2 feature extractor to classify food images into 101 classes from the Food101 dataset"
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)")], title= title, 
                    description= description, examples= examples_path)

demo.launch()