more live code clean up and options

1 files changed, 34 insertions, 6 deletions

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)