import os
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Force CPU
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'  # 0 = all messages are logged, 3 - INFO, WARNING, and ERROR messages are not printed

from keras.layers import Conv2D, Conv2DTranspose, UpSampling2D
from keras.layers import Activation, Dense, Dropout, Flatten, InputLayer
from keras.layers.normalization import BatchNormalization
from keras.callbacks import TensorBoard
from keras.models import Sequential
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from skimage.color import rgb2lab, lab2rgb, rgb2gray
from skimage.io import imsave
import numpy as np
import random
import tensorflow as tf
import cv2


folder_train = "Faces"
folder_src = "Input"
folder_dst = "Output"
model_file = "faces1.h5"
brightness_corr = 250
do_train = True

# Get train images
X = []
for filename in os.listdir(folder_train):
    X.append(img_to_array(load_img(folder_train + os.sep + filename)))
X = np.array([cv2.resize(i, (256, 256)) for i in X], dtype=float)/255.0
# X = np.array(X, dtype=float)
Xtrain = X

# Model
model = Sequential()
model.add(InputLayer(input_shape=(256, 256, 1)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))
model.compile(optimizer='rmsprop', loss='mse')

# Image transformer
datagen = ImageDataGenerator(shear_range=0.2, zoom_range=0.2, rotation_range=20, horizontal_flip=True)

# Generate training data
batch_size = 10
def image_a_b_gen(batch_size):
    for batch in datagen.flow(Xtrain, batch_size=batch_size):
        lab_batch = rgb2lab(batch)
        X_batch = lab_batch[:,:,:,0]
        Y_batch = lab_batch[:,:,:,1:] / 128
        yield (X_batch.reshape(X_batch.shape+(1,)), Y_batch)

if do_train:
    # Train model
    model.fit_generator(image_a_b_gen(batch_size), epochs=1, steps_per_epoch=400)

    # Save model
    model.save_weights(model_file)

# Load model
model.load_weights(model_file)

# Process images
color_me = []
for filename in os.listdir(folder_src):
    color_me.append(img_to_array(load_img(folder_src + os.sep + filename)))
color_me = np.array([cv2.resize(i, (256, 256)) for i in color_me], dtype=float)
# color_me = np.array(color_me, dtype=float)
color_me = rgb2lab(1.0/255*color_me)[:,:,:,0]
color_me = color_me.reshape(color_me.shape+(1,))

# Test model
output = model.predict(color_me)

# Output colorizations
for i in range(len(output)):
    cur = np.zeros((256, 256, 3))
    cur[:,:,0] = color_me[i][:,:,0]
    cur[:,:,1:] = output[i] * 128
    img_rgb = lab2rgb(cur)*brightness_corr
    imsave(folder_dst + os.sep + "img_%d.png" % i, img_rgb.astype(np.uint8))
    print("img_%d.png saved" % i)

import numpy as np
from typing import *
import cv2
import os

# Sources:
# https://www.learnopencv.com/convolutional-neural-network-based-image-colorization-using-opencv/
# https://www.pyimagesearch.com/2019/02/25/black-and-white-image-colorization-with-opencv-and-deep-learning/
# http://richzhang.github.io/colorization/

# Models:
# http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2.caffemodel
# http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v2/colorization_release_v2_norebal.caffemodel
# http://eecs.berkeley.edu/~rich.zhang/projects/2016_colorization/files/demo_v1/colorization_release_v1.caffemodel


prototxt = "models/colorization_deploy_v2.prototxt"
model = "models/colorization_release_v2.caffemodel"
points = "models/pts_in_hull.npy"
folder_in = "Input"
folder_out = "Output"


def load_model() -> Any:
    # Load serialized black and white colorizer model and cluster
    # The L channel encodes lightness intensity only
    # The a channel encodes green-red.
    # And the b channel encodes blue-yellow
    print("Loading model...")

    net = cv2.dnn.readNetFromCaffe(prototxt, model)
    pts = np.load(points)

    # Add the cluster centers as 1x1 convolutions to the model:

    class8 = net.getLayerId("class8_ab")
    conv8 = net.getLayerId("conv8_313_rh")
    pts = pts.transpose().reshape(2, 313, 1, 1)
    net.getLayer(class8).blobs = [pts.astype("float32")]
    net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")]
    return net


def colorize_image(net: Any, image_in: str, image_out: str):
    # Load the input image, scale it and convert it to Lab:
    image = cv2.imread(image_in)
    height, width, channels = image.shape
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    # Extracting "L"
    scaled = image.astype("float32") / 255.0
    lab = cv2.cvtColor(scaled, cv2.COLOR_RGB2LAB)

    # Resize to network size
    resized = cv2.resize(lab, (224, 224))
    L = cv2.split(resized)[0]
    L -= 50

    # Predicting "a" and "b"
    net.setInput(cv2.dnn.blobFromImage(L))
    ab = net.forward()[0, :, :, :].transpose((1, 2, 0))

    # Creating a colorized Lab photo (L + a + b)
    L = cv2.split(lab)[0]
    ab = cv2.resize(ab, (width, height))
    colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)

    # Convert to RGB
    colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2RGB)
    colorized = np.clip(colorized, 0, 1)
    colorized = (255 * colorized).astype("uint8")

    cv2.imwrite(image_out, cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR))
    print("Image %s saved" % image_out)


if __name__ == '__main__':

    net = load_model()
    for filename in os.listdir(folder_in):
       colorize_image(net, folder_in + os.sep + filename, folder_out + os.sep + filename.replace(".", "-out."))

def colorize_video(net: Any, video_in: str, video_out: str):
    vid_in = cv2.VideoCapture(video_in)
    frames, fps = int(vid_in.get(cv2.CAP_PROP_FRAME_COUNT)), int(round(vid_in.get(cv2.CAP_PROP_FPS)))
    width, height = int(vid_in.get(cv2.CAP_PROP_FRAME_WIDTH)), int(vid_in.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print("Video {}: {}x{}, {} frames, {} fps".format(video_in, width, height, frames, fps))
    size_out = (width, height)
    vid_out = cv2.VideoWriter(video_out, cv2.VideoWriter_fourcc(*"mp4v"), fps, size_out)  # .mp4
    # vid_out = cv2.VideoWriter(video_out, cv2.VideoWriter_fourcc(*'DIVX'), fps, size_out)  # .avi

    count = 0
    while True:
        success, frame = vid_in.read()
        if frame is None or success is False:
            break

        if (count % 10) == 0:
            print("Frame {} of {}".format(count, frames))

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
        scaled = frame.astype("float32") / 255.0
        lab = cv2.cvtColor(scaled, cv2.COLOR_RGB2LAB)
        resized = cv2.resize(lab, (224, 224))
        L = cv2.split(resized)[0]
        L -= 50

        net.setInput(cv2.dnn.blobFromImage(L))
        ab = net.forward()[0, :, :, :].transpose((1, 2, 0))
        ab = cv2.resize(ab, (frame.shape[1], frame.shape[0]))
        L = cv2.split(lab)[0]
        colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
        colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR)
        colorized = np.clip(colorized, 0, 1)
        colorized = (255 * colorized).astype("uint8")
        frame_out = colorized  # cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR)

        vid_out.write(frame_out)

        count += 1

    vid_in.release()
    vid_out.release()

    print("File %s saved" % video_out)