Need to apply post processing with onnx model #5854

kino0924 · 2023-05-25T12:26:31Z

kino0924
May 25, 2023

Hi, I recently deployed onnx model in my homelab

I tested triton + python(grpc) and its working well.
However, I want to put some extra post processing after inference but not sure its possible.

Here is my config and poc script.

config.pbtxt

name: "poc_onnx"
platform: "onnxruntime_onnx"
max_batch_size : 0
input [
  {
    name: "input"
    data_type: TYPE_FP32
    dims: [ 5, 3, 34, 260 ]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [ 5, 20, 11172 ]
  }
]

poc.py

def test_string():
    images = []
    for i in range(5):
        image = Image.open(f"./app/image_{i+1}.png")
        image = np.array(image)
        image = image.astype(np.float32) / 255.0
        image = (image - 0.5) / 0.5
        image = np.transpose(image, (2, 0, 1))
        images.append(image)

    client = grpcclient.InferenceServerClient(url="localhost:8001", verbose=False)

    inputs = []
    outputs = []

    input_data = np.stack(images, axis=0)

    input_name = "input"
    input_shape = input_data.shape
    input_dtype = "FP32"
    inputs.append(grpcclient.InferInput(input_name, input_shape, input_dtype))
    inputs[0].set_data_from_numpy(input_data)

    output_name = "output"
    outputs.append(grpcclient.InferRequestedOutput(output_name))

    response = client.infer(model_name="poc_onnx", inputs=inputs, outputs=outputs)

    output_data = response.as_numpy(output_name)
    selected_indices = np.argmax(output_data, axis=-1)

    final_strs = []
    for i in range(5):
        text = ""
        for idx in selected_indices[i]:
            text += char_11172[idx]
        final_strs.append(text)

    return {"result": final_strs}