UNET

Author: PPSantos

src/helpers.py

@@ -3,25 +3,6 @@
 import matplotlib.pyplot as plt
 from skimage import color, io
 
-
-def one_hot_encode(y, num_classes):
-    """
-        One-hot encoding.
-
-        Args:
-            - y : vector to encode.
-            - num_classes: number of classes.
-    """
-    N = y.shape[0]
-
-    encoded = np.zeros((N, num_classes))
-    
-    for i, label in enumerate(y):
-        encoded[i][label] = 1
-    
-    return encoded
-
-
 def load_batch(file_path):
     """
         Loads data given path to file.
@@ -85,3 +66,8 @@ def save_lab_images(img_batch, filename="images/output_{}.png"):
     for i in range(lab_unscaled.shape[0]):
         rgb = color.lab2rgb(lab_unscaled[i])
         io.imsave(filename.format(i), rgb)
+
+def save_gray_images(img_batch, filename="images/output_{}.png"):
+
+    for i in range(img_batch.shape[0]):
+        io.imsave(filename.format(i), img_batch[i])

src/main.py

@@ -7,30 +7,51 @@
 
 SEED = 24
 
-X_train, X_val, _ = load_CIFAR(SEED)
+X_train, Y_train, X_val, Y_val, X_test, Y_test = load_CIFAR(SEED)
 
 X_train = X_train[0:1000,:,:,:]
-X_val = X_val[0:50,:,:,:]
+Y_train = Y_train[0:1000,:,:,:]
 
-print(X_train.shape)
-print(X_val.shape)
+X_val = X_val[0:100,:,:,:]
+Y_val = Y_val[0:100,:,:,:]
+
+X_test = X_test[0:100,:,:,:]
+Y_test = Y_test[0:100,:,:,:]
+
+print('Train:')
+print('X_train:', X_train.shape)
+print('Y_train:', Y_train.shape)
+
+print('Validation:')
+print('X_val:', X_val.shape)
+print('Y_val:', Y_val.shape)
+
+print('Test:')
+print('X_test:', X_test.shape)
+print('Y_test:', Y_test.shape)
+
+save_gray_images(X_train[0:10,:,:,:], filename="images/train_before_gray_{}.png")
+save_lab_images(Y_train[0:10,:,:,:], filename="images/train_before_color_{}.png")
+
+save_gray_images(X_val[0:10,:,:,:], filename="images/val_before_gray_{}.png")
+save_lab_images(Y_val[0:10,:,:,:], filename="images/val_before_color_{}.png")
 
 np.random.seed(SEED)
 tf.random.set_random_seed(SEED)
 
-
 with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
 
     UNET = Model(sess, SEED)
-    
+
     UNET.compile()
 
-    UNET.train(X_train, X_val)
+    print('Training...')
+    UNET.train(X_train, Y_train, X_val, Y_val)
 
-    UNET.predict(X_train)
-    
-    """print(MLP.evaluate(X_train, Y_train))
-    print(MLP.evaluate(X_val, Y_val))
-    print(MLP.evaluate(X_test, Y_test))
-    """
-    
+    print('Predicting training set...')
+    pred = UNET.predict(X_train)
+    save_lab_images(pred[0:10,:,:,:], filename="images/after_train_{}.png")
+
+    print('Predicting validation set...')
+    pred = UNET.predict(X_val)
+    save_lab_images(pred[0:10,:,:,:], filename="images/after_val_{}.png")

src/model.py

@@ -14,7 +14,7 @@ def __init__(self, sess, seed):
         # Training settings.
         self.learning_rate = 0.001
         self.num_epochs = 200
-        self.batch_size = 100
+        self.batch_size = 128
 
         self.compiled = False
 
@@ -35,15 +35,15 @@ def compile(self):
         self.compiled = True
 
         # Placeholders.
-        self.X_rgb = tf.placeholder(tf.float32, shape=(None, 32, 32, 3), name='X_rgb')
-        self.X_gray = tf.placeholder(tf.float32, shape=(None, 32, 32, 1), name='X_gray')
+        self.X = tf.placeholder(tf.float32, shape=(None, 32, 32, 1), name='X')
+        self.Y = tf.placeholder(tf.float32, shape=(None, 32, 32, 3), name='Y')
 
         # Model.
         net = UNet(self.seed)
-        self.out = net.forward(self.X_rgb)
+        self.out = net.forward(self.X)
 
         # Loss and metrics.
-        self.loss = tf.reduce_sum(tf.square(self.out - self.X_rgb))
+        self.loss = tf.reduce_sum(tf.square(self.out - self.Y))
 
         # Optimizer.
         self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)
@@ -53,7 +53,7 @@ def compile(self):
 
         self.saver = tf.train.Saver()
 
-    def train(self, X_train, X_val):
+    def train(self, X_train, Y_train, X_val, Y_val):
 
         if not self.compiled:
             print('Compile model first.')
@@ -68,7 +68,7 @@ def train(self, X_train, X_val):
         merged = tf.summary.merge_all()
         date = str(datetime.datetime.now()).replace(" ", "_")[:19]
         train_writer = tf.summary.FileWriter('logs/' + date + '/train', self.sess.graph)
-        val_writer = tf.summary.FileWriter('logs/' + date + '/val')
+        val_writer = tf.summary.FileWriter('logs/' + date + '/val', self.sess.graph)
         train_writer.flush()
         val_writer.flush()
 
@@ -81,16 +81,9 @@ def train(self, X_train, X_val):
                     start = b * self.batch_size
                     end   = min(b * self.batch_size + self.batch_size, N)
                     batch_x = X_train[start:end,:,:,:]
+                    batch_y = Y_train[start:end,:,:,:]
 
-                    #if b != 0:
-                    #    end_t_2 = timer()
-                    #    print('data load: {0}'.format(end_t_2 - start_t_2))
-
-                    #start_t = timer()
-                    _, l = self.sess.run([self.optimizer, self.loss], feed_dict={self.X_rgb: batch_x})
-                    #end_t = timer()
-                    #print('sess.run: {0}'.format(end_t - start_t))
-                    #start_t_2 = timer()
+                    _, l = self.sess.run([self.optimizer, self.loss], feed_dict={self.X: batch_x ,self.Y: batch_y})
 
                     epoch_loss += l / num_batches
 
@@ -105,9 +98,9 @@ def train(self, X_train, X_val):
                     train_writer.flush()
 
                     # Add validation loss to val log.
-                    #summary = self.sess.run(merged, feed_dict={self.X_rgb: X_val})
-                    #val_writer.add_summary(summary, epoch)
-                    #val_writer.flush()
+                    summary = self.sess.run(merged,  feed_dict={self.X: X_val ,self.Y: Y_val})
+                    val_writer.add_summary(summary, epoch)
+                    val_writer.flush()
 
                 # Save model.
                 if self.save and epoch % self.save_interval == 0:
@@ -136,4 +129,4 @@ def predict(self, X):
             print('Compile model first.')
             return
 
-        return self.sess.run(self.out, feed_dict={self.X_rgb: X})
+        return self.sess.run(self.out, feed_dict={self.X: X})