from typing import Dict, List, Any
from PIL import Image
from fcnutr import FCNutr

import os
import tensorflow as tf


class PreTrainedPipeline():
    def __init__(self, path=""):
        self.crop_size = (224, 224)
        self.nutr_names = ('energy', 'fat', 'protein', 'carbs')
        self.img_size = 256
        self.model = FCNutr(self.nutr_names, self.crop_size, 4096, 3, False)
        self.model.compile()
        self.model(tf.zeros((1, self.crop_size[0], self.crop_size[1], 3)))
        self.model.load_weights(os.path.join(path, "fcnutr.h5"))

    def __call__(self, inputs: "Image.Image") -> List[Dict[str, Any]]:
        image = tf.keras.preprocessing.image.img_to_array(inputs)
        height = tf.shape(image)[0]
        width = tf.shape(image)[1]
        if width > height:
            image = tf.image.resize(image, (self.img_size, int(float(self.img_size * width) / float(height))))
        else:
            image = tf.image.resize(image, (int(float(self.img_size * height) / float(width)), self.img_size))

        image = tf.keras.applications.inception_v3.preprocess_input(image)
        image = tf.keras.layers.CenterCrop(*self.crop_size)(image)
        prediction = self.model(image[tf.newaxis, :])
        return [{"label": name, "score": float(prediction[name].numpy()[0, 0])} for name in self.nutr_names]