import json
from typing import Any, Dict, List

import tensorflow as tf
from tensorflow import keras
from PIL import Image
import base64


MODEL_FILENAME = "saved_model.pb"
CONFIG_FILENAME = "config.json"


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


        # Reload Keras SavedModel
        self.model = keras.models.load_model('./model.h5')

        # Number of labels
        self.num_labels = self.model.output_shape[1]

        self.id2label = self.id2label = {"0": "pet", "1":"no_pet"}

    def __call__(self, inputs: "Image.Image") -> List[Dict[str, Any]]:
        """
        Args:
            inputs (:obj:`PIL.Image`):
                The raw image representation as PIL.
                No transformation made whatsoever from the input. Make all necessary transformations here.
        Return:
            A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX" (str), mask: "base64 encoding of the mask" (str), "score": float}
                It is preferred if the returned list is in decreasing `score` order
        """
        # Resize image to expected size

        expected_input_size = self.model.input_shape

        
        with Image.open(inputs) as im:
            inputs = np.array(im)
        if expected_input_size[-1] == 1:  
            inputs = inputs.convert("L")


        
        target_size = (expected_input_size[1], expected_input_size[2])
        img = tf.image.resize(inputs, target_size)
        
        img_array = tf.keras.preprocessing.image.img_to_array(img)
        img_array = img_array[tf.newaxis, ...]
        
        predictions = self.model.predict(img_array)

        self.single_output_unit = (
            self.model.output_shape[1] == 1
        )  # if there are two classes
        
        labels = []
        for i in enumerate(predictions):
            
            labels.append({
                "label": str(i[0]),
                "mask": base64.b64encode(i[1]),
                "score": 1.0,
            })

        return sorted(labels)