text: #Ref: Sreenivas Sarwar Anik
  
  """
  
  Image colorization using Autoencoders. Transfer learning using VGG.
  
  """
  
  from keras.layers import Conv2D, UpSampling2D, Input
  from keras.models import Sequential, Model
  from keras.preprocessing.image import ImageDataGenerator, img_to_array, load_img
  from skimage.color import rgb2lab, lab2rgb, gray2rgb
  from skimage.transform import resize
  from skimage.io import imsave
  import numpy as np
  import tensorflow as tf
  import keras
  import os
  
  """
  Because we are going to replace the encoder part with VGG16, 
  we don’t need it as a classifier, we need it as a feature extractor so, 
  the last dense layers isn’t needed we have to pop them up.
  
  here, we iterate on each layer except the last dense layers so, 
  we add 19 layer to our model. the dimension of last layer volume is “7x7x512”. 
  we will be using that latent space volume as a feature vector to be input to the decoder.
   and the decoder is going to learn the mapping from the latent space vector to ab channels. 
   we want the layers of VGG16 with its original weights without changing them, 
   so that we set the trainable parameter in each layer to false because we don’t want to train them again.
  """
  from keras.applications.vgg16 import VGG16
  vggmodel = VGG16()
  newmodel = Sequential() 
  #num = 0
  for i, layer in enumerate(vggmodel.layers):
      if i<19:          #Only up to 19th layer to include feature extraction only
        newmodel.add(layer)
  newmodel.summary()
  for layer in newmodel.layers:
    layer.trainable=False   #We don't want to train these layers again, so False. 
  
  """
  VGG16 is expecting an image of 3 dimension with size 224x224 as an input, 
  in preprocessing we have to scale all images to 224 instead of 256
  """
  path = 'images/colorization/'
  #Normalize images - divide by 255
  train_datagen = ImageDataGenerator(rescale=1. / 255)
  
  train = train_datagen.flow_from_directory(path, target_size=(224, 224), batch_size=32, class_mode=None)
  
  
  #Convert from RGB to Lab
  """
  by iterating on each image, we convert the RGB to Lab. 
  Think of LAB image as a grey image in L channel and all color info stored in A and B channels. 
  The input to the network will be the L channel, so we assign L channel to X vector. 
  And assign A and B to Y.
  
  """
  
  X =[]
  Y =[]
  for img in train[0]:
    try:
        lab = rgb2lab(img)
        X.append(lab[:,:,0]) 
        Y.append(lab[:,:,1:] / 128) #A and B values range from -127 to 128, 
        #so we divide the values by 128 to restrict values to between -1 and 1.
    except:
       print('error')
  X = np.array(X)
  Y = np.array(Y)
  X = X.reshape(X.shape+(1,)) #dimensions to be the same for X and Y
  print(X.shape)
  print(Y.shape)
  
  
  #now we have one channel of L in each layer but, VGG16 is expecting 3 dimension, 
  #so we repeated the L channel two times to get 3 dimensions of the same L channel
  
  vggfeatures = []
  for i, sample in enumerate(X):
    sample = gray2rgb(sample)
    sample = sample.reshape((1,224,224,3))
    prediction = newmodel.predict(sample)
    prediction = prediction.reshape((7,7,512))
    vggfeatures.append(prediction)
  vggfeatures = np.array(vggfeatures)
  print(vggfeatures.shape)
  
  
  #Decoder
  model = Sequential()
  
  model.add(Conv2D(256, (3,3), activation='relu', padding='same', input_shape=(7,7,512)))
  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(UpSampling2D((2, 2)))
  model.add(Conv2D(16, (3,3), activation='relu', padding='same'))
  model.add(UpSampling2D((2, 2)))
  model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
  model.add(UpSampling2D((2, 2)))
  model.summary()
  
  
  model.compile(optimizer='Adam', loss='mse' , metrics=['accuracy'])
  model.fit(vggfeatures, Y, verbose=1, epochs=10, batch_size=128)
  
  model.save('colorize_autoencoder_VGG16.model')
  
  ############################################
  #Predicting using saved model.
  model = tf.keras.models.load_model('colorize_autoencoder_VGG16_10000.model',
                                     custom_objects=None,
                                     compile=True)
  testpath = 'images/colorization2/test_images/'
  files = os.listdir(testpath)
  for idx, file in enumerate(files):
      test = img_to_array(load_img(testpath+file))
      test = resize(test, (224,224), anti_aliasing=True)
      test*= 1.0/255
      lab = rgb2lab(test)
      l = lab[:,:,0]
      L = gray2rgb(l)
      L = L.reshape((1,224,224,3))
      #print(L.shape)
      vggpred = newmodel.predict(L)
      ab = model.predict(vggpred)
      #print(ab.shape)
      ab = ab*128
      cur = np.zeros((224, 224, 3))
      cur[:,:,0] = l
      cur[:,:,1:] = ab
      imsave('images/colorization2/vgg_result/result'+str(idx)+".jpg", lab2rgb(cur))
  