silaseic commited on
Commit
49638b7
1 Parent(s): 9b4f909

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +30 -0
main.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import requests
4
+ from torchvision import transforms
5
+
6
+ model = torch.hub.load("pytorch/vision:v0.6.0", "resnet18", pretrained=True).eval()
7
+ response = requests.get("https://git.io/JJkYN")
8
+ labels = response.text.split("\n")
9
+
10
+
11
+ def predict(inp):
12
+ inp = transforms.ToTensor()(inp).unsqueeze(0)
13
+ with torch.no_grad():
14
+ prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
15
+ confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
16
+ return confidences
17
+
18
+
19
+ def run():
20
+ demo = gr.Interface(
21
+ fn=predict,
22
+ inputs=gr.inputs.Image(type="pil"),
23
+ outputs=gr.outputs.Label(num_top_classes=3),
24
+ )
25
+
26
+ demo.launch(server_name="0.0.0.0", server_port=7860)
27
+
28
+
29
+ if __name__ == "__main__":
30
+ run()