diff options
| author | junyanz <junyanz@berkeley.edu> | 2017-04-18 03:38:47 -0700 |
|---|---|---|
| committer | junyanz <junyanz@berkeley.edu> | 2017-04-18 03:38:47 -0700 |
| commit | c99ce7c4e781712e0252c6127ad1a4e8021cc489 (patch) | |
| tree | ba99dfd56a47036d9c1f18620abf4efc248839ab /models/models.py | |
first commit
Diffstat (limited to 'models/models.py')
| -rw-r--r-- | models/models.py | 15 |
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 |
