### 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

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()]

# Load model architechture 
effnetB3 = torchvision.models.efficientnet_b3()

# Change the last layer to have 10 classes
effnetB3.classifier = nn.Sequential(
    nn.Dropout(p=0.3, inplace = True),
    nn.Linear(in_features=1536, out_features= 10, bias = True)
)

# Load image transform
effnetB3_Weights = torchvision.models.EfficientNet_B3_Weights.DEFAULT
effnetB3_transform = effnetB3_Weights.transforms()
image_transform = transforms.Compose([transforms.Grayscale(num_output_channels=3),
                                                  effnetB3_transform])

# Load trained model weights
effnetB3.load_state_dict(torch.load("models/Mini_MNIST_EffnetB3_epoch_50.pth",map_location=torch.device('cpu')))

def predict(img) -> Dict:
    """Transforms and performs a prediction on img and returns prediction and time taken.
    """

    if type(img) == numpy.ndarray:
      img = PIL.Image.fromarray(img)

    img = image_transform(img).unsqueeze(0)

    
    # Put model into evaluation mode and turn on inference mode
    effnetB3.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        y_pred = effnetB3(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 
    return pred_labels_and_probs
link_to_code ="https://github.com/henrythdu/EfficientNet_B3_MNIST"
title = "MNIST - EfficientNetB3 🔢"
description = "An EfficientNetB3 feature extractor computer vision model to classify images of digits. The goal is to utilize a smaller and less computationally intensive alternative to larger models, without significantly compromising on the predictive performance. "
article = f"Created using transfer learning from EfficientNetB3 - <a href={link_to_code}>link</a>"
sp = gr.Sketchpad(brush_radius=0.8)
demo = gr.Interface(
    fn=predict,
    inputs=sp,
    outputs=gr.Label(num_top_classes=3, label="Predictions"),
    title=title,
    description=description,
    article=article
    )

demo.launch()