cmdz

1 files changed, 58 insertions, 9 deletions

diff --git a/live-mogrify.py b/live-mogrify.py
index 5d4dde9..5ebbb16 100644
--- a/live-mogrify.py
+++ b/live-mogrify.py
@@ -57,16 +57,18 @@ def get_tag(opt, data_opt):
     tag = data_opt.tag
   return tag
 
-def load_first_frame(opt, data_opt):
+def create_render_dir(opt):
   print("create render_dir: {}".format(opt.render_dir))
   if os.path.exists(opt.render_dir):
       rmtree(opt.render_dir)
   mkdirs(opt.render_dir)
-  start_img_path = os.path.join(opt.render_dir, "frame_00000.png")
+
+def load_first_frame(opt, data_opt, i=0):
+  start_img_path = os.path.join(opt.render_dir, "frame_{:05}.png".format(i))
   if data_opt.just_copy:
     copyfile(opt.start_img, start_img_path)
-    A_im = None
     A_img = None
+    A_im = None
     A_offset = 0
   else:
     print("preload {}".format(opt.start_img))
@@ -137,7 +139,7 @@ def list_epochs(path):
 
 def list_sequences():
   print("> list sequences")
-  sequences = [name for name in os.listdir('./sequences') if os.path.isdir(os.path.join('./sequences/', name))]
+  sequences = sorted([name for name in os.listdir('./sequences') if os.path.isdir(os.path.join('./sequences/', 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))])
@@ -147,6 +149,10 @@ def list_sequences():
     })
   return results
 
+def read_sequence(path):
+  print("> read sequence {}".format(path))
+  return sorted([f for f in glob.glob(os.path.join('./sequences/', path, '*.png'))])
+
 class Listener():
   def __init__(self):
     opt, data_opt, data_opt_parser = load_opt()
@@ -181,6 +187,23 @@ class Listener():
       return list_epochs(payload)
     if cmd == 'list_sequences':
       return list_sequences()
+    if cmd == 'load_epoch':
+      name, epoch = payload.split(':')
+      print(">>> loading checkpoint {}, epoch {}".format(name, epoch))
+      self.data_opt.checkpoint = name
+      self.data_opt.epoch = epoch
+      self.data_opt.load_checkpoint = True
+      return 'ok'
+    if cmd == 'load_sequence' and os.path.exists('./sequences/' + payload):
+      self.data_opt.sequence_name = payload
+      self.data_opt.load_sequence = True
+    if cmd == 'get_status':
+      return {
+        'processing': self.data_opt.processing,
+        'checkpoint': self.data_opt.checkpoint,
+        'epoch': self.data_opt.epoch,
+        'sequence': self.data_opt.sequence_name,
+      }
     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)
@@ -204,12 +227,37 @@ def process_live_input(opt, data_opt, rpc_client, model):
   data_loader = CreateRecursiveDataLoader(opt)
   dataset = data_loader.load_data()
 
-  print("generating...")
-  A_offset, A_im, A_dir = load_first_frame(opt, data_opt)
+  create_render_dir(opt)
+  sequence = read_sequence(data_opt.sequence_name)
+  print("Got sequence {}, {} images, first: {}".format(data_out.sequence, len(sequence), sequence[0]))
+  # A_offset, A_im, A_dir = load_first_frame(opt, data_opt, 0)
+  # A_offset, A_im, A_dir = load_first_frame(opt, data_opt, i)
+  if len(sequence) == 0:
+    print("Got empty sequence...")
+    data_opt.processing = False
+    return
+  start_img_path = os.path.join(opt.render_dir, "frame_{:05}.png".format(i))
+  copyfile(sequence[0], start_img_path)
+
   last_im = None
+
+  print("generating...")
+  sequence_i = 1
   for i, data in enumerate(data_loader):
     if i >= opt.how_many:
       break
+    if data_opt.load_checkpoint is True:
+      model.save_dir = os.path.join(self.opt.checkpoints_dir, data_opt.checkpoint)
+      model.load_network(model.netG, 'G', data_opt.epoch)
+      data_opt.load_checkpoint = False
+    if data_opt.load_sequence is True:
+      data_opt.load_sequence = False
+      new_sequence = read_sequence(data_opt.sequence_name)
+      if len(new_sequence) != 0:
+        print("Got sequence {}, {} images, first: {}".format(data_opt.sequence_name, len(sequence), sequence[0]))
+        sequence = new_sequence
+        sequence_i = 1
+
     model.set_input(data)
     model.test()
     visuals = model.get_current_visuals()
@@ -223,8 +271,8 @@ def process_live_input(opt, data_opt, rpc_client, model):
     tmp_path = opt.render_dir + "frame_{:05d}_tmp.png".format(i+1)
     next_path = opt.render_dir + "frame_{:05d}.png".format(i+1)
     current_path = opt.render_dir + "ren_{:05d}.png".format(i+1)
-    if A_dir is not None:
-      sequence_path = A_dir.format(A_offset+i+1)
+    if sequence:
+      sequence_path = sequence[sequence_i]
 
     if data_opt.send_image == 'b':
       image_pil = Image.fromarray(im, mode='RGB')
@@ -238,7 +286,7 @@ def process_live_input(opt, data_opt, rpc_client, model):
       os.rename(tmp_path, current_path)
 
     if data_opt.recursive and last_im is not None:
-      if data_opt.sequence and A_dir is not None:
+      if data_opt.sequence and len(sequence):
         A_img = Image.open(sequence_path).convert('RGB')
         A_im = np.asarray(A_img)
         frac_a = data_opt.recursive_frac
@@ -292,6 +340,7 @@ def process_live_input(opt, data_opt, rpc_client, model):
     if data_opt.pause:
       data_opt.pause = False
       break
+    sequence_i += 1
   data_opt.processing = False
 
 if __name__ == '__main__':