66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
import tensorflow as tf
|
|
|
|
from PIL import Image
|
|
from load_models import load_generator
|
|
from tf_utils import allow_memory_growth
|
|
|
|
|
|
def inference(ckpt_dir, use_custom_cuda, res, out_fn=None):
|
|
# 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)
|
|
|
|
# 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)
|
|
fake_images = (tf.clip_by_value(fake_images, -1.0, 1.0) + 1.0) * 127.5
|
|
fake_images = tf.transpose(fake_images, perm=[0, 2, 3, 1])
|
|
fake_images = tf.cast(fake_images, tf.uint8)
|
|
fake_image = fake_images[0].numpy()
|
|
|
|
image = Image.fromarray(fake_image)
|
|
image = image.convert('RGB')
|
|
image.show()
|
|
if out_fn is not None:
|
|
image.save(out_fn)
|
|
return
|
|
|
|
|
|
def main():
|
|
allow_memory_growth()
|
|
|
|
checkpoints = [
|
|
{
|
|
'res': 512,
|
|
'ckpt_dir': '/workspace/Checkpoints/checkpoint_32100/stylegan2-test-2-512x512',
|
|
'use_custom_cuda': False,
|
|
'out_fn': 'out_0.png',
|
|
},
|
|
]
|
|
|
|
for run_item in checkpoints:
|
|
res = run_item['res']
|
|
ckpt_dir = run_item['ckpt_dir']
|
|
use_custom_cuda = run_item['use_custom_cuda']
|
|
out_fn = run_item['out_fn']
|
|
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, out_fn)
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|