-
Notifications
You must be signed in to change notification settings - Fork 2
/
predict.py
103 lines (84 loc) · 4.03 KB
/
predict.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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import numpy as np
import tensorflow as tf
import base64
from distutils.version import StrictVersion
from PIL import Image
import io
if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):
raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')
# What model to download.
MODEL_NAME = 'output'
# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
def load_image_into_numpy_array(image):
(im_width, im_height) = image.size
return np.array(image.getdata()).reshape(
(im_height, im_width, 3)).astype(np.uint8)
def run_inference_for_single_image(image, graph):
with graph.as_default():
with tf.Session() as sess:
# Get handles to input and output tensors
ops = tf.get_default_graph().get_operations()
all_tensor_names = {output.name for op in ops for output in op.outputs}
tensor_dict = {}
for key in [
'num_detections', 'detection_boxes', 'detection_scores',
'detection_classes'
]:
tensor_name = key + ':0'
if tensor_name in all_tensor_names:
tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
tensor_name)
image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
# Run inference
output_dict = sess.run(tensor_dict,
feed_dict={image_tensor: np.expand_dims(image, 0)})
# all outputs are float32 numpy arrays, so convert types as appropriate
output_dict['num_detections'] = int(output_dict['num_detections'][0])
output_dict['detection_classes'] = output_dict[
'detection_classes'][0].astype(np.uint8)
output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
output_dict['detection_scores'] = output_dict['detection_scores'][0]
return output_dict
def get_bounding_box_on_image(image,
ymin,
xmin,
ymax,
xmax,
use_normalized_coordinates=True):
im_width, im_height = image.size
if use_normalized_coordinates:
(left, right, top, bottom) = (xmin * im_width, xmax * im_width,
ymin * im_height, ymax * im_height)
else:
(left, right, top, bottom) = (xmin, xmax, ymin, ymax)
return (left, right, top, bottom)
# Take in base64 string and return PIL image
def stringToImage(base64_string):
imgdata = base64.b64decode(base64_string)
return Image.open(io.BytesIO(imgdata))
def image64_to_image(image64):
ImageData = stringToImage(image64)
return ImageData
def detect_image_64(image, probability=0.1):
image = image64_to_image(image)
image_np = load_image_into_numpy_array(image)
output_dict = run_inference_for_single_image(image_np, detection_graph)
boxes = []
for i in range(output_dict['num_detections']):
if output_dict['detection_scores'][i] > probability:
box_1 = get_bounding_box_on_image(image, output_dict['detection_boxes'][i][0],
output_dict['detection_boxes'][i][1],
output_dict['detection_boxes'][i][2],
output_dict['detection_boxes'][i][3])
first_box_class = output_dict['detection_classes'][i]
first_box_score = output_dict['detection_scores'][i]
boxes.append([box_1, first_box_class, first_box_score])
return boxes