| import time |
|
|
| import cv2 |
| import numpy as np |
| import triton_python_backend_utils as pb_utils |
|
|
|
|
| class TritonPythonModel: |
| @staticmethod |
| def _to_bytes(value) -> bytes: |
| if isinstance(value, np.bytes_): |
| return value.tobytes() |
| if isinstance(value, bytes): |
| return value |
| if isinstance(value, bytearray): |
| return bytes(value) |
| if isinstance(value, str): |
| return value.encode("utf-8") |
| return bytes(value) |
|
|
| def execute(self, requests): |
| logger = pb_utils.Logger |
| responses = [] |
| st = time.time() |
|
|
| for request in requests: |
| image_bytes_tensor = pb_utils.get_input_tensor_by_name(request, "IMAGE_BYTES_LIST") |
| image_bytes_list = image_bytes_tensor.as_numpy().reshape(-1) |
|
|
| images = [] |
| for idx, image_bytes in enumerate(image_bytes_list): |
| raw_bytes = self._to_bytes(image_bytes) |
| if len(raw_bytes) == 0: |
| raise pb_utils.TritonModelException(f"Empty image bytes at index {idx}.") |
|
|
| image_np = np.frombuffer(raw_bytes, np.uint8) |
| image = cv2.imdecode(image_np, cv2.IMREAD_COLOR) |
| if image is None: |
| raise pb_utils.TritonModelException( |
| f"Failed to decode image bytes at index {idx}." |
| ) |
|
|
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
| image = cv2.resize(image, (224, 224), interpolation=cv2.INTER_AREA) |
|
|
| x = image.astype(np.float32) / 255.0 |
| x = (x - 0.5) / 0.5 |
| x = np.transpose(x, (2, 0, 1)) |
| images.append(x) |
|
|
| if len(images) == 0: |
| pixel_values = np.empty((0, 3, 224, 224), dtype=np.float32) |
| else: |
| pixel_values = np.stack(images, axis=0).astype(np.float32) |
|
|
| responses.append( |
| pb_utils.InferenceResponse( |
| output_tensors=[pb_utils.Tensor("pixel_values", pixel_values)] |
| ) |
| ) |
|
|
| logger.log_info(f"siglip_preprocess execute duration : {int((time.time() - st) * 1000)} ms") |
| return responses |
|
|
