File size: 2,703 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import base64
import cv2
import numpy as np
import triton_python_backend_utils as pb_utils # type: ignore
class TritonPythonModel:
def initialize(self, args):
self.td_canvas_min_size = 640
self.td_canvas_max_size = 2240
def _preprocess(self, image_bytes: bytes):
image_array = np.frombuffer(image_bytes, dtype=np.uint8)
original_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR_BGR)
if original_image.shape[2] == 4:
original_image = original_image[:, :, :3]
h, w = original_image.shape[:2]
if min(h, w) < 640:
max_ratio = 1.0
else:
max_ratio = 1.5
ratio = min(self.td_canvas_max_size / max(h,w), max_ratio)
if ratio != 1.0:
ratio = float(int(ratio * 32)) / 32.0
resized_image = cv2.resize(original_image, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_CUBIC)
else:
resized_image = original_image
h, w = resized_image.shape[:2]
if h < self.td_canvas_min_size:
h_pad = self.td_canvas_min_size - h
else:
h_pad = min(64 - h%64, self.td_canvas_max_size - h)
if w < self.td_canvas_min_size:
w_pad = self.td_canvas_min_size - w
else:
w_pad = min(64 - w%64, self.td_canvas_max_size - w)
resized_image = np.pad(resized_image, ((0, (h_pad)), (0, w_pad), (0, 0)), 'constant', constant_values=0)
resized_image = resized_image.transpose(2, 0, 1)
resized_image = np.expand_dims(resized_image, axis=0)
resized_image = np.ascontiguousarray(resized_image)
out_tensor_0 = pb_utils.Tensor("resized_image", resized_image)
out_tensor_1 = pb_utils.Tensor("resize_ratio", np.array([ratio], dtype=np.float32))
out_tensor_2 = pb_utils.Tensor("original_image", original_image)
return out_tensor_0, out_tensor_1, out_tensor_2
def execute(self, requests):
logger = pb_utils.Logger
responses = []
for request in requests:
try:
input_tensor = pb_utils.get_input_tensor_by_name(request, "image_bytes").as_numpy()
image_bytes = base64.b64decode(input_tensor[0].decode('utf-8'))
out_tensor_0, out_tensor_1, out_tensor_2 = self._preprocess(image_bytes)
except pb_utils.TritonModelException as e:
responses.append(pb_utils.InferenceResponse(error=pb_utils.TritonError(str(e), pb_utils.TritonError.BAD_REQUEST)))
continue
responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0, out_tensor_1, out_tensor_2]))
return responses
|