96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
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()
|