shekharkoirala commited on
Commit
cebf2b0
1 Parent(s): 21defc7

added test resnet model

Browse files
Files changed (1) hide show
  1. app.py +22 -0
app.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import requests
3
+ import gradio as gr
4
+ from PIL import Image
5
+ from torchvision import transforms
6
+
7
+ model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()
8
+
9
+ # Download human-readable labels for ImageNet.
10
+ response = requests.get("https://git.io/JJkYN")
11
+ labels = response.text.split("\n")
12
+
13
+ def predict(inp):
14
+ inp = Image.fromarray(inp.astype('uint8'), 'RGB')
15
+ inp = transforms.ToTensor()(inp).unsqueeze(0)
16
+ with torch.no_grad():
17
+ prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
18
+ return {labels[i]: float(prediction[i]) for i in range(1000)}
19
+
20
+ inputs = gr.inputs.Image()
21
+ outputs = gr.outputs.Label(num_top_classes=3)
22
+ gr.Interface(fn=predict, inputs=inputs, outputs=outputs).launch()