!pip install git+https://github.com/tensorflow/examples.git
# GENERAL

import os
import os.path
from pathlib import Path
import time
import pandas as pd
import numpy as np
import random
import time
import math
import glob
import cv2
from IPython.display import clear_output
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML


# I PACKAGES

import tensorflow as tf
import tensorflow.keras
from tensorflow.keras import backend as K
from tensorflow.keras.models import Sequential
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPool2D, BatchNormalization,MaxPooling2D,BatchNormalization,\
                        Permute, TimeDistributed, Bidirectional,GRU, SimpleRNN,\
LSTM, GlobalAveragePooling2D, SeparableConv2D, ZeroPadding2D, Convolution2D, ZeroPadding2D,Reshape, Conv2DTranspose, LeakyReLU, ReLU, Conv2DTranspose
from tensorflow.keras import models
from tensorflow.keras import layers
from tensorflow_examples.models.pix2pix import pix2pix

# WARNINGS

from warnings import filterwarnings
filterwarnings("ignore",category=DeprecationWarning)
filterwarnings("ignore", category=FutureWarning) 
filterwarnings("ignore", category=UserWarning)


AUTO_MODEL_TUNE = tf.data.AUTOTUNE
"""
# BEFORE TRAINING
"""
FACE_PATH = Path("../input/pretty-face/face/face")
SEG_PATH = Path("../input/pretty-face/layout/layout")
FACE_LIST_RED = list(FACE_PATH.glob(r"*.png"))
SEG_LIST_RED = list(SEG_PATH.glob(r"*.png"))
print(len(FACE_LIST_RED))
print(len(SEG_LIST_RED))
FACE_LIST = FACE_LIST_RED[:70]
SEG_LIST = SEG_LIST_RED[:70]
print(len(FACE_LIST))
print(len(SEG_LIST))
FACE_SERIES = pd.Series(FACE_LIST,name="FACE").astype(str)
SEG_SERIES = pd.Series(SEG_LIST,name="SEG").astype(str)

MAIN_PROCESS_DATA = pd.concat([FACE_SERIES,SEG_SERIES],axis=1)
MAIN_PROCESS_DATA
for index_random in range(5):
    
    
    EXAMPLE_FACE = cv2.cvtColor(cv2.imread(MAIN_PROCESS_DATA["FACE"][index_random]),cv2.COLOR_BGR2RGB)
    EXAMPLE_SEG = cv2.cvtColor(cv2.imread(MAIN_PROCESS_DATA["SEG"][index_random]),cv2.COLOR_BGR2RGB)


    figure,axis = plt.subplots(1,2,figsize=(10,10))

    axis[0].set_title("FACE")
    axis[0].set_xlabel(EXAMPLE_FACE.shape)
    axis[0].imshow(EXAMPLE_FACE)

    axis[1].set_title("SEG")
    axis[1].set_xlabel(EXAMPLE_SEG.shape)
    axis[1].imshow(EXAMPLE_SEG)

    plt.tight_layout()
    plt.show()
"""
# DATA PROCESS
"""
BUFFER_SIZE = 400
BATCH_SIZE = 1
IMG_WIDTH = 256
IMG_HEIGHT = 256
OUTPUT_CHANNELS = 3
LAMBDA = 100
EPOCHS = 70
TYPE_ARRAY = "float32"
FACE_MAIN_LIST = []
SEG_MAIN_LIST = []

for x_face,x_seg in zip(MAIN_PROCESS_DATA.FACE,MAIN_PROCESS_DATA.SEG):
    
    FACE_IMG = cv2.cvtColor(cv2.imread(x_face),cv2.COLOR_BGR2RGB)
    SEG_IMG = cv2.cvtColor(cv2.imread(x_seg),cv2.COLOR_BGR2RGB)
    
    RESIZE_FACE = cv2.resize(FACE_IMG,(IMG_WIDTH,IMG_HEIGHT))
    RESIZE_SEG = cv2.resize(SEG_IMG,(IMG_WIDTH,IMG_HEIGHT))
    
    REDUCE_FACE = RESIZE_FACE / 255.
    REDUCE_SEG = RESIZE_SEG / 255.
    
    FACE_MAIN_LIST.append(REDUCE_FACE)
    SEG_MAIN_LIST.append(REDUCE_SEG)
print(np.shape(np.array(FACE_MAIN_LIST)))
print(np.shape(np.array(SEG_MAIN_LIST)))
for index_random in range(5):
    
    
    EXAMPLE_FACE = FACE_MAIN_LIST[index_random]
    EXAMPLE_SEG = SEG_MAIN_LIST[index_random]


    figure,axis = plt.subplots(1,2,figsize=(10,10))

    axis[0].set_title("FACE")
    axis[0].set_xlabel(EXAMPLE_FACE.shape)
    axis[0].imshow(EXAMPLE_FACE)

    axis[1].set_title("SEG")
    axis[1].set_xlabel(EXAMPLE_SEG.shape)
    axis[1].imshow(EXAMPLE_SEG)

    plt.tight_layout()
    plt.show()
ARRAY_FACE = np.array(FACE_MAIN_LIST,dtype=TYPE_ARRAY)
ARRAY_SEG = np.array(SEG_MAIN_LIST,dtype=TYPE_ARRAY)
print(ARRAY_FACE.shape)
print(ARRAY_SEG.shape)
TARGET_TENSOR = tf.data.Dataset.from_tensor_slices(ARRAY_SEG).batch(BATCH_SIZE)
OUTPUT_TENSOR = tf.data.Dataset.from_tensor_slices(ARRAY_FACE).batch(BATCH_SIZE)
print(TARGET_TENSOR.element_spec)
print(OUTPUT_TENSOR.element_spec)
"""
# MODEL PROCESS
"""
"""
### LAYER CREATION
"""
def DOWN_SAMPLE(filters,size,apply_batchnorm=True):
    
    Ini_Kernel = tf.random_normal_initializer(0.,0.03)
    
    Model = Sequential()
    
    Model.add(Conv2D(filters,
                     size,
                     strides=2,
                     padding="same",
                     use_bias=False,
                     kernel_initializer=Ini_Kernel))
    
    
    if apply_batchnorm:
        Model.add(BatchNormalization())
        
        
    return Model
Model_DOWN = DOWN_SAMPLE(3,4)
def UP_SAMPLE(filters,size,apply_dropout=False):
    
    Ini_Kernel = tf.random_normal_initializer(0.,0.03)
    
    Model = Sequential()
    
    Model.add(Conv2DTranspose(filters,
                              size,
                              strides=2,
                              padding="same",
                              use_bias=False,
                              kernel_initializer=Ini_Kernel))
    
    
    Model.add(BatchNormalization())
    
    if apply_dropout:
        
        Model.add(Dropout(0.2))
        
    Model.add(ReLU())
    
    return Model
Model_UP = UP_SAMPLE(3,4)
"""
### GENERATOR CREATION
"""
def Generator_Model(width_i,height_i,dimension_i):
    
    INPUT_LAYER = tf.keras.layers.Input(shape=[width_i,height_i,dimension_i])
    
    DOWN_LIST = [
        
        DOWN_SAMPLE(64,4,apply_batchnorm=False),
        DOWN_SAMPLE(128,4),
        DOWN_SAMPLE(256,4),
        DOWN_SAMPLE(512,4),
        DOWN_SAMPLE(512,4),
        DOWN_SAMPLE(512,4),
        DOWN_SAMPLE(512,4),
        DOWN_SAMPLE(512,4)
    ]
    
    
    UP_LIST = [
        
        UP_SAMPLE(512,4,apply_dropout=True),
        UP_SAMPLE(512,4,apply_dropout=True),
        UP_SAMPLE(512,4),
        UP_SAMPLE(512,4),
        UP_SAMPLE(256,4),
        UP_SAMPLE(128,4),
        UP_SAMPLE(64,4)
    ]
    
    
    Ini_Kernel = tf.random_normal_initializer(0.,0.03)
    
    TRANS_LAYER = Conv2DTranspose(OUTPUT_CHANNELS,4,strides=2,
                                 padding="same",
                                 kernel_initializer=Ini_Kernel,
                                 activation="tanh")
    
    LAST_X = INPUT_LAYER
    
    SKIPPING_PRE_LAYER = []
    
    for x_down in DOWN_LIST:
        
        LAST_X = x_down(LAST_X)
        SKIPPING_PRE_LAYER.append(LAST_X)
        
    SKIPPING_PRE_LAYER = reversed(SKIPPING_PRE_LAYER[:-1])
    
    for x_up,x_skip in zip(UP_LIST,SKIPPING_PRE_LAYER):
        
        LAST_X = x_up(LAST_X)
        LAST_X = tf.keras.layers.Concatenate()([LAST_X,x_skip])
        
    LAST_X = TRANS_LAYER(LAST_X)
    
    return tf.keras.Model(inputs=INPUT_LAYER,outputs=LAST_X)
Generation_Model = Generator_Model(IMG_WIDTH,IMG_HEIGHT,OUTPUT_CHANNELS) 
tf.keras.utils.plot_model(Generation_Model, show_shapes=True)
"""
### DISCRIMINATOR CREATION
"""
def Discriminator_Model(width_i,height_i,dimension_i):
    
    Ini_Kernel = tf.random_normal_initializer(0.,0.03)
    
    INPUT_LAYER = tf.keras.layers.Input(shape=[width_i,height_i,dimension_i],name="INPUT_LAYER")
    TARGET_LAYER = tf.keras.layers.Input(shape=[width_i,height_i,dimension_i],name="TARGET_LAYER")
    
    LAYER_X = tf.keras.layers.concatenate([INPUT_LAYER,TARGET_LAYER])
    
    DOWN_SAMPLE_1 = DOWN_SAMPLE(64,4,False)(LAYER_X)
    DOWN_SAMPLE_2 = DOWN_SAMPLE(64,4)(DOWN_SAMPLE_1)
    DOWN_SAMPLE_3 = DOWN_SAMPLE(64,4)(DOWN_SAMPLE_2)
    
    ZERO_PADD_LAY_1 = tf.keras.layers.ZeroPadding2D()(DOWN_SAMPLE_3)
    CONV2D_LAY_1 = Conv2D(512,
                        4,
                        strides=1,
                        kernel_initializer=Ini_Kernel,
                        use_bias=False)(ZERO_PADD_LAY_1)
    
    BATCH_LAY_1 = BatchNormalization()(CONV2D_LAY_1)
    ACT_LAY_1 = LeakyReLU()(BATCH_LAY_1)
    
    ZERO_PADD_LAY_2 = tf.keras.layers.ZeroPadding2D()(ACT_LAY_1)
    
    LAST_LAY = Conv2D(1,
                     4,
                     strides=1,
                     kernel_initializer=Ini_Kernel)(ZERO_PADD_LAY_2)
    
    return tf.keras.Model(inputs=[INPUT_LAYER,TARGET_LAYER],outputs=LAST_LAY)
Discrimination_Model = Discriminator_Model(IMG_WIDTH,IMG_HEIGHT,OUTPUT_CHANNELS)
tf.keras.utils.plot_model(Discrimination_Model, show_shapes=True)
"""
### LOSS FUNCTION CREATION
"""
LOSS_FUNCTION = tf.keras.losses.BinaryCrossentropy(from_logits=True)
def Generation_Loss(discriminator_out,generator_out,target_in):
    
    GAN_LOSS = LOSS_FUNCTION(tf.ones_like(discriminator_out),discriminator_out)
    
    MEAN_ABS_ERROR_OUT = tf.reduce_mean(tf.abs(target_in - generator_out))
    
    TOTAL_LOSS_OUT = GAN_LOSS + (LAMBDA * MEAN_ABS_ERROR_OUT)
    
    return TOTAL_LOSS_OUT,GAN_LOSS,MEAN_ABS_ERROR_OUT
def Discrimination_Loss(real_out,generated_out):
    
    REAL_LOSS = LOSS_FUNCTION(tf.ones_like(real_out),real_out)
    GEN_LOSS = LOSS_FUNCTION(tf.zeros_like(generated_out),generated_out)
    
    TOTAL_LOSS = REAL_LOSS + GEN_LOSS
    
    return TOTAL_LOSS
"""
### OPTIMIZATION AND GENERATION IMAGES CREATION
"""
GENERATOR_OPT = tf.keras.optimizers.RMSprop(lr=0.0003,clipvalue=1.0,decay=1e-8)
DISCRIMINATOR_OPT = tf.keras.optimizers.RMSprop(lr=0.0003,clipvalue=1.0,decay=1e-8)
os.mkdir("./OUT_CLASS")

def generate_images_single(model, test_input, tar, number_i):
    
    prediction = model(test_input, training=True)
    
    plt.figure(figsize=(15, 15))
    
    plt.imshow(prediction[0])
    plt.savefig('./OUT_CLASS/out_res_{:04d}.png'.format(number_i))
    plt.axis('off')
    
    plt.tight_layout()    
    plt.show()
os.mkdir("./MULTI_CLASS")

def generate_images_for_example(model, test_input, tar, number_i):
    
    prediction = model(test_input, training=True)
    
    RESULT_IN_OUT_LIST = [test_input[0], tar[0], prediction[0]]
    
    
    figure,axis = plt.subplots(1, 3,figsize=(14,14))

    axis[0].imshow(RESULT_IN_OUT_LIST[0])
    axis[0].axis('off')

    axis[1].imshow(RESULT_IN_OUT_LIST[1])
    axis[1].axis('off')

    axis[2].imshow(RESULT_IN_OUT_LIST[2])
    axis[2].axis('off')
    
    plt.savefig('./MULTI_CLASS/multi_res_{:04d}.png'.format(number_i))
    
    plt.tight_layout()
    plt.show()
"""
### TRAINING STEP CREATION - LAST STEP
"""
@tf.function

def Train_Process(INPUT_TENSOR,TARGET_TENSOR,EPOCH):
    
    with tf.GradientTape() as GEN_TAPE, tf.GradientTape() as DISC_TAPE:
        
        GENERATION_OUT = Generation_Model(INPUT_TENSOR,training=True)
        
        DISCRIMINATION_REAL = Discrimination_Model([INPUT_TENSOR,TARGET_TENSOR],training=True)
        DISCRIMINATION_FAKE = Discrimination_Model([INPUT_TENSOR,GENERATION_OUT],training=True)
        
        GEN_TOTAL_LOSS,GENERATOR_LOSS,GEN_MEAN_ABS_ERROR = Generation_Loss(DISCRIMINATION_FAKE,
                                                                          GENERATION_OUT,
                                                                          TARGET_TENSOR)
        
        
        DISC_LOSS = Discrimination_Loss(DISCRIMINATION_REAL,DISCRIMINATION_FAKE)
        
        
        
    Generator_Gradient = GEN_TAPE.gradient(GEN_TOTAL_LOSS,Generation_Model.trainable_variables)
    Discriminator_Gradient = DISC_TAPE.gradient(DISC_LOSS,Discrimination_Model.trainable_variables)
    
    GENERATOR_OPT.apply_gradients(zip(Generator_Gradient,Generation_Model.trainable_variables))
    DISCRIMINATOR_OPT.apply_gradients(zip(Discriminator_Gradient,Discrimination_Model.trainable_variables))
"""
### TRAINING SECTION
"""
EXAMPLE_INP,EXAMPLE_TAR = next(iter(TARGET_TENSOR.take(1))),next(iter(OUTPUT_TENSOR.take(1)))

counting_img = 0

for epoch in range(EPOCHS):

    n_count = 0
    
    for image_x, image_y in tf.data.Dataset.zip((TARGET_TENSOR.take(epoch),OUTPUT_TENSOR.take(epoch))):
        
        Train_Process(image_x, image_y, epoch)
        
        if n_count % 10 == 0:
            
            print ('.', end='')
            
        n_count += 1
        
    
        clear_output(wait=True)
        
        generate_images_for_example(Generation_Model,
                                    image_x,
                                    image_y,
                                    counting_img)
    
    
    
    
        counting_img = counting_img + 1
TEST_SEG_LIST_NEW = SEG_LIST_RED[600:605]
TEST_SERIES = pd.Series(TEST_SEG_LIST_NEW,name="TEST").astype(str)
TESTING_IMG_LIST = []

for x_seg in TEST_SERIES:

    SEG_IMG = cv2.cvtColor(cv2.imread(x_seg),cv2.COLOR_BGR2RGB)
    
    RESIZE_SEG = cv2.resize(SEG_IMG,(IMG_WIDTH,IMG_HEIGHT))
    
    REDUCE_SEG = RESIZE_SEG / 255.

    TESTING_IMG_LIST.append(REDUCE_SEG)
TEST_ARRAY = np.array(TESTING_IMG_LIST,dtype="float32")
TEST_TENSOR = tf.data.Dataset.from_tensor_slices(TEST_ARRAY).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
def generate_images_predict(model, test_input):
    
    prediction = model(test_input, training=True)
    
    plt.figure(figsize=(15, 15))
    
    plt.imshow(prediction[0])
    plt.savefig('./OUT_CLASS/out_res_{:04d}.png'.format(2))
    plt.axis('off')
    
    plt.tight_layout()    
    plt.show()
for x_counting_tensor in TEST_TENSOR.take(3):
    
    generate_images_predict(Generation_Model,x_counting_tensor)
TEST_SEG_LIST_NEW_2 = SEG_LIST_RED[610:655]
TEST_SERIES_2 = pd.Series(TEST_SEG_LIST_NEW_2,name="TEST_'").astype(str)
TESTING_IMG_LIST_2 = []

for x_seg in TEST_SERIES_2:

    SEG_IMG = cv2.cvtColor(cv2.imread(x_seg),cv2.COLOR_BGR2RGB)
    
    RESIZE_SEG = cv2.resize(SEG_IMG,(IMG_WIDTH,IMG_HEIGHT))
    
    REDUCE_SEG = RESIZE_SEG / 255.

    TESTING_IMG_LIST_2.append(REDUCE_SEG)
TEST_ARRAY_2 = np.array(TESTING_IMG_LIST_2,dtype="float32")
TEST_TENSOR_2 = tf.data.Dataset.from_tensor_slices(TEST_ARRAY_2).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
LIST_PREDICT_MULTI = []
os.mkdir("./MULTI_CLASS_PREDICT")
def generate_images_for_multi_predict(model, test_input):
    
    prediction = model(test_input, training=True)
    
    RESULT_IN_OUT_LIST = [test_input[0], prediction[0]]
    
    
    figure,axis = plt.subplots(1, 2,figsize=(15,10))

    axis[0].imshow(RESULT_IN_OUT_LIST[0])
    axis[0].axis('off')

    axis[1].imshow(RESULT_IN_OUT_LIST[1])
    axis[1].axis('off')

    SAVE_NUM = random.randint(0,120)
    plt.savefig('./MULTI_CLASS_PREDICT/multi_res_predict{:04d}.png'.format(SAVE_NUM))
    
    PRE_READ_IMG = cv2.cvtColor(cv2.imread('./MULTI_CLASS_PREDICT/multi_res_predict{:04d}.png'.format(SAVE_NUM)),cv2.COLOR_BGR2RGB)
    LIST_PREDICT_MULTI.append(PRE_READ_IMG)
    
    plt.tight_layout()
    plt.show()
for x_counting_tensor in TEST_TENSOR_2.take(20):
    
    generate_images_for_multi_predict(Generation_Model,x_counting_tensor)
os.mkdir("./SINGLE_CLASS_PREDICT")
SINGLE_READ_LIST = []

def generate_images_single_pre(model, test_input):
    
    plt.style.use("dark_background")
    
    prediction = model(test_input, training=True)
    
    plt.figure(figsize=(10, 10))
    plt.axis('off')
    plt.imshow(prediction[0])
    
    SAVE_NUM = random.randint(0,120)
    plt.savefig('./SINGLE_CLASS_PREDICT/sin_res_{:04d}.png'.format(SAVE_NUM))
    
    
    
    PRE_READ_IMG = cv2.cvtColor(cv2.imread('./SINGLE_CLASS_PREDICT/sin_res_{:04d}.png'.format(SAVE_NUM)),cv2.COLOR_BGR2RGB)
    SINGLE_READ_LIST.append(PRE_READ_IMG)
    
    
    plt.tight_layout()    
    plt.show()
for x_counting_tensor in TEST_TENSOR_2.take(32):
    
    generate_images_single_pre(Generation_Model,x_counting_tensor)
def displaying_plot(source):
    
    figure = plt.figure(figsize=(16,11))

    Image_List = []
    plt.style.use("dark_background")
    for indexing in source:
        
        Read_IMG = plt.imshow(indexing, animated=True,cmap="hot")
        plt.axis('off')
        Image_List.append([Read_IMG])

    Animation_Func = animation.ArtistAnimation(figure, Image_List, interval=350, repeat_delay=10200)
    
    plt.close()
    
    return Animation_Func
HTML(displaying_plot(SINGLE_READ_LIST).to_html5_video())
HTML(displaying_plot(LIST_PREDICT_MULTI).to_html5_video())
OUTPUT_RES_PATH = Path("./MULTI_CLASS")
OUTPUT_RES_LIST = list(OUTPUT_RES_PATH.glob(r"*.png"))
print(len(OUTPUT_RES_LIST))

OUT_RES_SERIES = pd.Series(OUTPUT_RES_LIST,name="RES").astype(str)
OUT_READ_RES = []

for x_seg in OUT_RES_SERIES:

    SEG_IMG = cv2.cvtColor(cv2.imread(x_seg),cv2.COLOR_BGR2RGB)
    
    RESIZE_SEG = cv2.resize(SEG_IMG,(IMG_WIDTH,IMG_HEIGHT))
    
    REDUCE_SEG = RESIZE_SEG / 255.

    OUT_READ_RES.append(REDUCE_SEG)
HTML(displaying_plot(OUT_READ_RES).to_html5_video())