akhaliq HF staff commited on
Commit
29a65ce
Β·
1 Parent(s): bff7599

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +73 -0
app.py ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+
4
+ from matplotlib import pyplot as plt
5
+ import matplotlib.patches as patches
6
+ import gradio as gr
7
+
8
+
9
+ # Images
10
+ torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000397133.jpg', 'example1.jpg')
11
+ torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000037777.jpg', 'example2.jpg')
12
+ torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000252219.jpg', 'example3.jpg')
13
+
14
+
15
+ ssd_model = torch.hub.load('AK391/DeepLearningExamples:torchhub', 'nvidia_ssd',pretrained=False,force_reload=True)
16
+
17
+ checkpoint = torch.hub.load_state_dict_from_url('https://api.ngc.nvidia.com/v2/models/nvidia/ssd_pyt_ckpt_amp/versions/20.06.0/files/nvidia_ssdpyt_amp_200703.pt', map_location="cpu")
18
+
19
+
20
+ ssd_model.load_state_dict(checkpoint['model'])
21
+
22
+ utils = torch.hub.load('AK391/DeepLearningExamples', 'nvidia_ssd_processing_utils',force_reload=True)
23
+
24
+ ssd_model.to('cpu')
25
+ ssd_model.eval()
26
+
27
+
28
+ def inference(img):
29
+
30
+ uris = [
31
+ img.name
32
+ ]
33
+
34
+ inputs = [utils.prepare_input(uri) for uri in uris]
35
+ tensor = utils.prepare_tensor(inputs)
36
+
37
+ with torch.no_grad():
38
+ detections_batch = ssd_model(tensor)
39
+
40
+ results_per_input = utils.decode_results(detections_batch)
41
+ best_results_per_input = [utils.pick_best(results, 0.40) for results in results_per_input]
42
+
43
+ classes_to_labels = utils.get_coco_object_dictionary()
44
+ for image_idx in range(len(best_results_per_input)):
45
+ fig, ax = plt.subplots(1)
46
+ # Show original, denormalized image...
47
+ image = inputs[image_idx] / 2 + 0.5
48
+ ax.imshow(image)
49
+ # ...with detections
50
+ bboxes, classes, confidences = best_results_per_input[image_idx]
51
+ for idx in range(len(bboxes)):
52
+ left, bot, right, top = bboxes[idx]
53
+ x, y, w, h = [val * 300 for val in [left, bot, right - left, top - bot]]
54
+ rect = patches.Rectangle((x, y), w, h, linewidth=1, edgecolor='r', facecolor='none')
55
+ ax.add_patch(rect)
56
+ ax.text(x, y, "{} {:.0f}%".format(classes_to_labels[classes[idx] - 1], confidences[idx]*100), bbox=dict(facecolor='white', alpha=0.5))
57
+ plt.axis('off')
58
+ plt.draw()
59
+ return plt
60
+
61
+ inputs = gr.inputs.Image(type='file', label="Original Image")
62
+ outputs = gr.outputs.Image(type="plot", label="Output Image")
63
+
64
+ title = "Single Shot MultiBox Detector model for object detection"
65
+ description = "Gradio demo for Single Shot MultiBox Detector model for object detection by Nvidia. To use it upload an image or click an example images images. Read more at the links below"
66
+ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1512.02325'>SSD: Single Shot MultiBox Detector</a> | <a href='https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/Detection/SSD'>Github Repo</a></p>"
67
+
68
+ examples = [
69
+ ['example1.jpg'],
70
+ ['example2.jpg'],
71
+ ['example3.jpg']
72
+ ]
73
+ gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples).launch(debug=True,enable_queue=True)