colorizer_train.py

# -*- coding: utf-8 -*-
"""colorizer_train.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1M1YzWG-tYBsQbxLABl2sTXSLh48gZDzV
"""

import tensorflow as tf

from tensorflow.keras import datasets, layers, models
import numpy as np

import matplotlib.pyplot as plt
import seaborn as sns

from keras.preprocessing.image import ImageDataGenerator

path = '/content/Grayscale-to-Color-AutoEncoder'
image_datagen = ImageDataGenerator(rescale = 1./255)

train = image_datagen.flow_from_directory(path, target_size=(256, 256),batch_size=7000,class_mode=None)

from skimage.color import rgb2lab

X =[]
Y =[]
for img in train[0]:
  try:
      lab = rgb2lab(img)
      X.append(lab[:,:,0])
      Y.append(lab[:,:,1:] / 128)
  except:
     print('error')
X = np.array(X)
Y = np.array(Y)
X = X.reshape(X.shape+(1,))
print(X.shape)
print(Y.shape)

input_shape = (256,256,1)

#Encoder
encoder_input = tf.keras.layers.Input(shape=(256, 256, 1,))
encoder_output = tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding='same', strides=2)(encoder_input)
encoder_output = tf.keras.layers.Conv2D(128, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = tf.keras.layers.Conv2D(128, (3,3), activation='relu', padding='same', strides=2)(encoder_output)
encoder_output = tf.keras.layers.Conv2D(256, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = tf.keras.layers.Conv2D(256, (3,3), activation='relu', padding='same', strides=2)(encoder_output)
encoder_output = tf.keras.layers.Conv2D(512, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = tf.keras.layers.Conv2D(512, (3,3), activation='relu', padding='same')(encoder_output)
encoder_output = tf.keras.layers.Conv2D(256, (3,3), activation='relu', padding='same')(encoder_output)

#Decoder
decoder_output = tf.keras.layers.Conv2D(128, (3,3), activation='relu', padding='same')(encoder_output)
decoder_output = tf.keras.layers.UpSampling2D((2, 2))(decoder_output)
decoder_output = tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = tf.keras.layers.UpSampling2D((2, 2))(decoder_output)
decoder_output = tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = tf.keras.layers.Conv2D(16, (3,3), activation='relu', padding='same')(decoder_output)
decoder_output = tf.keras.layers.Conv2D(2, (3, 3), activation='tanh', padding='same')(decoder_output)
decoder_output = tf.keras.layers.UpSampling2D((2, 2))(decoder_output)
model = tf.keras.Model(inputs=encoder_input, outputs=decoder_output)

#model = Sequential([Encoder,Decoder])
model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001), loss = 'mse',metrics = ['acc'])
model.fit(X, Y, epochs=50, verbose=2, validation_split = 0.1, batch_size = 16)

model.save('model_1k')