105 lines
3.1 KiB
Python
105 lines
3.1 KiB
Python
import tensorflow as tf
|
|
|
|
import os
|
|
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
|
|
|
|
|
|
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, 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]
|
|
filter_index = resolutions.index(res)
|
|
g_params = {
|
|
'z_dim': 256,
|
|
'w_dim': 256,
|
|
'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)
|
|
dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(1).repeat()
|
|
|
|
i = 0
|
|
for batch in dataset:
|
|
print(i)
|
|
# generate image
|
|
fake_images = generator([tf.random.normal(shape=[1, g_params['z_dim']]),
|
|
batch[1]],
|
|
training=False, truncation_psi=0.5)
|
|
|
|
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) + '.jpg'), img)
|
|
i += 1
|
|
|
|
if i > samples:
|
|
break
|
|
|
|
return
|
|
|
|
|
|
def main():
|
|
allow_memory_growth()
|
|
|
|
checkpoints = [
|
|
{
|
|
'res': 512,
|
|
'ckpt_dir': 'models/stylegan2-test-5-512x512',
|
|
'use_custom_cuda': True,
|
|
'folder': 'output',
|
|
'data_folder': '/home/workspace/data/dataset/stylegan_new/train'
|
|
},
|
|
]
|
|
|
|
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'
|
|
data_folder = run_item['data_folder']
|
|
|
|
print(message)
|
|
inference(ckpt_dir, use_custom_cuda, res, folder, data_folder)
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|