# AUTOGENERATED! DO NOT EDIT! File to edit: ../10_tensorflow_model.ipynb.

# %% auto 0
__all__ = ['crop_pad_pil', 'ukbb_model']

# %% ../10_tensorflow_model.ipynb 3
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow.compat.v1 as tf
from PIL import Image, ImageOps
import numpy as np

# %% ../10_tensorflow_model.ipynb 4
def crop_pad_pil(image, size):
    to_cut_w=((image.size[0]-size[0])/2)
    to_cut_h=((image.size[1]-size[1])/2)
    image = ImageOps.crop(image, (np.floor(to_cut_w), np.floor(to_cut_h), np.ceil(to_cut_w), np.ceil(to_cut_h)))
    return image

# %% ../10_tensorflow_model.ipynb 5
class ukbb_model:
    def __init__(self, model_path):
        tf.disable_eager_execution()
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())
        saver = tf.train.import_meta_graph(f'{model_path}.meta')
        saver.restore(self.sess, model_path)
        
    def prepareSize(self, image):
        X, Y = image.size
        image=crop_pad_pil(image,(int(np.ceil(X / 16.0)) * 16, int(np.ceil(Y / 16.0)) * 16))
        return image
    
    def image_to_input(self, image):
        img = image
        img = np.array(img)[:,:,0]
        img = img/255
        img = np.expand_dims(img, 0)
        img = np.expand_dims(img, -1)
        return img
    
    def predict(self, image):
        image = self.image_to_input(image)
        preds, classes = self.sess.run(['prob:0', 'pred:0'],
                   feed_dict={'image:0': image, 'training:0': False})
        preds = np.squeeze(preds, 0)
        preds = np.argmax(preds, axis = 2)
        preds = preds.astype(np.uint8)
        preds = Image.fromarray(preds)
        return preds
