plotting not working, losses fixed?

inference_custom.py

@@ -5,6 +5,7 @@ import sys
 import cv2
 import numpy as np
 from PIL import Image
+from dataset import DatasetSimple
 from load_models import load_generator
 from tf_utils import allow_memory_growth
 
@@ -38,7 +39,7 @@ def uv_to_mesh(uv):
     return mesh
 
 
-def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000):
+def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100):
     # create generator
     resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
     featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
@@ -52,17 +53,23 @@ def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000):
         'featuremaps': featuremaps[:filter_index + 1],
     }
     generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda)
+    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(1).repeat()
 
-    for i in range(samples):
+    i = 0
+    for batch in dataset:
         print(i)
         # generate image
         fake_images = generator([tf.random.normal(shape=[1, g_params['z_dim']]),
-                                tf.random.normal(shape=[1, g_params['labels_dim']])],
+                                batch[1]],
                                 training=False, truncation_psi=0.5)
 
-        np.savez_compressed(os.path.join(folder, str(i+100) + '.npz'), uv=fake_images)
+        np.savez_compressed(os.path.join(folder, str(i) + '.npz'), uv=tf.squeeze(fake_images), template=tf.squeeze(batch[1][2]), body=tf.squeeze(batch[3]), mask=tf.squeeze(batch[4]))
         img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0])
-        cv2.imwrite(os.path.join(folder, str(i+100) + '.jpg'), img)
+        cv2.imwrite(os.path.join(folder, str(i) + '.jpg'), img)
+        i += 1
+
+        if i > samples:
+            break
 
     return
 
@@ -73,9 +80,10 @@ def main():
     checkpoints = [
         {
             'res': 512,
-            'ckpt_dir': '/workspace/Checkpoints/checkpoint_52000/stylegan2-test-2-512x512',
-            'use_custom_cuda': False,
-            'folder': '/workspace/Checkpoints/checkpoint_52000/output',
+            'ckpt_dir': 'models/stylegan2-test-5-512x512',
+            'use_custom_cuda': True,
+            'folder': 'output',
+            'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
         },
     ]
 
@@ -85,9 +93,10 @@ def main():
         use_custom_cuda = run_item['use_custom_cuda']
         folder = run_item['folder']
         message = f'{res}x{res} with custom cuda' if use_custom_cuda else f'{res}x{res} without custom cuda'
+        data_folder = run_item['data_folder']
 
         print(message)
-        inference(ckpt_dir, use_custom_cuda, res, folder)
+        inference(ckpt_dir, use_custom_cuda, res, folder, data_folder)
     return