File size: 1,648 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 numpy as np
import triton_python_backend_utils as pb_utils # type: ignore
class TritonPythonModel:
def initialize(self, args):
_ = args
def execute(self, requests):
responses = []
for request in requests:
bbox_2d = pb_utils.get_input_tensor_by_name(request, "bbox_2d").as_numpy().astype(np.float32, copy=False)
function_2d = pb_utils.get_input_tensor_by_name(
request, "function_embedding_2d"
).as_numpy().astype(np.float32, copy=False)
vision_2d = pb_utils.get_input_tensor_by_name(
request, "vision_embedding_2d"
).as_numpy().astype(np.float32, copy=False)
text_2d = pb_utils.get_input_tensor_by_name(
request, "text_embedding_2d"
).as_numpy().astype(np.float32, copy=False)
n = int(bbox_2d.shape[0]) if bbox_2d.ndim >= 1 else 0
padding_mask = np.ones((1, n), dtype=np.int32)
responses.append(
pb_utils.InferenceResponse(
output_tensors=[
pb_utils.Tensor("bbox", np.expand_dims(bbox_2d, axis=0)),
pb_utils.Tensor("function_embedding", np.expand_dims(function_2d, axis=0)),
pb_utils.Tensor("vision_embedding", np.expand_dims(vision_2d, axis=0)),
pb_utils.Tensor("text_embedding", np.expand_dims(text_2d, axis=0)),
pb_utils.Tensor("padding_mask", padding_mask),
]
)
)
return responses
def finalize(self):
pass
|