Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import tensorflow as tf
- import cv2
- import os
- def get_frozen_graph(graph_file):
- """Read Frozen Graph file from disk."""
- with tf.gfile.FastGFile(graph_file, "rb") as f:
- graph_def = tf.GraphDef()
- graph_def.ParseFromString(f.read())
- return graph_def
- # The TensorRT inference graph file downloaded from Colab or your local machine.
- pb_fname = os.path.join(os.getcwd(), "faster_rcnn_inception_resnet_v2_atrous_coco_2018_01_28", "frozen_inference_graph.pb")
- trt_graph = get_frozen_graph(pb_fname)
- input_names = ['image_tensor']
- # Create session and load graph
- tf_config = tf.ConfigProto()
- tf_config.gpu_options.allow_growth = True
- tf_sess = tf.Session(config=tf_config)
- tf.import_graph_def(trt_graph, name='')
- tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ':0')
- tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')
- tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')
- tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')
- tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')
- IMAGE_PATH = os.path.join(os.getcwd(), "testimages", "000002_491724089556.png")
- image = cv2.imread(IMAGE_PATH)
- image = cv2.resize(image, (300, 300))
- scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={
- tf_input: image[None, ...]
- })
- boxes = boxes[0] # index by 0 to remove batch dimension
- scores = scores[0]
- classes = classes[0]
- num_detections = int(num_detections[0])
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement