cmdz

1 files changed, 13 insertions, 5 deletions

diff --git a/live-mogrify.py b/live-mogrify.py
index 1999124..a709bdc 100644
--- a/live-mogrify.py
+++ b/live-mogrify.py
@@ -20,6 +20,7 @@ import re
 import sys
 import math
 import subprocess
+import glob
 from time import sleep
 
 from rpc import CortexRPC
@@ -125,9 +126,14 @@ def process_image(opt, data_opt, im):
   return img
 
 def list_checkpoints():
-  pass
-def list_videos():
-  pass
+  return glob.glob('./checkpoints/*/latest_net_G*')
+def list_checkpoint_dir(path):
+  if not os.path.exists('./checkpoints/' + path):
+    return "not found"
+  return glob.glob(os.path.join('./checkpoints/' + path + '/*_net_G.pth'))
+def list_datasets():
+  # return glob.glob('./checkpoints/*/latest_net_G*')
+  return ['coming soon...']
 
 class Listener():
   def __init__(self):
@@ -159,8 +165,10 @@ class Listener():
     print("got command {}".format(cmd))
     if cmd == 'list_checkpoints':
       return list_checkpoints()
-    if cmd == 'list_videos':
-      return list_videos()
+    if cmd == 'list_checkpoint_dir':
+      return list_checkpoint_dir(payload)
+    if cmd == 'list_datasets':
+      return list_datasets()
     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)