| import numpy as np |
| import triton_python_backend_utils as pb_utils |
|
|
|
|
| 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 |
|
|
