| 28 | |
| 29 | # Controller |
| 30 | class SAMImageEncoder: |
| 31 | def __init__(self, triton_client: grpcclient.InferenceServerClient) -> None: |
| 32 | logger.info("Initialize SAMImageEncoder.") |
| 33 | self.triton_client = triton_client |
| 34 | |
| 35 | @torch.no_grad() |
| 36 | async def run( |
| 37 | self, file: UploadFile, inference_params: Dict[str, Any] |
| 38 | ) -> Dict[str, Any]: |
| 39 | logger.info("Run SAMImageEncoder.") |
| 40 | image = await file.read() |
| 41 | |
| 42 | # Preprocess |
| 43 | input_image = self.preprocess(image) |
| 44 | |
| 45 | # Prepare for inference. |
| 46 | triton_inputs = [grpcclient.InferInput("INPUT__0", input_image.shape, "FP32")] |
| 47 | triton_inputs[0].set_data_from_numpy(input_image) |
| 48 | triton_outputs = [grpcclient.InferRequestedOutput("OUTPUT__0")] |
| 49 | |
| 50 | # Run the inference. |
| 51 | result = await self.triton_client.infer( |
| 52 | inputs=triton_inputs, |
| 53 | outputs=triton_outputs, |
| 54 | **inference_params, |
| 55 | ) |
| 56 | # image embedding: numpy.ndarray [B, 256, 64, 64] |
| 57 | image_embedding = result.as_numpy("OUTPUT__0") |
| 58 | |
| 59 | # Postprocess |
| 60 | outputs = self.postprocess(image_embedding) |
| 61 | |
| 62 | return outputs |
| 63 | |
| 64 | def preprocess(self, image_byte, target_size=1024) -> torch.Tensor: |
| 65 | """ |
| 66 | Preprocess image to input to encoder. |
| 67 | |
| 68 | Return: |
| 69 | preprocessed: If longest size of target image is 1024, |
| 70 | shape of tensor is [B, 3, 1024, 1024]. |
| 71 | And dtype of tensor is float32. |
| 72 | """ |
| 73 | # Convert the bytes to numpy array |
| 74 | image = np.frombuffer(image_byte, dtype=np.uint8) |
| 75 | image = cv2.imdecode(image, cv2.IMREAD_COLOR)[:, :, ::-1] # RGB |
| 76 | |
| 77 | # Get image shape and convert type |
| 78 | if image.shape != (1024, 1024, 3): |
| 79 | origin_shape = image.shape[:2] |
| 80 | height, width = self.get_preprocess_shape(*origin_shape) |
| 81 | image = cv2.resize(image, dsize=(width, height)) |
| 82 | height, width = image.shape[:2] |
| 83 | image_fp = image.astype(np.float32) |
| 84 | |
| 85 | # Normalize |
| 86 | image_fp -= np.array([123.675, 116.28, 103.53], dtype=np.float32) # mean |
| 87 | image_fp /= np.array([58.395, 57.12, 57.375], dtype=np.float32) # std |