File size: 1,112 Bytes
279d47e |
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 |
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]
|