efficientdet_tf_ovino.py

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()