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]