akhaliq HF staff commited on
Commit
e5cdad5
1 Parent(s): ac5e9c6

Create app.py

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