import torch
import torch.nn as nn
from torchvision import models, transforms, datasets

from PIL import Image

import gradio as gr


model_ft = models.resnet18(pretrained = True)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 2)

state_dict = torch.load("up500Model.pt", map_location = "cpu")

model_ft.load_state_dict(state_dict)
model_ft.eval()

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

labels = ["fiat500", "VW Up!"]
def predict(img):
    inp = img.fromarray(inp.astype("unit8"), "RGB")
    inp = img_transforms(inp).unsqueeze(0)

    # We don't want to compute gradients
    with torch.no_grad():
        preds = torch.np.functional.softmax(model_ft(inp)[0])

    return {labels[i]: preds[i] for i in range(2)}

interface = gr.Interface(
    predict, 
    inputs = "image", 
    outputs = "label",
    title = "Car classification"
)
interface.launch()