import gradio as gr
import os
import torch
from timeit import default_timer as timer
from typing import Tuple, Dict

from model import create_model

#get the class names
classes = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']

#get the model and load its trained weights
model, transform = create_model(num_classes = len(classes))
model.load_state_dict(torch.load(f = "model_4epochs_90acc.pth" , map_location = torch.device("cpu")))

#prediction function for a single image
def predict(img):
    start_time = timer()
    img = transform(img).unsqueeze(0)
    
    #get the prediction probabilities and put them in a dictionary
    model.eval()
    with torch.inference_mode():
        y_prob = model(img).softmax(dim = 1)
        y_preds = {classes[i] : float(y_prob[0][i]) for i in range(len(classes)) }
        
    prediction_time = round(timer() - start_time, 5)
    
    return y_preds, prediction_time


#Gradio App ###
title = "Intel Scenery Classification"
description = "An efficientnet_b2 model for the classification of different image scenes from an intel dataset."
article = "Created by me, Richard Schattner."

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

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

# Launch the demo
demo.launch()