def convert(model_name, export_dir):
builder = tf.saved_model.builder.SavedModelBuilder(export_dir=export_dir)
with tf.Graph().as_default(), tf.Session().as_default() as sess:
tensors = freeze.unfreeze_into_current_graph(
os.path.join(paths.FROZEN_MODELS_DIR, model_name + '.pb'),
tensor_names=[
consts.INCEPTION_INPUT_TENSOR, consts.OUTPUT_TENSOR_NAME
])
raw_image_proto_info = tf.saved_model.utils.build_tensor_info(
tensors[consts.INCEPTION_INPUT_TENSOR])
probs_proto_info = tf.saved_model.utils.build_tensor_info(
tensors[consts.OUTPUT_TENSOR_NAME])
prediction_signature = (
tf.saved_model.signature_def_utils.build_signature_def(
inputs={'image_raw': raw_image_proto_info},
outputs={'probs': probs_proto_info},
method_name=tf.saved_model.signature_constants.
PREDICT_METHOD_NAME))
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
prediction_signature
})
builder.save()
def inception_model():
tensors = freeze.unfreeze_into_current_graph(
paths.IMAGENET_GRAPH_DEF,
tensor_names=[
consts.INCEPTION_INPUT_TENSOR, consts.INCEPTION_OUTPUT_TENSOR
])
def forward(sess, image_raw):
out = sess.run(tensors[consts.INCEPTION_OUTPUT_TENSOR],
{tensors[consts.INCEPTION_INPUT_TENSOR]: image_raw})
return out
def infer(model_name, img_raw):
with tf.Graph().as_default(), tf.Session().as_default() as sess:
tensors = freeze.unfreeze_into_current_graph(
os.path.join(paths.FROZEN_MODELS_DIR, model_name + '.pb'),
tensor_names=[consts.INCEPTION_INPUT_STRING_TENSOR, consts.OUTPUT_TENSOR_NAME])
_, one_hot_decoder = dataset.one_hot_label_encoder()
probs = sess.run(tensors[consts.OUTPUT_TENSOR_NAME],
feed_dict={tensors[consts.INCEPTION_INPUT_STRING_TENSOR]: img_raw})
breeds = one_hot_decoder(np.identity(consts.CLASSES_COUNT)).reshape(-1)
# print(breeds)
df = pd.DataFrame(data={'prob': probs.reshape(-1), 'breed': breeds})
return df.sort_values(['prob'], ascending=False)