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