more live code clean up and options

4 files changed, 57 insertions, 15 deletions

diff --git a/data/recursive_dataset.py b/data/recursive_dataset.py
index f56b871..0daa169 100644
--- a/data/recursive_dataset.py
+++ b/data/recursive_dataset.py
@@ -14,15 +14,12 @@ class RecursiveDataset(BaseDataset):
         self.dir_A = os.path.join(opt.dataroot)
         self.A_paths = make_dataset(self.dir_A)
         self.A_paths = sorted(self.A_paths)
-
-        self.stuff = [opt.dataroot + 'frame_00000.png']
-
         self.transform = get_transform(opt)
 
     def __getitem__(self, index):
         A_path = os.path.join(self.opt.render_dir, "frame_{:05d}.png".format(index))
         while not os.path.exists(A_path):
-            time.sleep(0.01)
+            time.sleep(self.opt.poll_delay)
         A_img = Image.open(A_path).convert('RGB')
         A = self.transform(A_img)
         if self.opt.which_direction == 'BtoA':
@@ -39,10 +36,5 @@ class RecursiveDataset(BaseDataset):
     def __len__(self):
         return 10000000
 
-    def append(self, path):
-        print('append', path)
-        # next_image = path
-        self.stuff.append(path)
-
     def name(self):
         return 'RecursiveImageDataset'
diff --git a/live-mogrify.py b/live-mogrify.py
index 815ce51..7ebba22 100644
--- a/live-mogrify.py
+++ b/live-mogrify.py
@@ -1,4 +1,5 @@
 import os
+sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../live-cortex/rpc/'))
 from options.test_options import TestOptions
 from options.dataset_options import DatasetOptions
 from data import CreateRecursiveDataLoader
@@ -17,15 +18,17 @@ from datetime import datetime
 import re
 import sys
 import math
-
 import subprocess
 from time import sleep
 
+from rpc import CortexRPC
 
 def clamp(n,a,b):
   return max(a, min(n, b))
+
 def lerp(n,a,b):
   return (b-a)*n+a
+
 def load_opt():
   opt = TestOptions().parse()
   data_opt = DatasetOptions().parse(opt.unknown)
@@ -37,6 +40,7 @@ def load_opt():
   data_opt.tag = get_tag(data_opt)
   opt.render_dir = opt.results_dir + opt.name + "/" + tag + "/"
   return opt, data_opt
+
 def get_tag(data_opt):
   if data_opt.tag == '':
     d = datetime.now()
@@ -47,6 +51,7 @@ def get_tag(data_opt):
     )
   else:
     tag = data_opt.tag
+
 def load_first_frame(opt, data_opt):
   start_img_path = os.path.join(opt.render_dir, "frame_00000.png")
   if data_opt.just_copy:
@@ -109,6 +114,12 @@ def process_image(im):
   return img
 
 def process_live_input():
+  def set_data_opt(key, value):
+    data_opt[key] = value
+  def get_opts():
+    return data_opt
+  rpc_client = CortexRPC(get_opts, set_data_opt)
+
   opt, data_opt = load_opt()
 
   A_offset, A_im, A_dir = load_first_frame(opt, data_opt)
@@ -130,16 +141,23 @@ def process_live_input():
       print('%04d: process image...' % (i))
 
     im = visuals['fake_B']
+    last_path = opt.render_dir + "frame_{:05d}.png".format(i)
     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)
 
-    # save rendered image
-    image_pil = Image.fromarray(im, mode='RGB')
-    image_pil.save(tmp_path)
-    os.rename(tmp_path, current_path)
+    if opt.send_image == 'b':
+      image_pil = Image.fromarray(im, mode='RGB')
+      rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), image_pil)
+
+    if opt.store_a is not True:
+      os.remove(last_path)
+    if opt.store_b is True:
+      image_pil = Image.fromarray(im, mode='RGB')
+      image_pil.save(tmp_path)
+      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:
@@ -172,12 +190,22 @@ def process_live_input():
       last_im = im.copy().astype('uint8')
       next_im = im
 
-    # image_pil = Image.fromarray(im, mode='RGB')
+    # image_pil = Image.fromarray(next_im, mode='RGB')
     # im = np.asarray(image_pil).astype('uint8')
     #print(im.shape, im.dtype)
 
     next_img = process_image(next_im)
 
+    if opt.send_image == 'sequence':
+      rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), A_img)
+    if opt.send_image == 'recursive':
+      pil_im = Image.fromarray(next_im)
+      rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), pil_im)
+    if opt.send_image == 'a':
+      rgb_im = cv2.cvtColor(next_img, cv2.COLOR_BGR2RGB)
+      pil_im = Image.fromarray(rgb_im)
+      rpc_client.send_pil_image("frame_{:05d}.png".format(i+1), pil_im)
+
     cv2.imwrite(tmp_path, next_img)
     os.rename(tmp_path, next_path)
 
diff --git a/options/base_options.py b/options/base_options.py
index f741927..f3e93c6 100644
--- a/options/base_options.py
+++ b/options/base_options.py
@@ -42,6 +42,7 @@ class BaseOptions():
         self.parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
         self.parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
         self.parser.add_argument('--center_crop', action='store_true', help='center crop instead of random crop')
+        self.parser.add_argument('--poll_delay', type=float, default=0.01, help='time to wait before checking for the next frame')
 
         self.initialized = True
 
diff --git a/options/dataset_options.py b/options/dataset_options.py
index 356ce60..a201903 100644
--- a/options/dataset_options.py
+++ b/options/dataset_options.py
@@ -72,6 +72,27 @@ class DatasetOptions(BaseOptions):
           help='scp destination'
         )
 
+        ## LIVE IMAGE PROCESSING
+
+        self.parser.add_argument(
+          '--send_image',
+          type=str,
+          default='b',
+          help='which image to send... a, b, a_mix, sequence'
+        )
+
+        self.parser.add_argument(
+          '--store_a',
+          action='store_true',
+          help='dont remove the generated A image after processing it'
+        )
+
+        self.parser.add_argument(
+          '--store_b',
+          action='store_true',
+          help='after generating a B image, save it to disk'
+        )
+
         ## IMAGE FILTERS
 
         ### RECURSION