image_colorizer_alpha_version_predict.py

def colorize(image_name):
    image = img_to_array(load_img(image_name))
    image = np.array(image, dtype=np.float32)
    X = rgb2lab(1.0/255*image)[:, :, 0]
    X = X.reshape(1, 400, 400, 1)

    model = ct.models.MLModel('model_alpha.mlmodel')
    input = np.zeros((1, 400, 400))
    input[0] = X[0][:, :, 0]
    output = model.predict({'input_1': X})["Identity"]
    output *= 128

    # Create & save output
    output_lab = np.zeros((400, 400, 3))
    output_lab[:, :, 0] = X[0][:, :, 0]
    output_lab[:, :, 1] = output[0][:, :, 0]
    output_lab[:, :, 2] = output[0][:, :, 1]

    output_rgb = lab2rgb(output_lab)
    imsave("output.jpg", output_rgb)