File size: 1,569 Bytes
aba0e05 |
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.input_A = self.Tensor(opt.batchSize, opt.input_nc, opt.fineSize, opt.fineSize)
self.netG = networks.define_G(opt.input_nc, opt.output_nc,
opt.ngf, opt.which_model_netG,
opt.norm, not opt.no_dropout,
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']
self.input_A.resize_(input_A.size()).copy_(input_A)
self.image_paths = input['A_paths']
def test(self):
self.real_A = Variable(self.input_A)
self.fake_B = self.netG.forward(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)])
|