summaryrefslogtreecommitdiff
path: root/models/models.py
diff options
context:
space:
mode:
authorjunyanz <junyanz@berkeley.edu>2017-04-18 03:38:47 -0700
committerjunyanz <junyanz@berkeley.edu>2017-04-18 03:38:47 -0700
commitc99ce7c4e781712e0252c6127ad1a4e8021cc489 (patch)
treeba99dfd56a47036d9c1f18620abf4efc248839ab /models/models.py
first commit
Diffstat (limited to 'models/models.py')
-rw-r--r--models/models.py15
1 files changed, 15 insertions, 0 deletions
diff --git a/models/models.py b/models/models.py
new file mode 100644
index 0000000..7e790d0
--- /dev/null
+++ b/models/models.py
@@ -0,0 +1,15 @@
+
+def create_model(opt):
+ model = None
+ print(opt.model)
+ if opt.model == 'cycle_gan':
+ from .cycle_gan_model import CycleGANModel
+ assert(opt.align_data == False)
+ model = CycleGANModel()
+ if opt.model == 'pix2pix':
+ from .pix2pix_model import Pix2PixModel
+ assert(opt.align_data == True)
+ model = Pix2PixModel()
+ model.initialize(opt)
+ print("model [%s] was created" % (model.name()))
+ return model