update dataset mode

4 files changed, 22 insertions, 23 deletions

diff --git a/models/base_model.py b/models/base_model.py
index 9b92bb4..36ceb43 100644
--- a/models/base_model.py
+++ b/models/base_model.py
@@ -1,6 +1,7 @@
 import os
 import torch
 
+
 class BaseModel():
     def name(self):
         return 'BaseModel'
diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py
index f8c4f9f..6fbb19f 100644
--- a/models/cycle_gan_model.py
+++ b/models/cycle_gan_model.py
@@ -10,6 +10,7 @@ from .base_model import BaseModel
 from . import networks
 import sys
 
+
 class CycleGANModel(BaseModel):
     def name(self):
         return 'CycleGANModel'
@@ -27,18 +28,18 @@ class CycleGANModel(BaseModel):
         # Code (paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
 
         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)
+                                        opt.ngf, opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids)
         self.netG_B = networks.define_G(opt.output_nc, opt.input_nc,
-                                    opt.ngf, opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids)
+                                        opt.ngf, opt.which_model_netG, opt.norm, opt.use_dropout, self.gpu_ids)
 
         if self.isTrain:
             use_sigmoid = opt.no_lsgan
             self.netD_A = networks.define_D(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)
             self.netD_B = networks.define_D(opt.input_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:
             which_epoch = opt.which_epoch
             self.load_network(self.netG_A, 'G_A', which_epoch)
@@ -58,10 +59,8 @@ class CycleGANModel(BaseModel):
             # initialize optimizers
             self.optimizer_G = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()),
                                                 lr=opt.lr, betas=(opt.beta1, 0.999))
-            self.optimizer_D_A = torch.optim.Adam(self.netD_A.parameters(),
-                                                lr=opt.lr, betas=(opt.beta1, 0.999))
-            self.optimizer_D_B = torch.optim.Adam(self.netD_B.parameters(),
-                                                lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizer_D_A = torch.optim.Adam(self.netD_A.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
+            self.optimizer_D_B = torch.optim.Adam(self.netD_B.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
 
             print('---------- Networks initialized -------------')
             networks.print_network(self.netG_A)
@@ -89,9 +88,9 @@ class CycleGANModel(BaseModel):
 
         self.real_B = Variable(self.input_B, volatile=True)
         self.fake_A = self.netG_B.forward(self.real_B)
-        self.rec_B  = self.netG_A.forward(self.fake_A)
+        self.rec_B = self.netG_A.forward(self.fake_A)
 
-    #get image paths
+    # get image paths
     def get_image_paths(self):
         return self.image_paths
 
@@ -114,7 +113,7 @@ class CycleGANModel(BaseModel):
 
     def backward_D_B(self):
         fake_A = self.fake_A_pool.query(self.fake_A)
-        self.loss_D_B =  self.backward_D_basic(self.netD_B, self.real_A, fake_A)
+        self.loss_D_B = self.backward_D_basic(self.netD_B, self.real_A, fake_A)
 
     def backward_G(self):
         lambda_idt = self.opt.identity
@@ -167,7 +166,6 @@ class CycleGANModel(BaseModel):
         self.backward_D_B()
         self.optimizer_D_B.step()
 
-
     def get_current_errors(self):
         D_A = self.loss_D_A.data[0]
         G_A = self.loss_G_A.data[0]
@@ -187,10 +185,10 @@ class CycleGANModel(BaseModel):
     def get_current_visuals(self):
         real_A = util.tensor2im(self.real_A.data)
         fake_B = util.tensor2im(self.fake_B.data)
-        rec_A  = util.tensor2im(self.rec_A.data)
+        rec_A = util.tensor2im(self.rec_A.data)
         real_B = util.tensor2im(self.real_B.data)
         fake_A = util.tensor2im(self.fake_A.data)
-        rec_B  = util.tensor2im(self.rec_B.data)
+        rec_B = util.tensor2im(self.rec_B.data)
         if self.opt.identity > 0.0:
             idt_A = util.tensor2im(self.idt_A.data)
             idt_B = util.tensor2im(self.idt_B.data)
diff --git a/models/models.py b/models/models.py
index 8fea4f4..efcd898 100644
--- a/models/models.py
+++ b/models/models.py
@@ -3,15 +3,16 @@ def create_model(opt):
     model = None
     print(opt.model)
     if opt.model == 'cycle_gan':
+        assert(opt.dataset_mode == 'unaligned')
         from .cycle_gan_model import CycleGANModel
-        #assert(opt.align_data == False)
         model = CycleGANModel()
     elif opt.model == 'pix2pix':
+        assert(opt.dataset_mode == 'aligned')
         from .pix2pix_model import Pix2PixModel
-        assert(opt.align_data == True)
         model = Pix2PixModel()
-    elif opt.model == 'one_direction_test':
-        from .one_direction_test_model import OneDirectionTestModel
+    elif opt.model == 'test':
+        assert(opt.dataset_mode == 'single')
+        from .test_model import TestModel
         model = OneDirectionTestModel()
     else:
         raise ValueError("Model [%s] not recognized." % opt.model)
diff --git a/models/one_direction_test_model.py b/models/test_model.py
index d4f6442..a356263 100644
--- a/models/one_direction_test_model.py
+++ b/models/test_model.py
@@ -5,9 +5,9 @@ from .base_model import BaseModel
 from . import networks
 
 
-class OneDirectionTestModel(BaseModel):
+class TestModel(BaseModel):
     def name(self):
-        return 'OneDirectionTestModel'
+        return 'TestModel'
 
     def initialize(self, opt):
         BaseModel.initialize(self, opt)
@@ -48,4 +48,3 @@ class OneDirectionTestModel(BaseModel):
         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)])
-