diff options
| -rw-r--r-- | live-mogrify.py | 20 |
1 files changed, 14 insertions, 6 deletions
diff --git a/live-mogrify.py b/live-mogrify.py index 85a9ee2..43c1fc1 100644 --- a/live-mogrify.py +++ b/live-mogrify.py @@ -128,15 +128,23 @@ def process_image(opt, data_opt, im): def list_checkpoints(): print("> list checkpoints") return sorted([f.split('/')[2] for f in glob.glob('./checkpoints/*/latest_net_G.pth')]) + def list_checkpoint_dir(path): print("> list checkpoint path {}".format(path)) if not os.path.exists(os.path.join('./checkpoints/', path)): return "not found" return sorted([f.split('/')[3].replace('_net_G.pth','') for f in glob.glob('./checkpoints/' + path + '/*_net_G.pth')]) -def list_datasets(): - print("> list datasets") - # return glob.glob('./checkpoints/*/latest_net_G*') - return ['coming soon...'] + +def list_sequences(): + print("> list sequences") + sequences = [name for name in os.listdir('./sequences') if os.path.isdir(name)] + results = [] + for path in sequences: + count = len([name for name in os.listdir(os.path.join('./sequences/', path)) if os.path.isfile(name)]) + results.append({ + 'name': path, + 'count': count, + }) class Listener(): def __init__(self): @@ -170,8 +178,8 @@ class Listener(): return list_checkpoints() if cmd == 'list_checkpoint_dir': return list_checkpoint_dir(payload) - if cmd == 'list_datasets': - return list_datasets() + if cmd == 'list_sequences': + return list_sequences() if cmd == 'play' and self.data_opt.processing is False: self.data_opt.pause = False process_live_input(self.opt, self.data_opt, self.rpc_client, self.model) |