diff options
| -rw-r--r-- | live-mogrify.py | 8 | ||||
| -rw-r--r-- | options/dataset_options.py | 7 |
2 files changed, 12 insertions, 3 deletions
diff --git a/live-mogrify.py b/live-mogrify.py index 6e62645..89ce10c 100644 --- a/live-mogrify.py +++ b/live-mogrify.py @@ -38,6 +38,8 @@ def load_opt(): opt = opt_parser.parse() data_opt_parser = DatasetOptions() data_opt = data_opt_parser.parse(opt.unknown) + global module_name + module_name = data_opt.module_name opt.nThreads = 1 # test code only supports nThreads = 1 opt.batchSize = 1 # test code only supports batchSize = 1 opt.serial_batches = True # no shuffle @@ -141,10 +143,10 @@ def list_epochs(path): def list_sequences(): print("> list sequences") - sequences = sorted([name for name in os.listdir('./sequences') if os.path.isdir(os.path.join('./sequences/', name))]) + sequences = sorted([name for name in os.listdir(os.path.join('./sequences/', module_name)) if os.path.isdir(os.path.join('./sequences/', module_name, name))]) results = [] for path in sequences: - count = len([name for name in os.listdir(os.path.join('./sequences/', path)) if os.path.isfile(os.path.join('./sequences/', path, name))]) + count = len([name for name in os.listdir(os.path.join('./sequences/', module_name, path)) if os.path.isfile(os.path.join('./sequences/', module_name, path, name))]) results.append({ 'name': path, 'count': count, @@ -153,7 +155,7 @@ def list_sequences(): def read_sequence(path): print("> read sequence {}".format(path)) - return sorted([f for f in glob.glob(os.path.join('./sequences/', path, '*.png'))]) + return sorted([f for f in glob.glob(os.path.join('./sequences/', module_name, path, '*.png'))]) class Listener(): def __init__(self): diff --git a/options/dataset_options.py b/options/dataset_options.py index 7ac5a8a..92b63ec 100644 --- a/options/dataset_options.py +++ b/options/dataset_options.py @@ -75,6 +75,13 @@ class DatasetOptions(BaseOptions): ## LIVE IMAGE PROCESSING self.parser.add_argument( + '--module-name', + type=str, + default='pix2pix', + help='module name... basically says where to look for sequences' + ) + + self.parser.add_argument( '--send-image', type=str, default='b', |