from tensorflow.python.saved_model import signature_constants

from tensorflow.python.saved_model import tag_constants

import tensorflow as tf

import os

export_dir = r'E:\files\tensorflow-yolov3\savemodel_pb\yolo3_hot_model/000001'

graph_pb = r'E:\files\tensorflow-yolov3\savemodel_pb\yolov3_base30.pb'

if not os.path.exists(export_dir):

os.makedirs(export_dir)

builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

with tf.gfile.GFile(graph_pb, "rb") as f:

graph_def = tf.GraphDef()

graph_def.ParseFromString(f.read())

sigs = {}

with tf.Session(graph=tf.Graph()) as sess: # name="" is important to ensure we don't get spurious prefixing

tf.import_graph_def(graph_def, name="")

g = tf.get_default_graph()

image_tensor = g.get_tensor_by_name("input/input_data:0")

detection_sbbox = g.get_tensor_by_name("pred_sbbox/concat_2:0")

detection_mbbox = g.get_tensor_by_name("pred_mbbox/concat_2:0")

detection_lbbox = g.get_tensor_by_name("pred_lbbox/concat_2:0")

print(type(detection_sbbox))

out = {'sbbox':detection_sbbox,'mbbox':detection_mbbox,'lbbox':detection_lbbox}

print(type(out))

sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \

tf.saved_model.signature_def_utils.predict_signature_def(

{"in": image_tensor},

{'sbbox':detection_sbbox,'mbbox':detection_mbbox,'lbbox':detection_lbbox},)

# {"scores":detection_scores},

# {"classes":detection_classes},

# {"nums":num_detections})

builder.add_meta_graph_and_variables(sess,

[tag_constants.SERVING],

signature_def_map = sigs)

builder.save()