Spaces:
Runtime error
Runtime error
File size: 4,221 Bytes
12ae7b0 f937c5e c9dfbd7 f937c5e 14baf72 12ae7b0 c9dfbd7 f937c5e 14baf72 f937c5e 12ae7b0 34a888d f937c5e f3cf058 14baf72 f937c5e 12ae7b0 d67e05e 12ae7b0 d67e05e 12ae7b0 d67e05e 12ae7b0 d67e05e 12ae7b0 d67e05e 12ae7b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
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"
class ImportGraph:
def __init__(self, checkpoint_dir):
self.graph = tf.Graph()
with self.graph.as_default():
test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
with tf.variable_scope("generator", reuse=False):
self.test_generated = generator.G_net(test_real).fake
self.saver = tf.train.Saver()
self.sess = tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
gpu_options=tf.GPUOptions(
allow_growth=True)))
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
self.saver.restore(self.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")
def test(self, style_name, sample_file, if_adjust_brightness, img_size=[256,256]):
result_dir = 'results/' + style_name
check_folder(result_dir)
sample_image = np.asarray(load_test_data(sample_file, img_size))
image_path = os.path.join(result_dir, '{0}'.format(os.path.basename(sample_file)))
fake_img = self.sess.run(self.test_generated, feed_dict={test_real: sample_image})
if if_adjust_brightness:
save_images(fake_img, image_path, sample_file)
else:
save_images(fake_img, image_path, None)
return image_path
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
def test(checkpoint_dir, style_name, test_file, if_adjust_brightness, img_size=[256,256]):
global g_sess
global test_generated
# 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
if __name__ == '__main__':
arg = parse_args()
print(arg.checkpoint_dir)
test(arg.checkpoint_dir, arg.save_dir, arg.test_dir, arg.if_adjust_brightness)
|