wahaha commited on
Commit
7d9456d
1 Parent(s): 149e6d6
Files changed (1) hide show
  1. test1.py +1 -43
test1.py CHANGED
@@ -8,45 +8,6 @@ import numpy as np
8
  from net import generator
9
  os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
10
 
11
- class ImportGraph:
12
- def __init__(self, checkpoint_dir):
13
- self.graph = tf.Graph()
14
-
15
- with self.graph.as_default():
16
-
17
- test_real = tf.placeholder(tf.float32, [1, None, None, 3], name='test')
18
- with tf.variable_scope("generator", reuse=False):
19
- self.test_generated = generator.G_net(test_real).fake
20
-
21
- self.saver = tf.train.Saver()
22
-
23
- self.sess = tf.Session(graph=self.graph, config=tf.ConfigProto(allow_soft_placement=True,
24
- gpu_options=tf.GPUOptions(
25
- allow_growth=True)))
26
-
27
- ckpt = tf.train.get_checkpoint_state(checkpoint_dir) # checkpoint file information
28
- if ckpt and ckpt.model_checkpoint_path:
29
- ckpt_name = os.path.basename(ckpt.model_checkpoint_path) # first line
30
- self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
31
- print(" [*] Success to read {}".format(os.path.join(checkpoint_dir, ckpt_name)))
32
- else:
33
- print(" [*] Failed to find a checkpoint")
34
-
35
- def test(self, style_name, sample_file, if_adjust_brightness, img_size=[256,256]):
36
- result_dir = 'results/' + style_name
37
- check_folder(result_dir)
38
-
39
- sample_image = np.asarray(load_test_data(sample_file, img_size))
40
- image_path = os.path.join(result_dir, '{0}'.format(os.path.basename(sample_file)))
41
-
42
- fake_img = self.sess.run(self.test_generated, feed_dict={test_real: sample_image})
43
- if if_adjust_brightness:
44
- save_images(fake_img, image_path, sample_file)
45
- else:
46
- save_images(fake_img, image_path, None)
47
-
48
- return image_path
49
-
50
 
51
  def stats_graph(graph):
52
  flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())
@@ -102,7 +63,4 @@ def test(checkpoint_dir, style_name, test_file, if_adjust_brightness, img_size=[
102
 
103
  return image_path
104
 
105
- if __name__ == '__main__':
106
- arg = parse_args()
107
- print(arg.checkpoint_dir)
108
- test(arg.checkpoint_dir, arg.save_dir, arg.test_dir, arg.if_adjust_brightness)
 
8
  from net import generator
9
  os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
10
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
 
12
  def stats_graph(graph):
13
  flops = tf.profiler.profile(graph, options=tf.profiler.ProfileOptionBuilder.float_operation())
 
63
 
64
  return image_path
65
 
66
+