"""

Basic components for using TF Serving.

"""

import os

os.environ["CUDA_VISIBLE_DEVICES"] = ""

import io

import numpy as np

from PIL import Image

import grpc

import tensorflow as tf

from grpc.beta import implementations

from tensorflow_serving.apis import predict_pb2

from tensorflow_serving.apis import prediction_service_pb2

from grpc._cython import cygrpc

import config

def decode_tensor(contents, downsample_factor):

    images = [Image.open(io.BytesIO(content)) for content in contents]

    dsize = (images[0].size[0] * downsample_factor, images[0].size[1] * downsample_factor)

    images = [image.resize(dsize, Image.BILINEAR) for image in images]

    mtx = np.array([np.asarray(image, dtype=np.uint8) for image in images], dtype=np.uint8)

    mtx = mtx.transpose((1, 2, 0))

    return mtx

class TFServing(object):

    """

    Tensorflow Serving client, send prediction request.

    """

    def __init__(self, host, port):

        super(TFServing, self).__init__()

        channel = self._insecure_channel(host, port)

        self._stub = prediction_service_pb2.beta_create_PredictionService_stub(channel)

    def _insecure_channel(self, host, port):

        channel = grpc.insecure_channel(

            target=host if port is None else '{}:{}'.format(host, port),

            options=[(cygrpc.ChannelArgKey.max_send_message_length, -1),

                     (cygrpc.ChannelArgKey.max_receive_message_length, -1)])

        return grpc.beta.implementations.Channel(channel)

    def predict(self, image_input, model_name):

        request = predict_pb2.PredictRequest()

        request.model_spec.name = model_name

        request.inputs[config.INPUT_KEY].CopyFrom(tf.contrib.util.make_tensor_proto(

            image_input.astype(np.uint8, copy=False)))

        try:

            result = self._stub.Predict(request, 1000.0)

            image_prob = np.array(result.outputs[config.PREDICT_KEY].int_val)

        except Exception as e:

            raise e

        else:

            return image_prob.astype(np.uint8)