tfserving备忘

import grpc
import numpy as np
import tensorflow as tf
from PIL import Image
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
import cv2
import matplotlib.pyplot as plt
 
def request_tfserving(inputs, server_url, model_name, signature_name, input_names, output_names):
    # 建立连接
    channel = grpc.insecure_channel(server_url)
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    # 开始设置请求数据
    request = predict_pb2.PredictRequest()
    request.model_spec.name = model_name  # 模型名称
    request.model_spec.signature_name = signature_name  # 签名名称（默认 serving_default）
    # 设置输入数据
    for input, input_name in zip(inputs, input_names):
        request.inputs[input_name].CopyFrom(tf.make_tensor_proto(input, shape=list(input.shape)))
        
    response = stub.Predict(request, 5.0)  # 其中第2个参数为请求的 timeout 时长
    res_from_server_np = []
    for output_name in output_names:
        res_from_server_np.append(tf.make_ndarray(response.outputs[output_name]))
    print(np.argmax(res_from_server_np[0][0,:],axis=2).shape)
    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 8))

    ax1.imshow(input[0,:].astype(np.uint8))
    ax2.imshow(np.argmax(res_from_server_np[0][0,:],axis=2))
    plt.show()
 
if __name__ == "__main__":
    img_path = ‘/home/ground_train/test/20112201814-1_00000036_img.png‘
    img = Image.open(img_path)
    img = img.resize((960, 544), Image.NEAREST).convert(‘RGB‘)
    img = image.img_to_array(img)
    img = np.expand_dims(img, axis=0).astype(‘float32‘)
    print(img.shape)
    request_tfserving(inputs=[img],#这里一定得是个list
                      server_url=‘ip:port‘,
                      model_name=‘ground_segmentation‘,
                      signature_name=‘serving_default‘,
                      input_names=[‘input‘],
                      output_names=[‘output‘])