import gradio as gr
import torch
from torchvision import models, transforms
from PIL import Image

model = models.resnet18()
model.fc = torch.nn.Linear(512, 10)
model.load_state_dict(torch.load("model.pt"))
model.eval()

transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])


def predict(image):
    image = transform(image).unsqueeze(0)
    with torch.no_grad():
        output = model(image)
    return int(output.argmax())


gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="number",
    title="My First Hugging Face Model",
).launch()