updated pipeline, added uv_mask_loss

inference_custom.py

@@ -9,6 +9,7 @@ from dataset import DatasetSimple
 from load_models import load_generator
 from tf_utils import allow_memory_growth
 
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
 
 def m2uv_array(m2uv):
     result = np.zeros((14475, 2), dtype=np.int32)
@@ -39,7 +40,7 @@ def uv_to_mesh(uv):
     return mesh
 
 
-def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100):
+def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50):
     # create generator
     resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
     featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
@@ -53,7 +54,9 @@ def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100):
         '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()
+    dataset = DatasetSimple(data_folder).cache().batch(1).repeat()
+
+    os.makedirs(os.path.join(folder), exist_ok=True)
 
     i = 0
     for batch in dataset:
@@ -80,9 +83,9 @@ def main():
     checkpoints = [
         {
             'res': 512,
-            'ckpt_dir': 'models/stylegan2-test-5-512x512',
+            'ckpt_dir': 'models/stylegan2-test-8-512x512',
             'use_custom_cuda': True,
-            'folder': 'output',
+            'folder': 'output/test-8-196000',
             'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
         },
     ]