import json
from typing import Any, Dict, List

import tensorflow as tf
from tensorflow import keras
from huggingface_hub import from_pretrained_keras, hf_hub_download
from PIL import Image
import base64
import numpy as np
from PIL import Image



class PreTrainedPipeline():
    def __init__(self, model_id: str):

        self.model = keras.models.load_model(model_id)

    def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:

        with Image.open(inputs) as img:
            img = np.array(img)

        im = tf.image.resize(img, (128, 128))
        im = tf.cast(im, tf.float32) / 255.0
        pred_mask = model.predict(im[tf.newaxis, ...])
        pred_mask_arg = tf.argmax(pred_mask, axis=-1)

        labels = []

        binary_masks = {}
        mask_codes = {}


        for cls in range(pred_mask.shape[-1]):

            binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2]))
            for row in range(pred_mask_arg[0][1].get_shape().as_list()[0]):

                for col in range(pred_mask_arg[0][2].get_shape().as_list()[0]):

                    if pred_mask_arg[0][row][col] == cls:
                        
                        binary_masks[f"mask_{cls}"][row][col] = 1
                    else:
                        binary_masks[f"mask_{cls}"][row][col] = 0


            mask_codes[f"mask_{cls}"] = base64.b64encode(binary_masks[f"mask_{cls}"])
    
    
        for i in range(pred_mask.shape[-1]): #for every class
            
            labels.append({
                "label": f"LABEL_{i}",
                "mask": mask_codes[f"mask_{i}"],
                "score": 1.0,
            })
            return labels