| import time |
| import sys |
|
|
| import numpy as np |
| import torch |
| import triton_python_backend_utils as pb_utils |
| if "/" not in sys.path: |
| sys.path.insert(0, "/") |
| from utils import non_max_suppression |
|
|
|
|
| class TritonPythonModel: |
| def initialize(self, args): |
| self.conf_thres = 0.31 |
| self.iou_thres = 0.55 |
|
|
| def execute(self, requests): |
| logger = pb_utils.Logger |
|
|
| responses = [] |
| st = time.time() |
| for request in requests: |
| output = pb_utils.get_input_tensor_by_name(request, "model_output").as_numpy() |
|
|
| output_tensor = torch.tensor(output).reshape(1, 22, 18900) |
| output_list = non_max_suppression( |
| output_tensor, |
| conf_thres=self.conf_thres, |
| iou_thres=self.iou_thres, |
| ) |
| detections = output_list[0] |
| if len(detections) == 0: |
| output_numpy = np.empty((1, 0, 4), dtype=np.float32) |
| else: |
| boxes = detections[:, :4].detach().cpu().numpy().astype(np.float32) |
| output_numpy = np.expand_dims(boxes, axis=0) |
|
|
| out_tensor_0 = pb_utils.Tensor("detections", output_numpy) |
|
|
| responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0])) |
|
|
| logger.log_info(f"od.widget.postprocess execute duration : {int((time.time() - st)*1000)} ms") |
| return responses |
|
|
