File size: 1,418 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 42 43 44 | import time
import sys
import numpy as np
import torch
import triton_python_backend_utils as pb_utils # type: ignore
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
|