demo / app.py
bhardwaj08sarthak's picture
Update app.py
0e11cd3 verified
raw
history blame
1.04 kB
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()