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

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

demo = gr.Interface(fn=pokemon_classifier, inputs=gr.Image(type="pil"), outputs="text")
demo.launch()