Minor corrections

model/tidev2.py

@@ -87,17 +87,28 @@ class ConvNeXtDecoderTiny(Model):
                  drop_path_rate=0.0,
                  layer_scale_init_value=1e-6,
                  model_name="convnext",
-                 latent_dim=None):
+                 latent_dim=None,
+                 image_dims=(320, 320),
+                 out_channels=3):
         super().__init__(name=model_name)
 
         if latent_dim is None:
             raise ValueError("latent_dim must be specified for decoder")
 
         # Intro layer (dense + reshape)
+        # self.intro = Sequential([
+        #     layers.Dense(10 * 10 * projection_dims[0], activation="relu"),
+        #     layers.Reshape((10, 10, projection_dims[0]))
+        # ], name=model_name + "_intro")
+        # TODO
+        downsample_factor = 4 * 2 * 2 * 2
+        input_height, input_width = image_dims
+        init_h = input_height // downsample_factor
+        init_w = input_width // downsample_factor
         self.intro = Sequential([
-            layers.Dense(10 * 10 * projection_dims[0], activation="relu"),
-            layers.Reshape((10, 10, projection_dims[0]))
-        ], name=model_name + "_intro")
+            layers.Dense(init_h * init_w * projection_dims[0], activation="relu"),
+            layers.Reshape((init_h, init_w, projection_dims[0]))
+        ])
 
         # Upsampling layers
         self.upsample_layers = [self.intro]
@@ -133,7 +144,7 @@ class ConvNeXtDecoderTiny(Model):
         ], name=model_name + "_top")
 
         self.top_layer = TopLayer(filters=96)
-        self.pred_layer = layers.Conv2DTranspose(3, kernel_size=1, activation="sigmoid",
+        self.pred_layer = layers.Conv2DTranspose(out_channels, kernel_size=1, activation="sigmoid",
                                                  padding="same", name="pred_layer")
 
     def call(self, inputs, training=False):

model/tidev2_utils.py

@@ -8,24 +8,24 @@ class TopLayer(layers.Layer):
         self.filters = filters
 
         self.conv_1x1 = layers.Conv2D(self.filters, (1, 1), activation='relu', strides=1, padding="same",
-                                      name="_top_layer")
+                                      name="top_layer_1x1")
         self.conv_2x2 = layers.Conv2D(self.filters//3, (2, 2), activation='relu', strides=1, padding="same",
-                                      name="_top_layer")
+                                      name="top_layer_2x2")
         self.conv_4x4 = layers.Conv2D(self.filters//3, (4, 4), activation='relu', strides=1, padding="same",
-                                      name="_top_layer")
+                                      name="top_layer_4x4")
         self.conv_8x8 = layers.Conv2D(self.filters//3, (8, 8), activation='relu', strides=1, padding="same",
-                                      name="_top_layer")
+                                      name="top_layer_8x8")
 
         self.concat = layers.Concatenate(axis=-1)
         self.point_wise_conv = layers.Conv2D(self.filters, (1, 1), 1, activation=None, use_bias=False,
-                                             padding='same', name="_top_layer")
+                                             padding='same', name="top_layer_point_wise")
         self.feat_fusion = layers.Conv2D(self.filters, (1, 1), 1, activation=None, use_bias=False,
-                                         padding='same', name="_top_layer")
+                                         padding='same', name="top_layer_fusion")
 
         self.addition = layers.Add()
         self.gelu = layers.Activation('gelu')
         self.final_conv = layers.Conv2D(self.filters, (1, 1),  activation='relu', strides=1, padding="same",
-                                        name="_top_layer")
+                                        name="top_layer_out")
 
     def call(self, inputs, training=False):
         x = self.conv_1x1(inputs, training=training)

train.py

@@ -4,7 +4,6 @@ import tensorflow as tf
 from json import dump
 from argparse import ArgumentParser
 
-
 from model import tidev2
 from model.vae import VAE
 from utils.callbacks import VisualizeCallback, CheckpointCallback
@@ -14,10 +13,11 @@ from utils.plots import visualize_from_latent_space
 
 if __name__ == '__main__':
     parser = ArgumentParser()
+
     parser.add_argument("--model_name", required=True, type=str, choices=['tide', 'tidev2'], help='VAE model')
     parser.add_argument("--output_path", default='./results/', type=str, help='Path to store the results')
     # VAE model
-    parser.add_argument("--input_shape", default=(320, 320, 3), type=tuple, help='Image shape for training')
+    parser.add_argument("--input_shape", default=[320, 320, 3], nargs=3, help='Image shape for training')
     parser.add_argument("--dim_latent", default=8, type=int, help='Dimensionality of latent space')
     # Training
     parser.add_argument("--epochs", default=5000, type=int, help='Number of training epochs')
@@ -33,6 +33,7 @@ if __name__ == '__main__':
     parser.add_argument("--crop_dim", default=None, type=tuple,
                         help='Dimensions for cropping images. Ignore if images are already cropped')
     args = parser.parse_args()
+    args.input_shape = tuple(map(int, args.input_shape))
 
     # Create folders & Save training config
     os.makedirs(args.output_path, exist_ok=True)
@@ -59,8 +60,11 @@ if __name__ == '__main__':
     # Create Model
     if args.model_name == 'tidev2':
         vae = VAE(tidev2.ConvNeXtEncoderTiny(latent_dim=args.dim_latent),
-                  tidev2.ConvNeXtDecoderTiny(latent_dim=args.dim_latent)
+                  tidev2.ConvNeXtDecoderTiny(latent_dim=args.dim_latent,
+                                             image_dims=args.input_shape[:2],
+                                             out_channels=args.input_shape[-1])
                   )
+        vae.build((None, *args.input_shape))
         vae.compile(optimizer=tf.keras.optimizers.Adam(args.learning_rate))
 
     # Training
@@ -85,3 +89,4 @@ if __name__ == '__main__':
             shuffle=True,
             initial_epoch=0)
 
+    print('Training finished')

utils/dataloader.py

@@ -1,6 +1,6 @@
 import os
-import random
 import numpy as np
+
 from PIL import Image
 from re import split, compile
 from tensorflow.keras.utils import Sequence
@@ -19,7 +19,7 @@ def list_filenames(data_path, img_extension='png', filename_prefix=None):
 
 
 class Dataset(Sequence):
-    def __init__(self, file_list, batch_size=32, crop_dim=None, resize_dim=None, shuffle=True):
+    def __init__(self, file_list, batch_size=32, crop_dim=None, resize_dim=None, shuffle=True, mode='RGB'):
         self.files_list = file_list
         self.batch_size = batch_size
 
@@ -28,6 +28,8 @@ class Dataset(Sequence):
         self.shuffle = shuffle
         self.on_epoch_end()
 
+        self.mode=mode
+
     def __len__(self):
         return int(np.ceil(len(self.files_list) / self.batch_size))
 
@@ -53,7 +55,10 @@ class Dataset(Sequence):
         return image.crop((left, top, right, bottom))
 
     def load_images(self, filepath):
-        image = Image.open(filepath).convert('RGB')
+        if self.mode=='RGB':
+            image = Image.open(filepath).convert('RGB')
+        else:
+            image = Image.open(filepath)
         if self.crop_dim:
             image = self.center_crop(image, crop_dim=self.crop_dim)
         if self.resize_dim:
@@ -61,4 +66,6 @@ class Dataset(Sequence):
 
         image = np.array(image).astype(np.float32)
         image = image / 255.0
-        return image
+        if image.ndim == 2:
+            image = np.expand_dims(image, -1)
+        return image

utils/plots.py

@@ -1,11 +1,11 @@
-import imageio
 import numpy as np
+from PIL import Image
 
 
 def visualize_from_latent_space(latent_dim, input_shape, vae, output_path, epoch="final", num_items=10,):
 
     image_size, _, img_channels = input_shape
-    figure = np.zeros((image_size * num_items, image_size * num_items, 3))
+    figure = np.zeros((image_size * num_items, image_size * num_items, img_channels))
 
     scale = 1.0
     grid_x = np.linspace(-scale, scale, num_items)
@@ -18,8 +18,10 @@ def visualize_from_latent_space(latent_dim, input_shape, vae, output_path, epoch
             x_decoded = vae.decoder.predict(random_z)
             image = x_decoded[0].reshape(input_shape)
             figure[i * image_size: (i + 1) * image_size, j * image_size: (j + 1) * image_size, ] = image
-    print(f'Saving collage in {output_path}/decoding-noise-ep{epoch}.jpg')
-    imageio.imsave(f'{output_path}/decoding-noise-ep{epoch}.jpg', (figure * 255).astype('uint8'))
-
-
+    print(f'Saving collage in {output_path}/decoding-noise-ep{epoch}.png')
+    figure = (figure * 255).astype('uint8')
+    if img_channels == 1:
+        figure = np.squeeze(figure, axis=-1)
+    figure = Image.fromarray(figure)
+    figure.save(f"{output_path}/decoding-noise-ep{epoch}.jpg")