File size: 1,587 Bytes
948d40c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import base64
import time
import sys
import numpy as np
import triton_python_backend_utils as pb_utils # type: ignore
if "/" not in sys.path:
sys.path.insert(0, "/")
from utils import decode_image, od_resize_image
class TritonPythonModel:
def execute(self, requests):
logger = pb_utils.Logger
responses = []
st = time.time()
for request in requests:
image_bytes = pb_utils.get_input_tensor_by_name(request, "image_bytes").as_numpy()
image_bytes = base64.b64decode(image_bytes[0].decode('utf-8'))
image = decode_image(image_bytes)
(
resized_image,
(original_width, original_height),
(resized_width, resized_height),
) = od_resize_image(image, (960, 960))
resized_image = resized_image[:, :, ::-1].transpose(2, 0, 1)
resized_image = np.ascontiguousarray(resized_image).astype(np.float32) / 255.0
resized_image = np.expand_dims(resized_image, axis=0)
out_tensor_0 = pb_utils.Tensor("resized_image", np.array(resized_image))
out_tensor_1 = pb_utils.Tensor("original_shape", np.array([original_width, original_height]).astype(np.uint16))
out_tensor_2 = pb_utils.Tensor("resized_shape", np.array([960, 960]).astype(np.uint16))
responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0, out_tensor_1, out_tensor_2]))
logger.log_info(f"OD Preprocess execute duration : {int((time.time() - st)*1000)} ms")
return responses
|