import torch
from utils import process_image, load_saved_model
import gradio as gr

MODEL_PATH = "model.pt"
CATEGORIES = ("Dog", "Cat")
model = load_saved_model(MODEL_PATH)

def predict_pet(image):
    with torch.no_grad():
        model.eval()
        x = process_image(image)
        probs = model(x).squeeze().tolist()

    return dict(zip(CATEGORIES, probs))

demo = gr.Interface(fn=predict_pet, inputs=gr.Image(label="Image"), outputs=gr.Label(label="Type of Pet"), allow_flagging="never",
                    title="Cat or Dog ?", examples="examples", description="This is a small image classification model for cats and dogs")
demo.launch()