pranavbapte's picture
Update app.py
ee2cd5c
import torch
import gradio as gr
from torch import nn
from torch.nn import functional as F
import torchvision
from PIL import Image
from torchvision import transforms
transformer = transforms.Compose([
transforms.Resize((256, 256)),
# transforms.RandomHorizontalFlip(),
#transforms.RandomRotation(degrees=10),
transforms.ToTensor()
#transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225
#])
])
model=torch.jit.load('model.pt',map_location=torch.device('cpu'))
#model=torch.jit.load('model1.pt')
classes=['Minivan Car', 'Muscle Car ', 'Sedan Car', 'Sports Car', 'None of the Above class']
def predict(inp):
inp=transformer(inp).unsqueeze(0)
#inp = transforms.ToTensor()(inp).unsqueeze(0)
with torch.no_grad():
prediction =F.softmax(model(inp)[0], dim=0)
confidences = {classes[i]: float(prediction[i]) for i in range(5)}
return confidences
# gr.Interface(fn=predict, inputs=gr.Image(type="pil"),outputs=gr.Label(num_top_classes=4),title='Image classification',interpretation='default').launch(debug='True')
gr.Interface(predict, gr.inputs.Image(type="pil"),outputs='label',title='Image classification').launch(debug='True')