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

#model initialization
model= models.resnet18(pretrained=True)

model.fc = torch.nn.Linear(model.fc.in_features,2)
state_dict=torch.load('up500Model.pt', map_location='cpu')
model.load_state_dict(state_dict)
model.eval()

#predictions
imgTransforms = 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 =['Fiat 500', 'VW Up']
def predict(inp):
    inp= Image.fromarray(inp.astype('unit8'), 'RGB')
    inp =imgTransforms(inp).unsqueeze(0)

    with torch.no_grad():
        predictions = torch.nn.functional.softmax(model(inp)[0])

    return {LABELS[i]: float(predictions[i]) for i in range(2)}
examples=[["fiat500.jpg"],["VWUP.jpg"]]
#app deploy
interface = gr.Interface(fn=predict, inputs='image', outputs='label', title='App', description= 'upload the car image', examples=examples, cache_examples= False)
interface.launch()