### 1. Imports and class names setup ### 
import gradio as gr
import os
import torch

from model import create_resnet_model, create_custom_model
from timeit import default_timer as timer
import torchvision
import torchvision.transforms as transforms

transformer = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(256),
    transforms.ToTensor(),
    transforms.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
])


model_name = 'resnet'
### 2. Model and transforms preparation ###

if model_name == 'custom':
    # Create model
    model  = create_custom_model()

    # Load saved weights
    model.load_state_dict(
        torch.load(
            f="./cnn-custom-model-version-4.pt",
            map_location=torch.device("cpu"),  # load to CPU
        )
    )
elif model_name == 'resnet':
    model  = create_resnet_model()

    # Load saved weights
    model.load_state_dict(
        torch.load(
            f="./cnn-resnet-version-1.pt",
            map_location=torch.device("cpu"),  # load to CPU
        )
    )
# else:


### 3. Predict function ###

def predict(img):
    """Transforms and performs a prediction on img and returns prediction and time taken.
    """
    # Transform the target image and add a batch dimension
    img = transformer(img).unsqueeze(0)
    
    # Put model into evaluation mode and turn on inference mode
    model.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        pred_prob = torch.sigmoid(model(img))
        
    pred_probs = {'Covid' : float(pred_prob), 'Non Covid' : (1-float(pred_prob))}

    # Return the prediction dictionary and prediction time 
    return pred_probs

### 4. Gradio app ###


# Create title, description and article strings
title = "Corona Prediction"
description = "A Convolutional Neural Network To classify whether a person have Corona or not using CT  Scans."
article = "Created by Thenujan Nagaratnam for DNN module at UoM"

# 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=2, label="Predictions")], # our fn has two outputs, therefore we have two outputs
                    examples=example_list, 
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch()