Minor corrections

.gitignore

@@ -0,0 +1,4 @@
+datasets_medmnist/
+train.sh
+train2.py
+results_medmnist/

generate_images.py

@@ -11,20 +11,27 @@ if __name__ == "__main__":
     parser.add_argument("--latent_dim", default=8, type=int, help='Dimensionality of latent space')
     parser.add_argument("--save_dir", default="./fake_images", type=str, help='Path to save synthetic images')
     parser.add_argument("--num_of_images", default=10, type=int, help='Number of images to generate')
+    parser.add_argument("--input_shape", default=[320, 320, 3], nargs=3, help='Image shape for training')
+
     args = parser.parse_args()
+    args.input_shape = tuple(map(int, args.input_shape))
 
     os.makedirs(args.save_dir, exist_ok=True)
 
     if not os.path.exists(args.weights_path):
         print("Not a valid path")
 
-    vae = init_vae_model(args.model_name, args.latent_dim)
-    noise_vector = get_noise_seeded((args.num_of_images, args.latent_dim))
+    vae = init_vae_model(args.model_name, args.latent_dim, args.input_shape)
+    # noise_vector = get_noise_seeded((args.num_of_images, args.latent_dim))
+
 
     # Load weights
     vae = load_weights(vae, args.weights_path)
     vae.trainable = False
 
     # Generate & Save images
-    fake_images = decode_noise(vae, noise_vector, return_list=True)
-    save_images(args.save_dir, fake_images)
+    for i in range(args.num_of_images):
+        print(f'Generating image for seed {i}/{args.num_of_images}, ')
+        noise_vector = get_noise_seeded((1, args.latent_dim), seed=i)
+        fake_images = decode_noise(vae, noise_vector, return_list=True)
+        save_images(args.save_dir, fake_images, seed=i)

utils/inference_utils.py

@@ -8,11 +8,12 @@ from model.vae import VAE
 from model import tidev2
 
 
-def init_vae_model(model_name, latent_dim):
+def init_vae_model(model_name, latent_dim, input_shape):
     if model_name == 'tidev2':
         vae_model = VAE(tidev2.ConvNeXtEncoderTiny(latent_dim=latent_dim),
-                        tidev2.ConvNeXtDecoderTiny(latent_dim=latent_dim)
+                        tidev2.ConvNeXtDecoderTiny(latent_dim=latent_dim, image_dims=input_shape[:2], out_channels=input_shape[-1])
                        )
+        vae_model.build((None, *input_shape))
         return vae_model
 
 
@@ -27,12 +28,11 @@ def load_weights(vae, weights_path):
         return vae
 
 
-def get_noise_seeded(noise_shape):
-    np.random.seed(0)
+def get_noise_seeded(noise_shape, seed=0):
+    np.random.seed(seed)
     random_z = np.random.normal(0, 1, noise_shape)
     return random_z
 
-
 def decode_noise(trained_vae, noise, return_list=False):
     print("Generating synthetic images ...")
     pred = trained_vae.decoder.predict(noise, batch_size=1)
@@ -44,9 +44,12 @@ def decode_noise(trained_vae, noise, return_list=False):
     return pred
 
 
-def save_images(save_folder, images):
+def save_images(save_folder, images, seed=None):
     print(f"Saving  synthetic images into {save_folder}")
     if isinstance(images, list):
         for i, image in enumerate(images):
             image = image.astype(np.uint8)
-            Image.fromarray(image).save(os.path.join(save_folder, f"image-{i}.jpg"))
+            if image.shape[-1] == 1:
+                image = np.squeeze(image, axis=-1)
+            save_filename = f"image-{i}.jpg" if seed is None else f"image-{seed}.jpg"
+            Image.fromarray(image).save(os.path.join(save_folder, save_filename))