def main(_):
if args.use_gpu:
print("[*] GPU\n")
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_idx
# gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_mem)
# with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:
with tf.Session(config=config) as sess:
model = lowlight_enhance(sess)
if args.phase == 'train':
lowlight_train(model)
elif args.phase == 'test':
lowlight_test(model)
else:
print('[!] Unknown phase')
exit(0)
else:
print("[*] CPU\n")
with tf.Session() as sess:
model = lowlight_enhance(sess)
if args.phase == 'train':
lowlight_train(model)
elif args.phase == 'test':
lowlight_test(model)
else:
print('[!] Unknown phase')
exit(0)
def main(input_file, use_gpu=False):
print('called main')
if use_gpu:
print("[*] GPU\n")
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_idx
gpu_options = tf.GPUOptions(
per_process_gpu_memory_fraction=args.gpu_mem)
with tf.Session(config=tf.ConfigProto(
gpu_options=gpu_options)) as sess:
model = lowlight_enhance(sess)
lowlight_test(input_file, model)
else:
print("[*] CPU\n")
with tf.Session() as sess:
model = lowlight_enhance(sess)
lowlight_test(input_file, model)
#if __name__ == '__main__':
# main(input_file, use_gpu=False)
#tf.app.run()