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
|
from torch.autograd import Variable
from collections import OrderedDict
import util.util as util
from .base_model import BaseModel
from . import networks
class TestModel(BaseModel):
def name(self):
return 'TestModel'
def initialize(self, opt):
assert(not opt.isTrain)
BaseModel.initialize(self, opt)
self.netG = networks.define_G(opt.input_nc, opt.output_nc,
opt.ngf, opt.which_model_netG,
opt.norm, not opt.no_dropout,
opt.init_type,
self.gpu_ids)
which_epoch = opt.which_epoch
self.load_network(self.netG, 'G', which_epoch)
print('---------- Networks initialized -------------')
networks.print_network(self.netG)
print('-----------------------------------------------')
def set_input(self, input):
# we need to use single_dataset mode
input_A = input['A']
if len(self.gpu_ids) > 0:
input_A.cuda(self.gpu_ids[0], async=True)
self.input_A = input_A
self.image_paths = input['A_paths']
def test(self):
self.real_A = Variable(self.input_A)
self.fake_B = self.netG(self.real_A)
# get image paths
def get_image_paths(self):
return self.image_paths
def get_current_visuals(self):
real_A = util.tensor2im(self.real_A.data)
fake_B = util.tensor2im(self.fake_B.data)
return OrderedDict([('real_A', real_A), ('fake_B', fake_B)])
|
