上述函数改为读入一张图片
时间: 2023-12-15 21:06:26 浏览: 65
图片信息读取
如果要将上述函数改为读入一张图片进行测试,可以将`test_files`改为一个包含单张图片路径的列表即可。修改后的代码如下:
```python
def test(checkpoint_dir, style_name, test_img, if_adjust_brightness, img_size=[256,256]):
# tf.reset_default_graph()
result_dir = 'results/'+style_name
check_folder(result_dir)
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)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)) as sess:
# tf.global_variables_initializer().run()
# 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: ' + test_img)
sample_image = np.asarray(load_test_data(test_img, img_size))
image_path = os.path.join(result_dir,'{0}'.format(os.path.basename(test_img)))
fake_img = sess.run(test_generated, feed_dict = {test_real : sample_image})
if if_adjust_brightness:
save_images(fake_img, image_path, test_img)
else:
save_images(fake_img, image_path, None)
end = time.time()
print(f'test-time: {end-begin} s')
```
在调用该函数时,传入的第三个参数`test_img`应为单张图片的路径。
阅读全文