Skip to content

Instantly share code, notes, and snippets.

@dkurt
Created July 29, 2020 12:20
Show Gist options
  • Save dkurt/8ea9bb1c9fbb001bbfc644244ba4611d to your computer and use it in GitHub Desktop.
Save dkurt/8ea9bb1c9fbb001bbfc644244ba4611d to your computer and use it in GitHub Desktop.
import numpy as np
import cv2 as cv
import tensorflow as tf
img = cv.imread('img1.jpg')
img = np.ascontiguousarray(img[:,:,[2, 1, 0]])
resized_img = cv.resize(img, (640, 640))
#
# Run TensorFlow
#
pb_file = 'efficientdet-d1_opt.pb'
graph_def = tf.compat.v1.GraphDef()
try:
with tf.io.gfile.GFile(pb_file, 'rb') as f:
graph_def.ParseFromString(f.read())
except:
with tf.gfile.FastGFile(pb_file, 'rb') as f:
graph_def.ParseFromString(f.read())
with tf.Session() as sess:
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
out = sess.run(sess.graph.get_tensor_by_name('detections:0'),
feed_dict={'image_arrays:0': resized_img.reshape(1, 640, 640, 3)})
#
# Run OpenVINO
#
from openvino.inference_engine import IECore, IENetwork
inp = resized_img.astype(np.float32)
inp[:,:,0] -= 123.675
inp[:,:,1] -= 116.28
inp[:,:,2] -= 103.53
inp /= 255
inp = np.ascontiguousarray(inp.transpose(2, 0, 1).reshape(1, 3, 640, 640))
ie = IECore()
net = ie.read_network('efficientdet-d1.xml', 'efficientdet-d1.bin')
exec_net = ie.load_network(net, 'CPU')
ie_out = exec_net.infer({'image_arrays': inp})
ie_out = next(iter(ie_out.values()))
print('=== TensorFlow ====================================')
for detection in out.reshape(-1, 7):
ymin, xmin, ymax, xmax = [v for v in detection[1:5]]
ymin = int(img.shape[0] * ymin / 640)
ymax = int(img.shape[0] * ymax / 640)
xmin = int(img.shape[1] * xmin / 640)
xmax = int(img.shape[1] * xmax / 640)
print('class:', detection[-1], 'confidence:', detection[-2], 'box:', ymin, xmin, ymax, xmax)
cv.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), thickness=3)
print('=== OpenVINO ======================================')
for detection in ie_out.reshape(-1, 7):
conf = detection[2]
if conf < 0.5:
continue
classId = detection[1] + 1
xmin = int(img.shape[1] * detection[3])
ymin = int(img.shape[0] * detection[4])
xmax = int(img.shape[1] * detection[5])
ymax = int(img.shape[0] * detection[6])
print('class:', classId, 'confidence:', conf, 'box:', ymin, xmin, ymax, xmax)
cv.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255))
cv.imwrite('res.png', img)
cv.waitKey()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment