import torch
import torchvision
import gradio as gr
import pathlib
import random

from torch import nn
from typing import Tuple, Dict
from PIL import Image
from timeit import default_timer as timer
from typing import Tuple, Dict


device = 'cuda' if torch.cuda.is_available() else 'cpu'

with open('class-names.txt', 'r') as f:
    class_names = f.read().split('\n')[:-1]

def load_model() -> Tuple[torch.nn.Module, torchvision.transforms.Compose]:
    weights = torchvision.models.ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1
    shufflenet_transforms = weights.transforms()
    shufflenet = torchvision.models.shufflenet_v2_x1_5(weights=weights)
    shufflenet.fc = nn.Linear(in_features=1024, out_features=len(class_names), bias=True)

    state_dict = torch.load('ShuffleNetV2.pt', map_location=device)
    shufflenet.load_state_dict(state_dict)

    return shufflenet, shufflenet_transforms

model, transforms = load_model()


def predict(img) -> Tuple[Dict, float]:
    start = timer()
    model.to(device)
    model.eval()
    with torch.inference_mode():
        transformed_img = transforms(img).to(device)

        logits = model(transformed_img.unsqueeze(dim=0))
        pred_prob = torch.softmax(logits, dim=1)
        pred_dict = {class_names[i]:pred_prob.squeeze(0)[i].item() for i in range(len(class_names))}

    pred_time = round(timer() - start, 5)

    return pred_dict, pred_time


example_paths = list(pathlib.Path('examples').glob("*/*.jpg"))
example_list = [[str(filepath)] for filepath in random.sample(example_paths, k=6)]


title = 'Birds Species Classifier 🐦'
description = 'A [ShuffleNetV2](https://pytorch.org/vision/main/models/shufflenetv2.html) feature extractor computer vision model to classify images of [525 species birds](https://www.kaggle.com/datasets/gpiosenka/100-bird-species/).'


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

demo.launch(debug=False)