#!/home/travis/build/triagemd/tensorflow-serving-client/.venv/bin/python
import argparse
import json

from tensorflow_serving_client import TensorflowServingClient
from keras_model_specs import ModelSpec

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--host', required=True, type=str, help='Hostname to query')
    parser.add_argument('--port', required=True, type=int, help='Port to query')
    parser.add_argument('--model_spec', required=True, type=str, help='Name of the model spec')
    parser.add_argument('--image', required=True, type=str, help='Image to send (JPG format)')
    args = parser.parse_args()

    model_spec = ModelSpec.get(args.model_spec)

    client = TensorflowServingClient(args.host, args.port)
    image_data = model_spec.load_image(args.image)
    response = client.make_prediction(image_data, 'image')
    predictions = response['class_probabilities'][0].tolist()

    print(json.dumps(predictions, indent=4, sort_keys=True))
