diff options
Diffstat (limited to 'models')
| -rw-r--r-- | models/__init__.py | 20 | ||||
| -rw-r--r-- | models/models.py | 20 |
2 files changed, 20 insertions, 20 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 diff --git a/models/models.py b/models/models.py deleted file mode 100644 index 39cc020..0000000 --- a/models/models.py +++ /dev/null @@ -1,20 +0,0 @@ -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 ValueError("Model [%s] not recognized." % opt.model) - model.initialize(opt) - print("model [%s] was created" % (model.name())) - return model |