| import tensorflow as tf | |
| from PIL import Image | |
| import io | |
| imported = tf.saved_model.load("./app") | |
| imported = imported.signatures["serving_default"] | |
| def get_image_from_bytes(binary_image: bytes) -> Image: | |
| """Convert image from bytes to PIL RGB format | |
| Args: | |
| binary_image (bytes): The binary representation of the image | |
| Returns: | |
| PIL.Image: The image in PIL RGB format | |
| """ | |
| input_image = Image.open(io.BytesIO(binary_image)).convert("RGB") | |
| return input_image | |
| def predict(input_image): | |
| """Reads file and returns prediction | |
| Args: | |
| x (_type_): _description_ | |
| Returns: | |
| _type_: _description_ | |
| """ | |
| tensor = tf.io.decode_image(input_image, channels=3) | |
| inference_shape = (240, 320) | |
| original_shape = tensor.shape[:2] | |
| input_tensor = tf.expand_dims(tensor, axis=0) | |
| input_tensor = tf.image.resize(input_tensor, inference_shape, | |
| preserve_aspect_ratio=True) | |
| saliency = imported(input_tensor)["output"] | |
| saliency = tf.image.resize(saliency, original_shape) | |
| return saliency.numpy()[0] | |