forked from NVlabs/SPADE
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
executable file
·49 lines (38 loc) · 1.48 KB
/
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
"""
Copyright (C) 2019 NVIDIA Corporation. All rights reserved.
Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
"""
import os
import torch
from collections import OrderedDict
import data
from options.test_options import TestOptions
from models.pix2pix_model import Pix2PixModel
from util.visualizer import Visualizer
from util import html
opt = TestOptions().parse()
dataloader = data.create_dataloader(opt)
model = Pix2PixModel(opt)
model.eval()
visualizer = Visualizer(opt)
# create a webpage that summarizes the all results
web_dir = os.path.join(opt.results_dir, opt.name,
'%s_%s' % (opt.phase, opt.which_epoch))
webpage = html.HTML(web_dir,
'Experiment = %s, Phase = %s, Epoch = %s' %
(opt.name, opt.phase, opt.which_epoch))
# test
for i, data_i in enumerate(dataloader):
if i * opt.batchSize >= opt.how_many:
break
with torch.no_grad():
input_semantics, real_image = model.preprocess_input(data_i)
generated = model.netG(input_semantics, z=None)
#generated = model(data_i, mode='inference')
img_path = data_i['path']
for b in range(generated.shape[0]):
print('process image... %s' % img_path[b])
visuals = OrderedDict([('input_label', data_i['label'][b]),
('synthesized_image', generated[b])])
visualizer.save_images(webpage, visuals, img_path[b:b + 1])
webpage.save()