imkaushalpatel commited on
Commit
3af7f4a
1 Parent(s): 983a1c5

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ from PIL import Image
4
+ from torchvision import transforms
5
+ import gradio as gr
6
+
7
+ model = torch.hub.load('pytorch/vision:v0.9.0', 'googlenet', pretrained=True)
8
+ model.eval()
9
+
10
+ torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
11
+
12
+ # sample execution (requires torchvision)
13
+ def inference(input_image):
14
+ preprocess = transforms.Compose([
15
+ transforms.Resize(256),
16
+ transforms.CenterCrop(224),
17
+ transforms.ToTensor(),
18
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
19
+ ])
20
+ input_tensor = preprocess(input_image)
21
+ input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
22
+
23
+ # move the input and model to GPU for speed if available
24
+ if torch.cuda.is_available():
25
+ input_batch = input_batch.to('cuda')
26
+ model.to('cuda')
27
+
28
+ with torch.no_grad():
29
+ output = model(input_batch)
30
+ # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
31
+ probabilities = torch.nn.functional.softmax(output[0], dim=0)
32
+ # Download ImageNet labels
33
+ os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
34
+ # Read the categories
35
+ with open("imagenet_classes.txt", "r") as f:
36
+ categories = [s.strip() for s in f.readlines()]
37
+ # Show top categories per image
38
+ top5_prob, top5_catid = torch.topk(probabilities, 5)
39
+ result = {}
40
+ for i in range(top5_prob.size(0)):
41
+ result[categories[top5_catid[i]]] = top5_prob[i].item()
42
+ return result
43
+
44
+ inputs = gr.inputs.Image(type='pil')
45
+ outputs = gr.outputs.Image(type="confidences",num_top_classes=5)
46
+
47
+ title = "GOOGLENET"
48
+ description = "Gradio demo for GOOGLENET, GoogLeNet was based on a deep convolutional neural network architecture codenamed Inception which won ImageNet 2014. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
49
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1409.4842'>Going Deeper with Convolutions</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/googlenet.py'>Github Repo</a></p>"
50
+
51
+ examples = [
52
+ ['dog.jpg']
53
+ ]
54
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()