changed the UI

app.py

@@ -3,12 +3,13 @@ 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.load_state_dict(torch.load("model_best.pt", map_location=torch.device(device)))
     model.eval()
 
     with open('labels.json') as f:
@@ -25,5 +26,12 @@ def pokemon_classifier(inp):
     
     return labels[str(pokemon)]
 
-demo = gr.Interface(fn=pokemon_classifier, inputs=gr.Image(type="pil"), outputs="text")
+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()