cmdz

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)