Spaces:
Runtime error
Runtime error
| import argparse | |
| from tools.utils import * | |
| import os | |
| from tqdm import tqdm | |
| from glob import glob | |
| import time | |
| import numpy as np | |
| from net import generator | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "-1" | |
| def stats_graph(graph): | |
| flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation()) | |
| # params = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.trainable_variables_parameter()) | |
| print('FLOPs: {}'.format(flops.total_float_ops)) | |
| g_sess = None | |
| test_generated = None | |
| test_real = None | |
| def test(checkpoint_dir, style_name, test_file, if_adjust_brightness, img_size=[256,256]): | |
| global g_sess | |
| global test_generated | |
| global test_real | |
| # tf.reset_default_graph() | |
| result_dir = 'results/'+style_name | |
| check_folder(result_dir) | |
| if g_sess is None: | |
| test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test') | |
| with tf.variable_scope("generator", reuse=False): | |
| test_generated = generator.G_net(test_real).fake | |
| saver = tf.train.Saver() | |
| gpu_options = tf.GPUOptions(allow_growth=True) | |
| g_sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)) | |
| # load model | |
| ckpt = tf.train.get_checkpoint_state(checkpoint_dir) # checkpoint file information | |
| if ckpt and ckpt.model_checkpoint_path: | |
| ckpt_name = os.path.basename(ckpt.model_checkpoint_path) # first line | |
| saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name)) | |
| print(" [*] Success to read {}".format(os.path.join(checkpoint_dir, ckpt_name))) | |
| else: | |
| print(" [*] Failed to find a checkpoint") | |
| return | |
| # stats_graph(tf.get_default_graph()) | |
| begin = time.time() | |
| # print('Processing image: ' + sample_file) | |
| sample_image = np.asarray(load_test_data(test_file, img_size)) | |
| image_path = os.path.join(result_dir,'{0}'.format(os.path.basename(test_file))) | |
| fake_img = g_sess.run(test_generated, feed_dict = {test_real : sample_image}) | |
| if if_adjust_brightness: | |
| save_images(fake_img, image_path, test_file) | |
| else: | |
| save_images(fake_img, image_path, None) | |
| end = time.time() | |
| print(f'test-time: {end-begin} s') | |
| return image_path | |