akhaliq HF staff commited on
Commit
7b01c04
1 Parent(s): 38144ab

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +57 -0
app.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from PIL import Image
3
+ from torchvision import transforms
4
+ import gradio as gr
5
+ import os
6
+
7
+ os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")
8
+
9
+ target_platform = "proxyless_cpu"
10
+ # proxyless_gpu, proxyless_mobile, proxyless_mobile14 are also avaliable.
11
+ model = torch.hub.load('mit-han-lab/ProxylessNAS', target_platform, pretrained=True)
12
+ model.eval()
13
+ # Download an example image from the pytorch website
14
+ torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
15
+ # sample execution (requires torchvision)
16
+
17
+ def inference(input_image):
18
+ preprocess = transforms.Compose([
19
+ transforms.Resize(256),
20
+ transforms.CenterCrop(224),
21
+ transforms.ToTensor(),
22
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
23
+ ])
24
+ input_tensor = preprocess(input_image)
25
+ input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
26
+
27
+ # move the input and model to GPU for speed if available
28
+ if torch.cuda.is_available():
29
+ input_batch = input_batch.to('cuda')
30
+ model.to('cuda')
31
+
32
+ with torch.no_grad():
33
+ output = model(input_batch)
34
+ # The output has unnormalized scores. To get probabilities, you can run a softmax on it.
35
+ probabilities = torch.nn.functional.softmax(output[0], dim=0)
36
+
37
+ # Read the categories
38
+ with open("imagenet_classes.txt", "r") as f:
39
+ categories = [s.strip() for s in f.readlines()]
40
+ # Show top categories per image
41
+ top5_prob, top5_catid = torch.topk(probabilities, 5)
42
+ result = {}
43
+ for i in range(top5_prob.size(0)):
44
+ result[categories[top5_catid[i]]] = top5_prob[i].item()
45
+ return result
46
+
47
+ inputs = gr.inputs.Image(type='pil')
48
+ outputs = gr.outputs.Label(type="confidences",num_top_classes=5)
49
+
50
+ title = "PROXYLESSNAS"
51
+ description = "Gradio demo for PROXYLESSNAS, Proxylessly specialize CNN architectures for different hardware platforms. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
52
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1812.00332'>ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware</a> | <a href='https://github.com/mit-han-lab/ProxylessNAS'>Github Repo</a></p>"
53
+
54
+ examples = [
55
+ ['dog.jpg']
56
+ ]
57
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()