Spaces:
Runtime error
Runtime error
import torch | |
from PIL import Image | |
from torchvision import transforms | |
import gradio as gr | |
import os | |
import torch | |
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True) | |
transformer = transforms.Compose([ | |
transforms.Resize((224, 224)), | |
transforms.RandomHorizontalFlip(), | |
transforms.RandomRotation(degrees=10), | |
transforms.ToTensor(), | |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
]) | |
input_tensor = preprocess(input_image) | |
input_batch = input_tensor.unsqueeze(0) | |
classes=['Bus','bicycle','car'] | |
if torch.cuda.is_available(): | |
input_batch = input_batch.to('cuda') | |
model.to('cuda') | |
with torch.no_grad(): | |
output = model(input_batch) | |
# The output has unnormalized scores. To get probabilities, you can run a softmax on it. | |
probabilities = torch.nn.functional.softmax(output[0], dim=0) | |
# Read the categories | |
with open("imagenet_classes.txt", "r") as f: | |
categories = [s.strip() for s in f.readlines()] | |
# Show top categories per image | |
top5_prob, top5_catid = torch.topk(probabilities, 3) | |
result = {} | |
for i in range(top5_prob.size(0)): | |
result[categories[top5_catid[i]]] = top5_prob[i].item() | |
return result | |
inputs = gr.inputs.Image(type='pil') | |
outputs = gr.outputs.Label(type="confidences",num_top_classes=3) | |
gr.Interface(inference, inputs, outputs, analytics_enabled=False).launch() | |