### 1. Imports and class names setup ### 
import gradio as gr
import os
import torch
from torch import nn
from typing import Dict
import torchvision
from torchvision import transforms
import PIL
import numpy

# import os
# os.system("pip uninstall -y gradio")
# os.system("pip install gradio==3.50.2")

with open("class_names.txt", "r") as f: # reading them in from class_names.txt
    class_names = [food_name.strip() for food_name in  f.readlines()]


mobileNetV3_transform = torchvision.models.MobileNet_V3_Large_Weights.DEFAULT.transforms()

mobileNetV3 = torchvision.models.mobilenet_v3_large()
    
mobileNetV3.classifier = nn.Sequential(
  nn.Linear(in_features=960, out_features=1280, bias=True),
  nn.Hardswish(),
  nn.Dropout(p=0.2, inplace=True),
  nn.Linear(in_features=1280, out_features=len(class_names), bias=True)
)

mobileNetV3.load_state_dict(torch.load("models/mobileNetV3_quickdraw_animals_epoch_149_Adam.pth",map_location=torch.device('cpu')))



def convert_img_to_tensor(img):
    convert_img = torch.from_numpy(img)
    convert_img = convert_img.repeat(3, 1, 1)
    return convert_img

def predict(img) -> Dict:
    """Transforms and performs a prediction on img and returns prediction.
    """
    img = convert_img_to_tensor(img)
    img = mobileNetV3_transform(img).unsqueeze(0)

    # Put model into evaluation mode and turn on inference mode
    mobileNetV3.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        y_pred = mobileNetV3(img)
        pred_probs = torch.softmax(y_pred, dim=1)

    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}

    # Return the prediction dictionary and prediction time
    return pred_labels_and_probs

link_to_code = "https://github.com/henrythdu/MobileNetV3-QuickDraw-Animals"
title = "MobileNetV3 - Quick Draw - Animals 🔢"
description = "An MobileNetV3 feature extractor computer vision model to classify doodling of animals from Quick Draw dataset. The goal is to have a small model without significantly compromising on the predictive performance"
article = f"Created using transfer learning from - <a href={link_to_code}>MobileNetV3</a>"

sp = gr.Sketchpad(shape = (28,28), brush_radius = 1)
demo = gr.Interface(
    fn=predict,
    inputs=sp,
    outputs= gr.Label(num_top_classes=5, label="Predictions"),
    title=title,
    description=description,
    article=article
)
demo.launch(debug=True)