import numpy as np
import tensorflow as tf
import cv2

import matplotlib.pyplot as plt

def run(in_image):
    model = tf.keras.models.load_model('autoencoder-gray2color.keras')
    #in_image = cv2.imread(im_path)
    height, width, channels = in_image.shape
    # in_image = cv2.cvtColor(in_image, cv2.COLOR_BGR2GRAY)
    
    in_image = cv2.resize(in_image, (160, 160))
    in_image = in_image.astype('float32') / 255.0
    predicted = np.clip(model.predict(in_image.reshape(1,160, 160,3)),0.0,1.0).reshape(160, 160,3)
    
    # cv2.imwrite(f'colored-{im_path}', cv2.resize(predicted * 255., (height, width)))
    # plot_inference_images(in_image,predicted)
    out_img = cv2.resize(predicted * 255., (width, height))
    return out_img.astype(int)