from PIL import Image
from io import BytesIO
import base64
import tensorflow as tf


IMAGE_HEIGHT = 28
IMAGE_WIDTH = 28
LABELS = ["੦", "੧", "੨", "੩", "੪", "੫", "੬", "੭", "੮", "੯"]
DEFAULT_TOP_K = 3


class EndpointHandler():
    def __init__(self, path=""):
        self.model = tf.keras.models.load_model(path)

    def __call__(self, data):
        image = Image.open(BytesIO(base64.b64decode(data.pop("inputs").pop("image"))))
        tensors = self.to_tensors(image)
        predictions = self.model.predict(tensors)
        top_k_scores, top_k_label_ids = tf.nn.top_k(predictions, k=data.pop("parameters", {}).pop("topK", DEFAULT_TOP_K))

        return [
            {
                "label": LABELS[(int(label_id))],
                "score": float(score),
            }
            for label_id, score in zip(top_k_label_ids[0], top_k_scores[0])
        ]

    @staticmethod    
    def to_tensors(image):
        if image.mode == "RGBA":
            img = Image.new("RGB", image.size, (255, 255, 255))
            img.paste(image, mask=image.split()[3])
            image = img
        elif image.mode != "RGB":
            image = image.convert("RGB")

        image = tf.image.resize(image, size=(IMAGE_HEIGHT, IMAGE_WIDTH), antialias=True)
        image = tf.image.rgb_to_grayscale(image)
        return tf.reshape(image, shape=(1, IMAGE_HEIGHT, IMAGE_WIDTH))