diff options
Diffstat (limited to 'models/__init__.py')
| -rw-r--r-- | models/__init__.py | 20 |
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 |
