summaryrefslogtreecommitdiff
path: root/models/__init__.py
diff options
context:
space:
mode:
Diffstat (limited to 'models/__init__.py')
-rw-r--r--models/__init__.py20
1 files changed, 20 insertions, 0 deletions
diff --git a/models/__init__.py b/models/__init__.py
index e69de29..681c6de 100644
--- a/models/__init__.py
+++ b/models/__init__.py
@@ -0,0 +1,20 @@
+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
+ model = CycleGANModel()
+ elif opt.model == 'pix2pix':
+ assert(opt.dataset_mode == 'aligned')
+ from .pix2pix_model import Pix2PixModel
+ model = Pix2PixModel()
+ elif opt.model == 'test':
+ assert(opt.dataset_mode == 'single')
+ from .test_model import TestModel
+ model = TestModel()
+ else:
+ raise NotImplementedError('model [%s] not implemented.' % opt.model)
+ model.initialize(opt)
+ print("model [%s] was created" % (model.name()))
+ return model