import gradio as gr
import torch
import torchvision.transforms.functional as TF
from model import NeuralNetwork
import json
import os

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

def pokemon_classifier(inp):
    model = NeuralNetwork()
    model.load_state_dict(torch.load("model_best.pt", map_location=torch.device(device)))
    model.eval()

    with open('labels.json') as f:
        labels = json.load(f)

    x = TF.to_tensor(inp)
    x = TF.resize(x, 64, antialias=True)
    x = x.to(device)
    x = x.unsqueeze(0)
    
    with torch.no_grad():
        y_pred = model(x)
    pokemon = torch.argmax(y_pred, dim=1).item()
    
    return labels[str(pokemon)]

with gr.Blocks() as demo:
    gr.Markdown("# Gen 1 Pokemon classifier")
    with gr.Row():
        inp = gr.Image(type="pil")
        out = gr.Textbox()
    btn = gr.Button("Run")
    btn.click(fn=pokemon_classifier, inputs=inp, outputs=out)

demo.launch()