diff --git a/Autoencoder.py b/Autoencoder.py
index 1f29067..5bb68a5 100644
--- a/Autoencoder.py
+++ b/Autoencoder.py
@@ -7,14 +7,14 @@ from Decoder import *
 class Autoencoder(tf.keras.Model):
     def __init__(self):
         super(Autoencoder, self).__init__()
-
         self.encoder = Encoder()
         self.decoder = Decoder()
 
+        self.optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)
+
         self.loss_function = tf.keras.losses.MeanSquaredError()
 
-        self.optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
-
+        self.metric_mean = tf.keras.metrics.Mean(name="loss")
 
     @tf.function
     def call(self, x, training=False):
@@ -31,16 +31,24 @@ class Autoencoder(tf.keras.Model):
 
         gradients = tape.gradient(loss, self.trainable_variables)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-        return loss
+
+        self.metric_mean.update_state(loss)
+
+        
 
     def test(self, test_data):
+
+        self.metric_mean.reset_states()
+
         # test over complete test data
-        test_loss_aggregator = []
-        for input, target in test_data: # ignore label            
+        for input, target in test_data:           
             prediction = self(input)
             
-            sample_test_loss = self.loss_function(target, prediction)
-            test_loss_aggregator.append(sample_test_loss.numpy())
+            loss = self.loss_function(target, prediction)
+            self.metric_mean.update_state(loss)
 
-        test_loss = tf.reduce_mean(test_loss_aggregator)
-        return test_loss
\ No newline at end of file
+        mean_loss = self.metric_mean.result()
+        self.metric_mean.reset_states()
+        return mean_loss
+
+    
diff --git a/Colorful_Image_Colorization/__pycache__/config.cpython-38.pyc b/Colorful_Image_Colorization/__pycache__/config.cpython-38.pyc
new file mode 100644
index 0000000..944263b
Binary files /dev/null and b/Colorful_Image_Colorization/__pycache__/config.cpython-38.pyc differ
diff --git a/Colorful_Image_Colorization/__pycache__/model.cpython-38.pyc b/Colorful_Image_Colorization/__pycache__/model.cpython-38.pyc
new file mode 100644
index 0000000..697d427
Binary files /dev/null and b/Colorful_Image_Colorization/__pycache__/model.cpython-38.pyc differ
diff --git a/Colorful_Image_Colorization/config.py b/Colorful_Image_Colorization/config.py
new file mode 100644
index 0000000..d518840
--- /dev/null
+++ b/Colorful_Image_Colorization/config.py
@@ -0,0 +1,14 @@
+img_rows, img_cols = 256, 256
+channel = 3
+batch_size = 32
+epochs = 10000
+patience = 50
+num_train_samples = 529202
+num_valid_samples = 4268
+num_classes = 313
+kernel = 3
+weight_decay = 1e-3
+epsilon = 1e-8
+nb_neighbors = 5
+# temperature parameter T
+T = 0.38
diff --git a/Colorful_Image_Colorization/model.06-2.5489.hdf5 b/Colorful_Image_Colorization/model.06-2.5489.hdf5
new file mode 100644
index 0000000..6353dbc
Binary files /dev/null and b/Colorful_Image_Colorization/model.06-2.5489.hdf5 differ
diff --git a/Colorful_Image_Colorization/model.py b/Colorful_Image_Colorization/model.py
new file mode 100644
index 0000000..712d3ee
--- /dev/null
+++ b/Colorful_Image_Colorization/model.py
@@ -0,0 +1,92 @@
+import keras.backend as K
+import tensorflow as tf
+from keras.layers import Input, Conv2D, BatchNormalization, UpSampling2D
+from keras.models import Model
+from keras.regularizers import l2
+
+from Colorful_Image_Colorization.config import img_rows, img_cols, num_classes, kernel
+
+l2_reg = l2(1e-3)
+
+
+def build_model():
+    input_tensor = Input(shape=(img_rows, img_cols, 1))
+    x = Conv2D(64, (kernel, kernel), activation='relu', padding='same', name='conv1_1', kernel_initializer="he_normal",
+               kernel_regularizer=l2_reg)(input_tensor)
+    x = Conv2D(64, (kernel, kernel), activation='relu', padding='same', name='conv1_2', kernel_initializer="he_normal",
+               kernel_regularizer=l2_reg, strides=(2, 2))(x)
+    x = BatchNormalization()(x)
+
+    x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv2_1', kernel_initializer="he_normal",
+               kernel_regularizer=l2_reg)(x)
+    x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv2_2', kernel_initializer="he_normal",
+               kernel_regularizer=l2_reg,
+               strides=(2, 2))(x)
+    x = BatchNormalization()(x)
+
+    x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv3_1',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv3_2',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv3_3', kernel_initializer="he_normal",
+               strides=(2, 2))(x)
+    x = BatchNormalization()(x)
+
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', name='conv4_1',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', name='conv4_2',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', name='conv4_3',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = BatchNormalization()(x)
+
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv5_1',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv5_2',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv5_3',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = BatchNormalization()(x)
+
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv6_1',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv6_2',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(512, (kernel, kernel), activation='relu', padding='same', dilation_rate=2, name='conv6_3',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = BatchNormalization()(x)
+
+    x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv7_1',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv7_2',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(256, (kernel, kernel), activation='relu', padding='same', name='conv7_3',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = BatchNormalization()(x)
+
+    x = UpSampling2D(size=(2, 2))(x)
+    x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv8_1',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv8_2',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = Conv2D(128, (kernel, kernel), activation='relu', padding='same', name='conv8_3',
+               kernel_initializer="he_normal", kernel_regularizer=l2_reg)(x)
+    x = BatchNormalization()(x)
+
+    outputs = Conv2D(num_classes, (1, 1), activation='softmax', padding='same', name='pred')(x)
+
+    model = Model(inputs=input_tensor, outputs=outputs, name="ColorNet")
+    return model
+
+
+if __name__ == '__main__':
+    with tf.device("/cpu:0"):
+        encoder_decoder = build_model()
+    print(encoder_decoder.summary())
+    #plot_model(encoder_decoder, to_file='encoder_decoder.svg', show_layer_names=True, show_shapes=True)
+
+    parallel_model = encoder_decoder
+    print(parallel_model.summary())
+    #plot_model(parallel_model, to_file='parallel_model.svg', show_layer_names=True, show_shapes=True)
+
+    K.clear_session()
diff --git a/Colorful_Image_Colorization/pts_in_hull.npy b/Colorful_Image_Colorization/pts_in_hull.npy
new file mode 100644
index 0000000..239ff11
Binary files /dev/null and b/Colorful_Image_Colorization/pts_in_hull.npy differ
diff --git a/Encoder.py b/Encoder.py
index fc75cc1..e95f393 100644
--- a/Encoder.py
+++ b/Encoder.py
@@ -41,5 +41,4 @@ class Encoder(tf.keras.Model): # <-- Needed to make parameters trainable and to
         #print(x.shape)
         #print("-------------")
         #exit()
-        return x
-    
+        return x
\ No newline at end of file
diff --git a/Plots/ColoredImages_1.png b/Plots/ColoredImages_1.png
new file mode 100644
index 0000000..5cbbfa7
Binary files /dev/null and b/Plots/ColoredImages_1.png differ
diff --git a/Plots/ColoredImages_2.png b/Plots/ColoredImages_2.png
new file mode 100644
index 0000000..14039c9
Binary files /dev/null and b/Plots/ColoredImages_2.png differ
diff --git a/Plots/ColoredImages_compareModels_1.png b/Plots/ColoredImages_compareModels_1.png
new file mode 100644
index 0000000..acf6898
Binary files /dev/null and b/Plots/ColoredImages_compareModels_1.png differ
diff --git a/Plots/ColoredImages_compareModels_2.png b/Plots/ColoredImages_compareModels_2.png
new file mode 100644
index 0000000..e57843a
Binary files /dev/null and b/Plots/ColoredImages_compareModels_2.png differ
diff --git a/Plots/ColorizedImagesLossPlot_comparedModels_1.png b/Plots/ColorizedImagesLossPlot_comparedModels_1.png
new file mode 100644
index 0000000..7241378
Binary files /dev/null and b/Plots/ColorizedImagesLossPlot_comparedModels_1.png differ
diff --git a/Plots/ColorizedImagesLossPlot_comparedModels_2.png b/Plots/ColorizedImagesLossPlot_comparedModels_2.png
new file mode 100644
index 0000000..ea41e1b
Binary files /dev/null and b/Plots/ColorizedImagesLossPlot_comparedModels_2.png differ
diff --git a/Plots/CreatePlot_TrainTestLoss.py b/Plots/CreatePlot_TrainTestLoss.py
new file mode 100644
index 0000000..c36dfff
--- /dev/null
+++ b/Plots/CreatePlot_TrainTestLoss.py
@@ -0,0 +1,32 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+def main():
+
+    df_train_loss = pd.read_csv('run-.-tag-Train loss.csv', sep=',')
+    train_loss = df_train_loss["Value"]
+
+    df_test_loss = pd.read_csv('run-.-tag-Test loss.csv', sep=',')
+    test_loss = df_test_loss["Value"]
+
+    x = np.arange(len(train_loss))
+    
+    plt.plot(x, train_loss, label="Train loss", color="r")
+    plt.plot(x, test_loss, label="Test loss", color="b")
+
+    plt.legend()
+    plt.grid(True)
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+
+    plt.savefig("TrainTestLoss.png")
+    plt.show()
+    
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt received")
\ No newline at end of file
diff --git a/Plots/CreatePlot_compareModels.py b/Plots/CreatePlot_compareModels.py
new file mode 100644
index 0000000..4e68e50
--- /dev/null
+++ b/Plots/CreatePlot_compareModels.py
@@ -0,0 +1,182 @@
+import tensorflow as tf
+import tensorflow_datasets as tfds
+import matplotlib.pyplot as plt
+
+
+import sys
+sys.path.append("..")
+
+from Autoencoder import Autoencoder
+from Training import prepare_data, getRGB
+
+import numpy as np
+import os
+
+
+#from Training import prepare_data, getRGB
+
+from Colorful_Image_Colorization.model import build_model
+from Colorful_Image_Colorization.config import img_rows, img_cols
+from Colorful_Image_Colorization.config import nb_neighbors, T, epsilon
+import cv2 as cv
+
+def main():
+
+    # Create Imagenet
+    labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
+    imagenet_labels = np.array(open(labels_path).read().splitlines())
+
+    data_dir = '/home/timwitte/Downloads/'
+    write_dir = '../imagenet'
+
+    # Construct a tf.data.Dataset
+    download_config = tfds.download.DownloadConfig(
+                        extract_dir=os.path.join(write_dir, 'extracted'),
+                        manual_dir=data_dir
+                    )
+    download_and_prepare_kwargs = {
+        'download_dir': os.path.join(write_dir, 'downloaded'),
+        'download_config': download_config,
+    }
+
+    train_dataset, test_dataset= tfds.load('imagenet2012', 
+                data_dir=os.path.join(write_dir, 'data'),         
+                split=['train', 'validation'], 
+                shuffle_files=True, 
+                download=True,
+                as_supervised=True,
+                download_and_prepare_kwargs=download_and_prepare_kwargs)
+    
+    test_dataset = test_dataset.take(32).apply(prepare_data)
+
+    # Load our model
+    model_our = Autoencoder()
+    model_our.build((1, 256, 256, 1)) # need a batch size
+    model_our.load_weights("../saved_models/trainied_weights_epoch_12")
+
+    # Load model to compare
+    model_weights_path = '../Colorful_Image_Colorization/model.06-2.5489.hdf5'
+    model_toCompare = build_model()
+    model_toCompare.load_weights(model_weights_path)
+
+    loss_function = tf.keras.losses.MeanSquaredError()
+
+    for img_L, img_AB_orginal in test_dataset.take(1):
+
+        img_rgb_orginal = getRGB(img_L, img_AB_orginal)
+
+        img_AB_reconstructed_our = model_our.predict(img_L.numpy())
+        img_rgb_reconstructed_our = getRGB(img_L, img_AB_reconstructed_our)
+
+        NUM_IMGS = 5
+        fig, axs = plt.subplots(NUM_IMGS, 4)
+
+        axs[0, 0].set_title("Input", fontsize=30)
+        axs[0, 1].set_title("Richard Zhang $\it{et\ al.}$", fontsize=30,)
+        axs[0, 2].set_title("Ours", fontsize=30)
+        axs[0, 3].set_title("Ground Truth", fontsize=30)
+        losses1 = []
+        losses2 = []
+        for i in range(NUM_IMGS):
+
+            img_AB_reconstructed_toCompare = getABFromModel(model_toCompare, img_L[i].numpy())
+            img_rgb_reconstructed_toCompare = getRGB(img_L[i], img_AB_reconstructed_toCompare, batch_mode=False)
+
+            axs[i, 0].imshow(img_L[i], cmap="gray")
+            axs[i, 0].set_axis_off()
+
+            axs[i, 1].imshow(img_rgb_reconstructed_toCompare)
+            axs[i, 1].set_axis_off()
+
+            axs[i, 2].imshow(img_rgb_reconstructed_our[i])
+            axs[i, 2].set_axis_off()
+
+            axs[i, 3].imshow(img_rgb_orginal[i])
+            axs[i, 3].set_axis_off()
+
+            loss_our = loss_function(img_rgb_orginal[i], img_rgb_reconstructed_our[i])
+            loss_toCompare = loss_function(img_rgb_orginal[i], img_rgb_reconstructed_our)
+
+            losses1.append(loss_our)
+            losses2.append(loss_toCompare)
+       
+
+        plt.tight_layout()
+
+        fig.set_size_inches(20, 25)
+        fig.savefig("ColoredImages_compareModels.png")
+
+        # Reset plot
+        plt.clf()
+        plt.cla()
+        fig = plt.figure()
+
+        # Create bar plot
+        x_axis = np.arange(NUM_IMGS)
+        width = 0.2
+        plt.bar(x_axis - width/2., losses2, width=width/2, label = "Richard Zhang $\it{et\ al.}$")
+        plt.bar(x_axis - width/2. + 1/float(2)*width, losses1, width=width/2, label = 'Ours')
+       
+        
+        plt.xticks(x_axis,[f"No. {i}" for i in range(NUM_IMGS)])
+        
+        plt.title("Loss of colorized images")
+        plt.xlabel("Image")
+        plt.ylabel("Loss")
+
+        plt.legend()
+        plt.tight_layout()
+        plt.savefig("ColorizedImagesLossPlot_comparedModels.png")
+
+ 
+
+
+def getABFromModel(model, grey_img):
+    # code taken from https://github.com/foamliu/Colorful-Image-Colorization/blob/master/demo.py
+    q_ab = np.load("../Colorful_Image_Colorization/pts_in_hull.npy")
+    nb_q = q_ab.shape[0]
+
+    grey_img = np.expand_dims(grey_img, axis=0)
+    
+    X_colorized = model.predict((grey_img+1)/2)
+
+    
+    h, w = img_rows // 4, img_cols // 4
+    X_colorized = X_colorized.reshape((h * w, nb_q))
+
+    # Reweight probas
+    X_colorized = np.exp(np.log(X_colorized + epsilon) / T)
+    X_colorized = X_colorized / np.sum(X_colorized, 1)[:, np.newaxis]
+
+    # Reweighted
+    q_a = q_ab[:, 0].reshape((1, 313))
+    q_b = q_ab[:, 1].reshape((1, 313))
+
+    X_a = np.sum(X_colorized * q_a, 1).reshape((h, w))
+    X_b = np.sum(X_colorized * q_b, 1).reshape((h, w))
+
+    X_a = cv.resize(X_a, (img_rows, img_cols), cv.INTER_CUBIC)
+    X_b = cv.resize(X_b, (img_rows, img_cols), cv.INTER_CUBIC)
+
+    # Before: -90 <=a<= 100, -110 <=b<= 110
+    # After: 38 <=a<= 228, 18 <=b<= 238
+    X_a = X_a + 128
+    X_b = X_b + 128
+  
+    out_lab = np.zeros((256, 256, 2), dtype=np.float32)
+    grey_img = np.reshape(grey_img, newshape=(256,256))
+  
+
+    out_lab[:, :, 0] = X_a
+    out_lab[:, :, 1] = X_b
+ 
+    out_lab[:, :, 0] = -1.0 + 2*(out_lab[:, :, 0] - 38.0)/190
+    out_lab[:, :, 1] = -1.0 + 2*(out_lab[:, :, 1] - 20.0)/203
+
+    return out_lab
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt received")
\ No newline at end of file
diff --git a/Plots/CreatePlot_showImages.py b/Plots/CreatePlot_showImages.py
new file mode 100644
index 0000000..76d1c09
--- /dev/null
+++ b/Plots/CreatePlot_showImages.py
@@ -0,0 +1,84 @@
+import tensorflow as tf
+import tensorflow_datasets as tfds
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+
+import sys
+sys.path.append("..")
+from Autoencoder import Autoencoder
+from Training import prepare_data, getRGB
+
+def main():
+
+    labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
+    imagenet_labels = np.array(open(labels_path).read().splitlines())
+
+    data_dir = '/home/timwitte/Downloads/'
+    write_dir = '../imagenet'
+
+    # Construct a tf.data.Dataset
+    download_config = tfds.download.DownloadConfig(
+                        extract_dir=os.path.join(write_dir, 'extracted'),
+                        manual_dir=data_dir
+                    )
+    download_and_prepare_kwargs = {
+        'download_dir': os.path.join(write_dir, 'downloaded'),
+        'download_config': download_config,
+    }
+
+    train_dataset, test_dataset= tfds.load('imagenet2012', 
+                data_dir=os.path.join(write_dir, 'data'),         
+                split=['train', 'validation'], 
+                shuffle_files=True, 
+                download=True,
+                as_supervised=True,
+                download_and_prepare_kwargs=download_and_prepare_kwargs)
+    
+    test_dataset = test_dataset.take(32).apply(prepare_data)
+
+    autoencoder = Autoencoder()
+
+    autoencoder.build((1, 256, 256, 1)) # need a batch size
+    autoencoder.load_weights("../saved_models/trainied_weights_epoch_12")
+    autoencoder.summary()
+    
+    autoencoder.encoder.summary()
+    autoencoder.decoder.summary()
+
+    for img_L, img_AB_orginal in test_dataset.take(1):
+
+        img_AB_reconstructed = autoencoder(img_L)
+
+        img_rgb_orginal = getRGB(img_L, img_AB_orginal)
+        img_rgb_reconstructed = getRGB(img_L, img_AB_reconstructed)
+        
+        NUM_IMGS = 5
+        fig, axs = plt.subplots(NUM_IMGS, 3)
+
+        axs[0, 0].set_title("Input", fontsize=30)
+        axs[0, 1].set_title("Output", fontsize=30)
+        axs[0, 2].set_title("Ground Truth", fontsize=30)
+
+        for i in range(NUM_IMGS):
+            
+            axs[i, 0].imshow(img_L[i], cmap="gray")
+            axs[i, 0].set_axis_off()
+
+            axs[i, 1].imshow(img_rgb_reconstructed[i])
+            axs[i, 1].set_axis_off()
+
+            axs[i, 2].imshow(img_rgb_orginal[i])
+            axs[i, 2].set_axis_off()
+        
+        plt.tight_layout()
+
+        fig.set_size_inches(15, 25)
+        fig.savefig("ColoredImages.png")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt received")
\ No newline at end of file
diff --git a/Plots/Layers/CreatePlot_layerStructure.py b/Plots/Layers/CreatePlot_layerStructure.py
new file mode 100644
index 0000000..d125a6f
--- /dev/null
+++ b/Plots/Layers/CreatePlot_layerStructure.py
@@ -0,0 +1,33 @@
+import tensorflow as tf
+
+from EncoderLayers import *
+from DecoderLayers import *
+
+import sys
+sys.path.append("../..")
+
+from Colorful_Image_Colorization.model import *
+
+def main():
+
+    encoder_layers = EncoderLayers()
+    decoder_layers = DecoderLayers()
+
+    inputs = tf.keras.Input(shape=(256,256, 1), name="Grey image") 
+    encoder = tf.keras.Model(inputs=[inputs],outputs=encoder_layers.call(inputs))
+
+    embedding = tf.keras.Input(shape=(32,32, 3), name="Embedding") 
+    decoder = tf.keras.Model(inputs=[embedding],outputs=decoder_layers.call(embedding))
+
+    tf.keras.utils.plot_model(encoder,show_shapes=True, show_layer_names=True, to_file="EncoderLayer.png")
+    tf.keras.utils.plot_model(decoder,show_shapes=True, show_layer_names=True, to_file="DecoderLayer.png")
+    
+    ModelToCompare_layers = build_model()
+    modelToCompare = tf.keras.Model(inputs=[inputs],outputs=ModelToCompare_layers.call(inputs))
+    tf.keras.utils.plot_model(modelToCompare,show_shapes=True, show_layer_names=True, to_file="ModelToCompare.png")
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt received")
\ No newline at end of file
diff --git a/Plots/Layers/DecoderLayer.png b/Plots/Layers/DecoderLayer.png
new file mode 100644
index 0000000..d5a235f
Binary files /dev/null and b/Plots/Layers/DecoderLayer.png differ
diff --git a/Plots/Layers/DecoderLayers.py b/Plots/Layers/DecoderLayers.py
new file mode 100644
index 0000000..86eccff
--- /dev/null
+++ b/Plots/Layers/DecoderLayers.py
@@ -0,0 +1,31 @@
+import tensorflow as tf
+
+class DecoderLayers(tf.keras.Model):
+    def __init__(self):
+        super(DecoderLayers, self).__init__()
+
+        self.layer_list = [
+            tf.keras.layers.Conv2DTranspose(105, kernel_size=(3,3), strides=2, padding='same', name="Conv2D_Trans_0"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_0"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_0"),
+
+            tf.keras.layers.Conv2DTranspose(90, kernel_size=(3,3), strides=2, padding='same', name="Conv2D_Trans_1"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_1"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_1"),
+
+            tf.keras.layers.Conv2DTranspose(75, kernel_size=(3,3), strides=2, padding='same', name="Conv2D_Trans_2"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_2"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_2"),
+
+            # bottleneck to RGB
+
+            tf.keras.layers.Conv2DTranspose(2, kernel_size=(1,1), strides=1, padding='same', name="Conv2D_Trans_3"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_3"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_3"),
+        ]
+        
+    
+    def call(self, x):
+        for layer in self.layer_list:
+            x = layer(x)
+        return x
\ No newline at end of file
diff --git a/Plots/Layers/EncoderLayer.png b/Plots/Layers/EncoderLayer.png
new file mode 100644
index 0000000..2543263
Binary files /dev/null and b/Plots/Layers/EncoderLayer.png differ
diff --git a/Plots/Layers/EncoderLayers.py b/Plots/Layers/EncoderLayers.py
new file mode 100644
index 0000000..742ca82
--- /dev/null
+++ b/Plots/Layers/EncoderLayers.py
@@ -0,0 +1,29 @@
+import tensorflow as tf
+
+class EncoderLayers(tf.keras.Model):
+    def __init__(self):
+        super(EncoderLayers, self).__init__()
+
+        self.layer_list = [
+            tf.keras.layers.Conv2D(75, kernel_size=(3, 3), strides=2, padding='same', name="Conv2D_0"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_0"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_0"),
+
+            tf.keras.layers.Conv2D(90, kernel_size=(3, 3), strides=2, padding='same', name="Conv2D_1"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_1"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_1"),
+
+            tf.keras.layers.Conv2D(105, kernel_size=(3, 3), strides=2, padding='same',name="Conv2D_2"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_2"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_2"),
+
+            tf.keras.layers.Conv2D(3, kernel_size=(1, 1), strides=1, padding='same', name="Conv2D_3"),
+            tf.keras.layers.BatchNormalization(name="BatchNormalization_3"),
+            tf.keras.layers.Activation(tf.nn.tanh, name="tanh_3"),
+        ]
+        
+    
+    def call(self, x):
+        for layer in self.layer_list:
+            x = layer(x)
+        return x
\ No newline at end of file
diff --git a/Plots/Layers/ModelToCompare.png b/Plots/Layers/ModelToCompare.png
new file mode 100644
index 0000000..1c05847
Binary files /dev/null and b/Plots/Layers/ModelToCompare.png differ
diff --git a/Plots/Layers/__pycache__/DecoderLayers.cpython-38.pyc b/Plots/Layers/__pycache__/DecoderLayers.cpython-38.pyc
new file mode 100644
index 0000000..1cd898a
Binary files /dev/null and b/Plots/Layers/__pycache__/DecoderLayers.cpython-38.pyc differ
diff --git a/Plots/Layers/__pycache__/EncoderLayers.cpython-38.pyc b/Plots/Layers/__pycache__/EncoderLayers.cpython-38.pyc
new file mode 100644
index 0000000..f6f88fc
Binary files /dev/null and b/Plots/Layers/__pycache__/EncoderLayers.cpython-38.pyc differ
diff --git a/Plots/Layers/__pycache__/ModelToCompare.cpython-38.pyc b/Plots/Layers/__pycache__/ModelToCompare.cpython-38.pyc
new file mode 100644
index 0000000..efa16f3
Binary files /dev/null and b/Plots/Layers/__pycache__/ModelToCompare.cpython-38.pyc differ
diff --git a/Plots/TrainTestLossPlot.png b/Plots/TrainTestLossPlot.png
new file mode 100644
index 0000000..2bbed3a
Binary files /dev/null and b/Plots/TrainTestLossPlot.png differ
diff --git a/Plots/__pycache__/DecoderLayers.cpython-38.pyc b/Plots/__pycache__/DecoderLayers.cpython-38.pyc
new file mode 100644
index 0000000..f2ef30d
Binary files /dev/null and b/Plots/__pycache__/DecoderLayers.cpython-38.pyc differ
diff --git a/Plots/__pycache__/EncoderLayers.cpython-38.pyc b/Plots/__pycache__/EncoderLayers.cpython-38.pyc
new file mode 100644
index 0000000..8438a95
Binary files /dev/null and b/Plots/__pycache__/EncoderLayers.cpython-38.pyc differ
diff --git a/Plots/run-.-tag-Test loss.csv b/Plots/run-.-tag-Test loss.csv
new file mode 100644
index 0000000..5d43b83
--- /dev/null
+++ b/Plots/run-.-tag-Test loss.csv	
@@ -0,0 +1,15 @@
+Wall time,Step,Value
+1647118261.250931,0,0.016809336841106415
+1647122922.100936,1,0.013759356923401356
+1647127578.618303,2,0.01362006925046444
+1647132237.117106,3,0.014061697758734226
+1647136894.337419,4,0.013472857885062695
+1647141548.038196,5,0.01342787966132164
+1647146209.292402,6,0.013368184678256512
+1647150861.434495,7,0.013420150615274906
+1647155517.411057,8,0.01329082902520895
+1647160207.188101,9,0.013379388488829136
+1647164926.587916,10,0.013525118120014668
+1647169761.401568,11,0.01334059052169323
+1647174430.458649,12,0.013532023876905441
+1647179221.740573,13,0.01326887309551239
diff --git a/Plots/run-.-tag-Train loss.csv b/Plots/run-.-tag-Train loss.csv
new file mode 100644
index 0000000..158be26
--- /dev/null
+++ b/Plots/run-.-tag-Train loss.csv	
@@ -0,0 +1,15 @@
+Wall time,Step,Value
+1647118226.379471,0,0.01686934195458889
+1647122884.759068,1,0.016268473118543625
+1647127535.214221,2,0.013647115789353848
+1647132192.879982,3,0.013552550226449966
+1647136850.629965,4,0.01349611859768629
+1647141510.180662,5,0.013455081731081009
+1647146165.61258,6,0.01342522632330656
+1647150823.542946,7,0.013399843126535416
+1647155473.882963,8,0.013378930278122425
+1647160163.770788,9,0.013358119875192642
+1647164886.509832,10,0.013342463411390781
+1647169721.204018,11,0.013329868204891682
+1647174386.649681,12,0.013316545635461807
+1647179177.404204,13,0.013304967433214188
diff --git a/Training.py b/Training.py
new file mode 100644
index 0000000..90cf5cf
--- /dev/null
+++ b/Training.py
@@ -0,0 +1,175 @@
+from sklearn.utils import shuffle
+import tensorflow as tf
+import tensorflow_datasets as tfds
+import matplotlib.pyplot as plt
+import numpy as np
+import tqdm
+
+from Decoder import *
+
+import os
+
+from Autoencoder import Autoencoder
+import tensorflow_io as tfio
+
+def getRGB(L, AB, batch_mode=True):
+    # Remove normalization
+    L = (L + 1)*50
+    AB = ((AB - 1)*255/2)+128
+
+    if batch_mode:
+        L = tf.reshape(L, (32, 256,256,1))
+        LAB = tf.concat([L, AB], 3)
+    else:
+        L = tf.reshape(L, (256,256,1))
+        LAB = tf.concat([L, AB], 2)
+    rgb = tfio.experimental.color.lab_to_rgb(LAB)
+
+    return rgb
+
+def main():
+    
+
+    labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
+    imagenet_labels = np.array(open(labels_path).read().splitlines())
+
+    data_dir = '/home/timwitte/Downloads/'
+    write_dir = './imagenet'
+
+    # Construct a tf.data.Dataset
+    download_config = tfds.download.DownloadConfig(
+                        extract_dir=os.path.join(write_dir, 'extracted'),
+                        manual_dir=data_dir
+                    )
+    download_and_prepare_kwargs = {
+        'download_dir': os.path.join(write_dir, 'downloaded'),
+        'download_config': download_config,
+    }
+
+    train_dataset, test_dataset= tfds.load('imagenet2012', 
+                data_dir=os.path.join(write_dir, 'data'),         
+                split=['train', 'validation'], 
+                shuffle_files=True, 
+                download=True,
+                as_supervised=True,
+                download_and_prepare_kwargs=download_and_prepare_kwargs)
+    
+    train_dataset = train_dataset.apply(prepare_data)
+    test_dataset = test_dataset.apply(prepare_data).take(500) # take 500 batches
+    
+
+    # for L, AB in train_dataset.take(1):
+        
+    #     print(L.shape)
+    #     print(AB.shape)
+
+    #     print(np.min(L[0]))
+    #     print(np.max(L[0]))
+    #     print("######################")
+    #     print(np.min(AB[0]))
+    #     print(np.max(AB[0]))
+
+    #     rgb = getRGB(L, AB)
+        
+    #     plt.imshow(rgb[0])
+    #     plt.show()
+   
+    #     exit()
+
+    autoencoder = Autoencoder()
+    num_epochs = 75
+
+    file_path = "test_logs/test"
+    summary_writer = tf.summary.create_file_writer(file_path)
+
+    for img_L_tensorBoard, img_AB_tensorBoard in test_dataset.take(1):
+        pass
+
+    with summary_writer.as_default():
+
+        tf.summary.image(name="grey_images",data = img_L_tensorBoard, step=0, max_outputs=32)
+        img_RBG = getRGB(img_L_tensorBoard, img_AB_tensorBoard)
+        tf.summary.image(name="colored_images",data = img_RBG, step=0, max_outputs=32)
+        
+        imgs = autoencoder(img_L_tensorBoard)
+        tf.summary.image(name="recolored_images",data = imgs, step=0, max_outputs=32)
+
+        autoencoder.summary()
+      
+        train_loss = autoencoder.test(train_dataset.take(100))
+        
+        tf.summary.scalar(name="Train loss", data=train_loss, step=0)
+
+        test_loss = autoencoder.test(test_dataset)
+        tf.summary.scalar(name="Test loss", data=test_loss, step=0)
+
+
+        for epoch in range(num_epochs):
+            
+            print(f"Epoch {epoch}")
+           
+        
+            for img_L, img_AB in tqdm.tqdm(train_dataset,position=0, leave=True): 
+                autoencoder.train_step(img_L, img_AB)
+               
+             
+            tf.summary.scalar(name="Train loss", data=autoencoder.metric_mean.result(), step=epoch+1)
+            autoencoder.metric_mean.reset_states()
+
+            test_loss = autoencoder.test(test_dataset)
+            tf.summary.scalar(name="Test loss", data=test_loss, step=epoch+1)
+
+            img_AB = autoencoder(img_L_tensorBoard)
+
+            img_RBG = getRGB(img_L_tensorBoard, img_AB)
+
+            tf.summary.image(name="recolored_images",data = img_RBG, step=epoch + 1, max_outputs=32)
+            
+            # save model
+            autoencoder.save_weights(f"./saved_models/trainied_weights_epoch_{epoch}", save_format="tf")
+
+def prepare_data(data):
+
+    # Remove label 
+    data = data.map(lambda img, label: img )
+
+    # resize
+    data = data.map(lambda img: tf.image.resize(img, [256,256]) )
+
+    #convert data from uint8 to float32
+    data = data.map(lambda img: tf.cast(img, tf.float32) )
+
+    # tfio.experimental.color.rgb_to_lab expects its input to be a float normalized between 0 and 1.
+    data = data.map(lambda img: (img/255.) )
+    data = data.map(lambda img: tfio.experimental.color.rgb_to_lab(img) )
+
+    # X = L channel
+    # Y = (A,B) channel
+    data = data.map(lambda img: (img[:, :, 0], tf.stack([img[:, :, 1], img[:, :, 2]], axis=2)))
+
+    # Reshape R channel -> grey
+    data = data.map(lambda L, AB: ( tf.reshape(L, shape=(256,256,1)) , AB))  
+    
+    # Normalize between [-1, 1]
+    data = data.map(lambda L, AB: ( (L/50.0) - 1., 1 + (2*(AB - 128)/255) ))  
+    
+    # add gray scaled image
+    #data = data.map(lambda img: (tf.image.rgb_to_grayscale(img), img))
+
+    #cache this progress in memory, as there is no need to redo it; it is deterministic after all
+    #data = data.cache("cachefile")
+
+    #shuffle, batch, prefetch
+    data = data.shuffle(7000) 
+    data = data.batch(32)
+
+    AUTOTUNE = tf.data.AUTOTUNE
+    data = data.prefetch(AUTOTUNE)
+    #return preprocessed dataset
+    return data
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("KeyboardInterrupt received")
\ No newline at end of file
diff --git a/imagenet/PLACE YOUR IMAGENET HERE b/imagenet/PLACE YOUR IMAGENET HERE
new file mode 100644
index 0000000..e69de29
diff --git a/saved_models/.DS_Store b/saved_models/.DS_Store
deleted file mode 100644
index 7deef98..0000000
Binary files a/saved_models/.DS_Store and /dev/null differ
diff --git a/saved_models/trainied_weights_epoch_0.data-00000-of-00001 b/saved_models/trainied_weights_epoch_0.data-00000-of-00001
new file mode 100644
index 0000000..442f321
Binary files /dev/null and b/saved_models/trainied_weights_epoch_0.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_0.index b/saved_models/trainied_weights_epoch_0.index
new file mode 100644
index 0000000..93c4f64
Binary files /dev/null and b/saved_models/trainied_weights_epoch_0.index differ
diff --git a/saved_models/trainied_weights_epoch_1.data-00000-of-00001 b/saved_models/trainied_weights_epoch_1.data-00000-of-00001
index 82f95cd..0f2be2d 100644
Binary files a/saved_models/trainied_weights_epoch_1.data-00000-of-00001 and b/saved_models/trainied_weights_epoch_1.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_1.index b/saved_models/trainied_weights_epoch_1.index
index b541cdb..f8a1873 100644
Binary files a/saved_models/trainied_weights_epoch_1.index and b/saved_models/trainied_weights_epoch_1.index differ
diff --git a/saved_models/trainied_weights_epoch_10.data-00000-of-00001 b/saved_models/trainied_weights_epoch_10.data-00000-of-00001
new file mode 100644
index 0000000..908b1f2
Binary files /dev/null and b/saved_models/trainied_weights_epoch_10.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_10.index b/saved_models/trainied_weights_epoch_10.index
new file mode 100644
index 0000000..af9d171
Binary files /dev/null and b/saved_models/trainied_weights_epoch_10.index differ
diff --git a/saved_models/trainied_weights_epoch_11.data-00000-of-00001 b/saved_models/trainied_weights_epoch_11.data-00000-of-00001
new file mode 100644
index 0000000..ea9f96f
Binary files /dev/null and b/saved_models/trainied_weights_epoch_11.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_11.index b/saved_models/trainied_weights_epoch_11.index
new file mode 100644
index 0000000..6432076
Binary files /dev/null and b/saved_models/trainied_weights_epoch_11.index differ
diff --git a/saved_models/trainied_weights_epoch_19.data-00000-of-00001 b/saved_models/trainied_weights_epoch_19.data-00000-of-00001
deleted file mode 100644
index a25bc5d..0000000
Binary files a/saved_models/trainied_weights_epoch_19.data-00000-of-00001 and /dev/null differ
diff --git a/saved_models/trainied_weights_epoch_19.index b/saved_models/trainied_weights_epoch_19.index
deleted file mode 100644
index a59234c..0000000
Binary files a/saved_models/trainied_weights_epoch_19.index and /dev/null differ
diff --git a/saved_models/trainied_weights_epoch_2.data-00000-of-00001 b/saved_models/trainied_weights_epoch_2.data-00000-of-00001
new file mode 100644
index 0000000..6dabb84
Binary files /dev/null and b/saved_models/trainied_weights_epoch_2.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_2.index b/saved_models/trainied_weights_epoch_2.index
new file mode 100644
index 0000000..ac71b7b
Binary files /dev/null and b/saved_models/trainied_weights_epoch_2.index differ
diff --git a/saved_models/trainied_weights_epoch_3.data-00000-of-00001 b/saved_models/trainied_weights_epoch_3.data-00000-of-00001
new file mode 100644
index 0000000..2a15ece
Binary files /dev/null and b/saved_models/trainied_weights_epoch_3.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_3.index b/saved_models/trainied_weights_epoch_3.index
new file mode 100644
index 0000000..ac14249
Binary files /dev/null and b/saved_models/trainied_weights_epoch_3.index differ
diff --git a/saved_models/trainied_weights_epoch_4.data-00000-of-00001 b/saved_models/trainied_weights_epoch_4.data-00000-of-00001
new file mode 100644
index 0000000..d9637db
Binary files /dev/null and b/saved_models/trainied_weights_epoch_4.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_4.index b/saved_models/trainied_weights_epoch_4.index
new file mode 100644
index 0000000..449ba6c
Binary files /dev/null and b/saved_models/trainied_weights_epoch_4.index differ
diff --git a/saved_models/trainied_weights_epoch_5.data-00000-of-00001 b/saved_models/trainied_weights_epoch_5.data-00000-of-00001
new file mode 100644
index 0000000..7dfa708
Binary files /dev/null and b/saved_models/trainied_weights_epoch_5.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_5.index b/saved_models/trainied_weights_epoch_5.index
new file mode 100644
index 0000000..09c5068
Binary files /dev/null and b/saved_models/trainied_weights_epoch_5.index differ
diff --git a/saved_models/trainied_weights_epoch_6.data-00000-of-00001 b/saved_models/trainied_weights_epoch_6.data-00000-of-00001
index 8f371ba..13e1b2e 100644
Binary files a/saved_models/trainied_weights_epoch_6.data-00000-of-00001 and b/saved_models/trainied_weights_epoch_6.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_6.index b/saved_models/trainied_weights_epoch_6.index
index 0fd28f6..92f5bbd 100644
Binary files a/saved_models/trainied_weights_epoch_6.index and b/saved_models/trainied_weights_epoch_6.index differ
diff --git a/saved_models/trainied_weights_epoch_7.data-00000-of-00001 b/saved_models/trainied_weights_epoch_7.data-00000-of-00001
new file mode 100644
index 0000000..35aff74
Binary files /dev/null and b/saved_models/trainied_weights_epoch_7.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_7.index b/saved_models/trainied_weights_epoch_7.index
new file mode 100644
index 0000000..f8dbd9b
Binary files /dev/null and b/saved_models/trainied_weights_epoch_7.index differ
diff --git a/saved_models/trainied_weights_epoch_8.data-00000-of-00001 b/saved_models/trainied_weights_epoch_8.data-00000-of-00001
new file mode 100644
index 0000000..f8691a4
Binary files /dev/null and b/saved_models/trainied_weights_epoch_8.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_8.index b/saved_models/trainied_weights_epoch_8.index
new file mode 100644
index 0000000..ca33937
Binary files /dev/null and b/saved_models/trainied_weights_epoch_8.index differ
diff --git a/saved_models/trainied_weights_epoch_9.data-00000-of-00001 b/saved_models/trainied_weights_epoch_9.data-00000-of-00001
new file mode 100644
index 0000000..7c8522d
Binary files /dev/null and b/saved_models/trainied_weights_epoch_9.data-00000-of-00001 differ
diff --git a/saved_models/trainied_weights_epoch_9.index b/saved_models/trainied_weights_epoch_9.index
new file mode 100644
index 0000000..294aa8b
Binary files /dev/null and b/saved_models/trainied_weights_epoch_9.index differ
diff --git a/test_logs/test/events.out.tfevents.1647118190.tim-PC.3290.0.v2 b/test_logs/test/events.out.tfevents.1647118190.tim-PC.3290.0.v2
new file mode 100644
index 0000000..b465893
Binary files /dev/null and b/test_logs/test/events.out.tfevents.1647118190.tim-PC.3290.0.v2 differ