yuhe6 commited on
Commit
32a0c88
1 Parent(s): 7944111

Upload 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
+ os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
8
+
9
+ model = torch.hub.load('huawei-noah/ghostnet', 'ghostnet_1x', pretrained=True)
10
+ model.eval()
11
+ # Download an example image from the pytorch website
12
+ torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
13
+
14
+ def inference(input_image):
15
+ preprocess = transforms.Compose([
16
+ transforms.Resize(256),
17
+ transforms.CenterCrop(224),
18
+ transforms.ToTensor(),
19
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
20
+ ])
21
+ input_tensor = preprocess(input_image)
22
+ input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
23
+
24
+ # move the input and model to GPU for speed if available
25
+ if torch.cuda.is_available():
26
+ input_batch = input_batch.to('cuda')
27
+ model.to('cuda')
28
+
29
+ with torch.no_grad():
30
+ output = model(input_batch)
31
+ # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
32
+ probabilities = torch.nn.functional.softmax(output[0], dim=0)
33
+
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.Label(type="confidences",num_top_classes=5)
46
+
47
+ title = "GHOSTNET"
48
+ description = "Gradio demo for GHOSTNET, Efficient networks by generating more features from cheap operations. 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/1911.11907'>GhostNet: More Features from Cheap Operations</a> | <a href='https://github.com/huawei-noah/CV-Backbones'>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()