import gradio as gr
import torch
from torchvision import transforms

model = torch.jit.load("./models/cat_dog_cnn.pt")
model.eval()


transform = transforms.Compose([
  transforms.Resize((224,224)),
  transforms.ToTensor(),
  transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
  ])

CLASSES = ["Cat", "Dog", "Panda"]

def classify_image(inp):
  inp = transform(inp).unsqueeze(0)
  out = model(inp)
  return CLASSES[out.argmax().item()]

iface = gr.Interface(fn=classify_image,
                     inputs=gr.Image(type="pil", label="Input Image"),
                     outputs="text",
                     examples=[

                       "./app_data/cat.jpg",
                       "./app_data/dog.jpg",
                       "./app_data/panda.jpg",
                     
                               
                     ])
iface.launch()