-
Notifications
You must be signed in to change notification settings - Fork 70
/
run_test.py
98 lines (79 loc) · 3.23 KB
/
run_test.py
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
import tensorflow as tf
import os
import utils
import style_transfer_tester
import argparse
import time
"""parsing and configuration"""
def parse_args():
desc = "Tensorflow implementation of 'Perceptual Losses for Real-Time Style Transfer and Super-Resolution'"
parser = argparse.ArgumentParser(description=desc)
parser.add_argument('--style_model', type=str, default='models/wave.ckpt', help='location for model file (*.ckpt)',
required=True)
parser.add_argument('--content', type=str, default='content/female_knight.jpg',
help='File path of content image (notation in the paper : x)', required=True)
parser.add_argument('--output', type=str, default='result.jpg',
help='File path of output image (notation in the paper : y_c)', required=True)
parser.add_argument('--max_size', type=int, default=None, help='The maximum width or height of input images')
return check_args(parser.parse_args())
"""checking arguments"""
def check_args(args):
# --style_model
try:
#Tensorflow r0.12 requires 3 files related to *.ckpt
assert os.path.exists(args.style_model + '.index') and os.path.exists(args.style_model + '.meta') and os.path.exists(
args.style_model + '.data-00000-of-00001')
except:
print('There is no %s'%args.style_model)
print('Tensorflow r0.12 requires 3 files related to *.ckpt')
print('If you want to restore any models generated from old tensorflow versions, this assert might be ignored')
return None
# --content
try:
assert os.path.exists(args.content)
except:
print('There is no %s' % args.content)
return None
# --max_size
try:
if args.max_size is not None:
assert args.max_size > 0
except:
print('The maximum width or height of input image must be positive')
return None
# --output
dirname = os.path.dirname(args.output)
try:
if len(dirname) > 0:
os.stat(dirname)
except:
os.mkdir(dirname)
return args
"""main"""
def main():
# parse arguments
args = parse_args()
if args is None:
exit()
# load content image
content_image = utils.load_image(args.content, max_size=args.max_size)
# open session
soft_config = tf.ConfigProto(allow_soft_placement=True)
soft_config.gpu_options.allow_growth = True # to deal with large image
sess = tf.Session(config=soft_config)
# build the graph
transformer = style_transfer_tester.StyleTransferTester(session=sess,
model_path=args.style_model,
content_image=content_image,
)
# execute the graph
start_time = time.time()
output_image = transformer.test()
end_time = time.time()
# save result
utils.save_image(output_image, args.output)
# report execution time
shape = content_image.shape #(batch, width, height, channel)
print('Execution time for a %d x %d image : %f msec' % (shape[0], shape[1], 1000.*float(end_time - start_time)/60))
if __name__ == '__main__':
main()