/
object_detection.py
80 lines (61 loc) · 2.59 KB
/
object_detection.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import numpy as np
from PIL import Image
def preprocess(image):
ratio = 800.0 / min(image.size[0], image.size[1])
image = image.resize((int(ratio * image.size[0]), int(ratio * image.size[1])), Image.BILINEAR)
image = np.array(image)[:, :, [2, 1, 0]].astype('float32')
image = np.transpose(image, [2, 0, 1])
mean_vec = np.array([102.9801, 115.9465, 122.7717])
for i in range(image.shape[0]):
image[i, :, :] = image[i, :, :] - mean_vec[i]
import math
padded_h = int(math.ceil(image.shape[1] / 32) * 32)
padded_w = int(math.ceil(image.shape[2] / 32) * 32)
padded_image = np.zeros((3, padded_h, padded_w), dtype=np.float32)
padded_image[:, :image.shape[1], :image.shape[2]] = image
image = padded_image
return image
import tempfile
from six.moves.urllib.request import urlopen
from six import BytesIO
def display_image(image):
fig = plt.figure(figsize=(20,15))
plt.grid(False)
plt.imshow(image)
plt.show()
def download_image(url):
_, filename = tempfile.mkstemp(suffix=".jpg")
response = urlopen(url)
image_data = response.read()
image_data = BytesIO(image_data)
pil_image = Image.open(image_data)
pil_image_rgb = pil_image.convert("RGB")
pil_image_rgb.save(filename, format="JPEG", quality=90)
print("Image downloaded to %s." % filename)
return filename
image_url = 'https://upload.wikimedia.org/wikipedia/commons/6/60/Naxos_Taverna.jpg'
downloaded_image_path = download_image(image_url)
img = Image.open(downloaded_image_path)
img_data = preprocess(img)
import onnxruntime
classes = [line.rstrip('\n') for line in open('coco_classes.txt')]
ort_session = onnxruntime.InferenceSession("model.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: img_data}
ort_outs = ort_session.run(None, ort_inputs)
import matplotlib.pyplot as plt
import matplotlib.patches as patches
def display_objdetect_image(image, boxes, labels, scores, score_threshold=0.7):
# Resize boxes
ratio = 800.0 / min(image.size[0], image.size[1])
boxes /= ratio
_, ax = plt.subplots(1, figsize=(12,9))
image = np.array(image)
ax.imshow(image)
# Showing boxes with score > 0.7
for box, label, score in zip(boxes, labels, scores):
if score > score_threshold:
rect = patches.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], linewidth=1, edgecolor='b', facecolor='none')
ax.annotate(classes[label] + ':' + str(np.round(score, 2)), (box[0], box[1]), color='w', fontsize=12)
ax.add_patch(rect)
plt.show()
display_objdetect_image(img, ort_outs[0], ort_outs[1], ort_outs[2])