diff options
Diffstat (limited to 'models')
| -rw-r--r-- | models/models.py | 2 | ||||
| -rw-r--r-- | models/pix2pix_model.py | 19 | ||||
| -rw-r--r-- | models/test_model.py | 31 |
3 files changed, 24 insertions, 28 deletions
diff --git a/models/models.py b/models/models.py index efcd898..d5bb9d8 100644 --- a/models/models.py +++ b/models/models.py @@ -13,7 +13,7 @@ def create_model(opt): elif opt.model == 'test': assert(opt.dataset_mode == 'single') from .test_model import TestModel - model = OneDirectionTestModel() + model = TestModel() else: raise ValueError("Model [%s] not recognized." % opt.model) model.initialize(opt) diff --git a/models/pix2pix_model.py b/models/pix2pix_model.py index 4581d33..e44529b 100644 --- a/models/pix2pix_model.py +++ b/models/pix2pix_model.py @@ -8,6 +8,7 @@ from util.image_pool import ImagePool from .base_model import BaseModel from . import networks + class Pix2PixModel(BaseModel): def name(self): return 'Pix2PixModel' @@ -23,12 +24,12 @@ class Pix2PixModel(BaseModel): # load/define networks self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, - opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids) + opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids) if self.isTrain: use_sigmoid = opt.no_lsgan self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, - opt.which_model_netD, - opt.n_layers_D, opt.norm, use_sigmoid, self.gpu_ids) + opt.which_model_netD, + opt.n_layers_D, opt.norm, use_sigmoid, self.gpu_ids) if not self.isTrain or opt.continue_train: self.load_network(self.netG, 'G', opt.which_epoch) if self.isTrain: @@ -71,7 +72,7 @@ class Pix2PixModel(BaseModel): self.fake_B = self.netG.forward(self.real_A) self.real_B = Variable(self.input_B, volatile=True) - #get image paths + # get image paths def get_image_paths(self): return self.image_paths @@ -83,7 +84,7 @@ class Pix2PixModel(BaseModel): self.loss_D_fake = self.criterionGAN(self.pred_fake, False) # Real - real_AB = torch.cat((self.real_A, self.real_B), 1)#.detach() + real_AB = torch.cat((self.real_A, self.real_B), 1) self.pred_real = self.netD.forward(real_AB) self.loss_D_real = self.criterionGAN(self.pred_real, True) @@ -118,10 +119,10 @@ class Pix2PixModel(BaseModel): def get_current_errors(self): return OrderedDict([('G_GAN', self.loss_G_GAN.data[0]), - ('G_L1', self.loss_G_L1.data[0]), - ('D_real', self.loss_D_real.data[0]), - ('D_fake', self.loss_D_fake.data[0]) - ]) + ('G_L1', self.loss_G_L1.data[0]), + ('D_real', self.loss_D_real.data[0]), + ('D_fake', self.loss_D_fake.data[0]) + ]) def get_current_visuals(self): real_A = util.tensor2im(self.real_A.data) diff --git a/models/test_model.py b/models/test_model.py index a356263..65aa088 100644 --- a/models/test_model.py +++ b/models/test_model.py @@ -10,37 +10,32 @@ class TestModel(BaseModel): 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) - nb = opt.batchSize - size = opt.fineSize - self.input_A = self.Tensor(nb, opt.input_nc, size, size) - - assert(not self.isTrain) - self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, - opt.ngf, opt.which_model_netG, - opt.norm, opt.use_dropout, - self.gpu_ids) + self.netG = networks.define_G(opt.input_nc, opt.output_nc, + opt.ngf, opt.which_model_netG, + opt.norm, opt.use_dropout, + self.gpu_ids) which_epoch = opt.which_epoch - #AtoB = self.opt.which_direction == 'AtoB' - #which_network = 'G_A' if AtoB else 'G_B' - self.load_network(self.netG_A, 'G', which_epoch) + self.load_network(self.netG, 'G', which_epoch) print('---------- Networks initialized -------------') - networks.print_network(self.netG_A) + networks.print_network(self.netG) print('-----------------------------------------------') def set_input(self, input): - AtoB = self.opt.which_direction == 'AtoB' - input_A = input['A' if AtoB else 'B'] + # 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' if AtoB else 'B_paths'] + self.image_paths = input['A_paths'] def test(self): self.real_A = Variable(self.input_A) - self.fake_B = self.netG_A.forward(self.real_A) + self.fake_B = self.netG.forward(self.real_A) - #get image paths + # get image paths def get_image_paths(self): return self.image_paths |