You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
import grpc
|
|
import inference_pb2
|
|
import inference_pb2_grpc
|
|
import management_pb2
|
|
import management_pb2_grpc
|
|
import json
|
|
|
|
|
|
def get_inference_stub():
|
|
channel = grpc.insecure_channel('localhost:7070')
|
|
stub = inference_pb2_grpc.InferenceAPIsServiceStub(channel)
|
|
return stub
|
|
|
|
|
|
def get_management_stub():
|
|
channel = grpc.insecure_channel('localhost:7071')
|
|
stub = management_pb2_grpc.ManagementAPIsServiceStub(channel)
|
|
return stub
|
|
|
|
|
|
def list_models(management_stub):
|
|
list_model_request_object = management_pb2.ListModelsRequest(limit=10)
|
|
return management_stub.ListModels(list_model_request_object)
|
|
|
|
|
|
def make_prediction(inference_stub, image_path):
|
|
with open(image_path, "rb") as f:
|
|
image_bytes = f.read()
|
|
input_data = {"data": image_bytes}
|
|
prediction_request = inference_pb2.PredictionsRequest(model_name="digitmodel", input=input_data)
|
|
inference_prediction = inference_stub.Predictions(prediction_request)
|
|
return inference_prediction
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
inference_stub = get_inference_stub()
|
|
management_stub = get_management_stub()
|
|
|
|
output = list_models(management_stub)
|
|
print(f"output: {output}")
|
|
|
|
message = json.loads(output.msg)
|
|
print(f"message: {message}")
|
|
|
|
prediction = make_prediction(inference_stub, "torchserve_grpc/images/test.png")
|
|
|
|
print(prediction)
|