-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathoverlay.py
118 lines (96 loc) · 3.87 KB
/
overlay.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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sun Dec 3 15:39:18 2017
@author: atabak
"""
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import cv2
import os
import numpy as np
import tensorflow as tf
from PIL import Image
from matplotlib import pyplot as plt
from object_detection.utils import label_map_util, visualization_utils
tf.logging.set_verbosity(tf.logging.INFO)
flags = tf.app.flags
flags.DEFINE_string('images_path', '', 'path to resized test images folder')
flags.DEFINE_string('ckpt_path', '', 'path to network frozen graph')
flags.DEFINE_string('save_path', '', 'path to save infered bounding boxes')
FLAGS = flags.FLAGS
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 main(_):
path_to_ckpt = FLAGS.ckpt_path
path_to_labels = 'config/label_map.pbtxt'
num_classes = 3
path_to_test_images_dir = FLAGS.images_path
test_images_path = os.listdir(path_to_test_images_dir)
img_size = (12,12)
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(path_to_ckpt, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name('num_detections:0')
label_map = label_map_util.load_labelmap(path_to_labels)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=num_classes, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
with tf.Session(graph=detection_graph) as sess:
for image_path in test_images_path:
image = Image.open(os.path.join(path_to_test_images_dir,image_path))
image_np = load_image_into_numpy_array(image)
image_np_expanded = np.expand_dims(image_np, axis=0)
(boxes, scores, classes, num) = sess.run(
[detection_boxes, detection_scores, detection_classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
width=960
height=540
#NEW
font = cv2.FONT_HERSHEY_DUPLEX
fontsize=1
fontweight=1
#img = cv2.imread(full_path,cv2.IMREAD_COLOR)
classes_text= np.squeeze(classes).astype(np.int32)
#print(classes_text)
for i in xrange(0,len(np.squeeze(boxes))):
bbox=np.squeeze(boxes)[i]
xmin=bbox[1]
ymin=bbox[0]
xmax=bbox[3]
ymax=bbox[2]
score=str("%.2f"%np.squeeze(scores)[i])
if(np.squeeze(scores)[i]<=0.75) :
continue
if classes_text[i]==1:
class_='cylinder'
color=(5,15,202)
elif classes_text[i]==2:
class_='sphere'
color=(0,128,0)
elif classes_text[i]==3:
class_='box'
color=(255,0,0)
cv2.rectangle(image_np,(int(xmin*width),int(ymin*height)),
(int(xmax*width),int(ymax*height)),color,2)
cv2.putText(image_np, class_+"("+score+")", (int(xmin*width), int(ymin*height-8)),
font, fontsize, color, fontweight, cv2.LINE_AA)
#cv2.imshow('image',image_np)
#cv2.waitKey(0)
#cv2.destroyAllWindows()
cv2.imwrite(os.path.join(FLAGS.save_path,image_path).replace('test','new_test'),
cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR))
#ENDNEW
#plt.savefig(os.path.join(FLAGS.save_path,image_path), format='png', bbox_inches='tight')
if __name__ == '__main__':
tf.app.run()