import torch
import torchvision
from torchvision.models import efficientnet_v2_s, EfficientNet_V2_S_Weights

from torch import nn

from PIL import Image
from model import create_effnet_v2_model

import gradio as gr
import os
from timeit import default_timer as timer

class_names = ['Honda', 'Hyundai', 'Toyota']

effnet_v2, transforms = create_effnet_v2_model(num_classes=len(class_names), weights_path="efficient_net_s_carvision_3.pth")

def predict(image):

    start_time = timer()

    # image = Image.open(image_path)
    image = transforms(image).unsqueeze(0)
    # image = image.to(device)
    output = effnet_v2(image)
    
    effnet_v2.eval()
    with torch.inference_mode():
        probs = torch.softmax(output, dim=1)

    pred_labels_and_probs = {class_names[i]: float(probs[0, i]) for i in range(len(class_names))}

    pred_time = round(timer() - start_time, 5)
    return pred_labels_and_probs, pred_time


### 4. Gradio app ###

# Create title, description and article strings
title = "CarVision 🚗🚘🚙🏎️"
description = """
An EfficientNet V2 model to classify cars as **Honda, Hyundai or Toyota**

[Dataset Used](https://github.com/nicolas-gervais/predicting-car-price-from-scraped-data/tree/master/picture-scraper)

Currently running at **85% Accuracy**
"""
article = "Created by [Akshay Ballal](https://www.akshaymakes.com)"

# Create examples list from "examples/" directory
example_list = [["examples/" + example] for example in os.listdir("examples")]

# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs=gr.Image(type="pil"), # what are the inputs?
                    outputs=[gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?
                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
                    # Create examples list from "examples/" directory
                    examples=example_list, 
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch()