Test client for Tensorflow Serving running Darkflow

test_client.py

#!/usr/bin/env python2.7
# Based on https://github.com/tensorflow/serving/blob/master/tensorflow_serving/example/inception_client.py

"""Send JPEG image to tensorflow_model_server loaded with Darkflow model.
"""

from __future__ import print_function

from grpc.beta import implementations
import tensorflow as tf

from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2

import numpy as np
import sys
import cv2

tf.app.flags.DEFINE_string('server', 'localhost:11000',
                           'PredictionService host:port')
tf.app.flags.DEFINE_string('image', sys.argv[1], sys.argv[1])
FLAGS = tf.app.flags.FLAGS

def main(_):
    host, port = FLAGS.server.split(':')
    channel = implementations.insecure_channel(host, int(port))
    stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)

    data = cv2.imread(sys.argv[1])
    data = cv2.resize(data, (1024, 1024))
    data = data.astype(np.float32)
    data = tf.contrib.util.make_tensor_proto(data, shape=[1, 1024, 1024, 3])

    request = predict_pb2.PredictRequest()
    request.model_spec.name = 'my-model' # Export directory of SavedModel
    request.model_spec.signature_name = "predict"

    request.inputs['input'].CopyFrom(data)
    result = stub.Predict(request, 10.0)  # 10 secs timeout
    print(result)

if __name__ == '__main__':
  tf.app.run()