akhaliq HF staff commited on
Commit
9297431
1 Parent(s): 5abfd48

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +55 -0
app.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from PIL import Image
3
+ from torchvision import transforms
4
+ import gradio as gr
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', 'resnext50_32x4d', pretrained=True)
9
+ # or
10
+ # model = torch.hub.load('pytorch/vision:v0.9.0', 'resnext101_32x8d', pretrained=True)
11
+ model.eval()
12
+
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
+ !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
+
38
+ # Show top categories per image
39
+ top5_prob, top5_catid = torch.topk(probabilities, 5)
40
+ result = {}
41
+ for i in range(top5_prob.size(0)):
42
+ result[categories[top5_catid[i]]] = top5_prob[i].item()
43
+ return result
44
+
45
+ inputs = gr.inputs.Image(type='pil')
46
+ outputs = gr.outputs.Label(type="confidences",num_top_classes=5)
47
+
48
+ title = "RESNEXT"
49
+ description = "Gradio demo for RESNEXT, Next generation ResNets, more efficient and accurate. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
50
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1611.05431'>Aggregated Residual Transformations for Deep Neural Networks</a> | <a href='https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py'>Github Repo</a></p>"
51
+
52
+ examples = [
53
+ ['dog.jpg']
54
+ ]
55
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()