import gradio as gr
import torch
from model import create_effnet_b2_model
from timeit import default_timer as timer
from typing import Tuple
from PIL import Image
import requests
from io import BytesIO

# Setup class names
class_names = ['paper', 'rock', 'scissors']

# Model and transforms
model, transforms = create_effnet_b2_model(num_classes=len(class_names))

# Load saved weights
model.load_state_dict(torch.load(f="final_model.pth", map_location=torch.device('cpu')))

# Predict function
def predict(img) -> Tuple[str, float]:
    # Start the timer
    start_time = timer()

    # Transform image and add batch dimension
    transformed_image = transforms(img).unsqueeze(dim=0)

    # Put model into evaluation mode and turn on inference mode
    model.eval()
    with torch.inference_mode():
        # Pass image through model to get raw logits
        target_image_pred = model(transformed_image)
        # Convert prediction logits into prediction probabilities
        target_image_pred_prob = torch.softmax(target_image_pred, dim = 1) # e.g. tensor([[0.0522, 0.9266, 0.0212]])

    # Crete prediction label and prediction probability dictionary for each class
    pred_labels_and_probs = {class_names[i]: float(target_image_pred_prob[0][i]) for i in range(len(class_names))}

    # Calculate the prediction time
    total_time = round(timer() - start_time, 5)

    # Return prediction dictionary and prediction time
    return pred_labels_and_probs, total_time

# Gradio app
title = "Rock paper scissors"
description = "Detects hand sign (rock, paper, scissors) from provided image."
article = "Dataset taken from https://www.kaggle.com/datasets/alexandredj/rock-paper-scissors-dataset?select=scissors. Used transfer learning to get 98% accuracy using EfficientNetB2 model."

# Specify the example image paths
example_paths = [
    "https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_1.png",
    "https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_2.png",
    "https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_3.png",
    "https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_4.png",
    "https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_5.png",
    "https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_6.png",
]

# Create example list
example_list = [[path] for path in example_paths]

demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type='pil'),
                    outputs=[gr.Label(num_top_classes=3, label="Predictions"), gr.Number(label="Prediction time")],
                    examples=example_list,
                    title=title,
                    description=description,
                    article=article)

demo.launch(share=True)