import tensorflow as tf
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model.utils import build_tensor_info

model_path = './frozen/inception_v3_2016_08_28_frozen.pb'
target_dir = './models/inception/3'

# (1)
with tf.gfile.FastGFile(model_path, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    _ = tf.import_graph_def(graph_def, name='')

input_name = 'input'
output_name = 'InceptionV3/Predictions/Reshape_1'

with tf.Session() as sess:
    # (2)
    model_input = build_tensor_info(sess.graph.get_tensor_by_name(input_name + ':0'))
    model_output = build_tensor_info(sess.graph.get_tensor_by_name(output_name + ':0'))

    # (3)
    signature_definition = signature_def_utils.build_signature_def(
        inputs={input_name: model_input},
        outputs={output_name: model_output},
        method_name=signature_constants.PREDICT_METHOD_NAME)

    # (4)
    builder = saved_model_builder.SavedModelBuilder(target_dir)
    builder.add_meta_graph_and_variables(sess, [tag_constants.SERVING], signature_def_map={
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_definition
    }, clear_devices=True)
    builder.save()