init

inference_custom.py

@@ -0,0 +1,95 @@
+import tensorflow as tf
+
+import os
+import sys
+import cv2
+import numpy as np
+from PIL import Image
+from load_models import load_generator
+from tf_utils import allow_memory_growth
+
+
+def m2uv_array(m2uv):
+    result = np.zeros((14475, 2), dtype=np.int32)
+    for k, v in m2uv.items():
+        if len(v) == 2:
+            result[k] = list(v)
+        elif len(v) == 1:
+            x = v.pop()
+            result[k] = [x, x]
+        else:
+            assert False
+
+    return result
+
+
+def uv_to_mesh(uv):
+    uv_res = 256
+    m2uv = np.load('MeshUtils/m2uv.npy', allow_pickle=True).item()
+    m2uv = m2uv_array(m2uv)
+    uv_map = np.load('Meshutils/stretched_uv_map_3.npy')
+    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
+
+    mesh = np.zeros((14475, 3))
+    uv_pixels = uv_map_discreet[m2uv]
+    # counter = np.ones((14475))
+    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
+    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
+    return mesh
+
+
+def inference(ckpt_dir, use_custom_cuda, res, folder=None, samples=1000):
+    # create generator
+    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
+    featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
+    filter_index = resolutions.index(res)
+    g_params = {
+        'z_dim': 512,
+        'w_dim': 512,
+        'labels_dim': 0,
+        'n_mapping': 8,
+        'resolutions': resolutions[:filter_index + 1],
+        'featuremaps': featuremaps[:filter_index + 1],
+    }
+    generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda)
+
+    for i in range(samples):
+        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']])],
+                                training=False, truncation_psi=0.5)
+
+        np.savez_compressed(os.path.join(folder, str(i+100) + '.npz'), uv=fake_images)
+        img = np.transpose(np.round((fake_images + 1) * 128)[0], [1,2,0])
+        cv2.imwrite(os.path.join(folder, str(i+100) + '.jpg'), img)
+
+    return
+
+
+def main():
+    allow_memory_growth()
+
+    checkpoints = [
+        {
+            'res': 512,
+            'ckpt_dir': '/workspace/Checkpoints/checkpoint_52000/stylegan2-test-2-512x512',
+            'use_custom_cuda': False,
+            'folder': '/workspace/Checkpoints/checkpoint_52000/output',
+        },
+    ]
+
+    for run_item in checkpoints:
+        res = run_item['res']
+        ckpt_dir = run_item['ckpt_dir']
+        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'
+
+        print(message)
+        inference(ckpt_dir, use_custom_cuda, res, folder)
+    return
+
+
+if __name__ == '__main__':
+    main()