refactor mogrify script

2 files changed, 192 insertions, 0 deletions

diff --git a/live-mogrify.py b/live-mogrify.py
new file mode 100644
index 0000000..815ce51
--- /dev/null
+++ b/live-mogrify.py
@@ -0,0 +1,185 @@
+import os
+from options.test_options import TestOptions
+from options.dataset_options import DatasetOptions
+from data import CreateRecursiveDataLoader
+from models import create_model
+# from util.visualizer import Visualizer
+from util.util import mkdirs, crop_image
+from util import html
+from shutil import move, copyfile
+from PIL import Image, ImageOps
+from skimage.transform import resize
+from scipy.misc import imresize
+from shutil import copyfile, rmtree
+import numpy as np
+import cv2
+from datetime import datetime
+import re
+import sys
+import math
+
+import subprocess
+from time import sleep
+
+
+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)
+  opt.nThreads = 1   # test code only supports nThreads = 1
+  opt.batchSize = 1  # test code only supports batchSize = 1
+  opt.serial_batches = True  # no shuffle
+  opt.no_flip = True  # no flip
+
+  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()
+    tag = data_opt.tag = "{}_{}_{}".format(
+      opt.name,
+      'live',
+      d.strftime('%Y%m%d%H%M')
+    )
+  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:
+    copyfile(opt.start_img, start_img_path)
+  else:
+    print("preload {}".format(opt.start_img))
+    A_img = Image.open(opt.start_img).convert('RGB')
+    A_im = np.asarray(A_img)
+    A = process_image(A_im)
+    cv2.imwrite(start_img_path, A)
+
+  numz = re.findall(r'\d+', os.path.basename(opt.start_img))
+  print(numz)
+  if len(numz) > 0:
+    A_offset = int(numz[0])
+    print(A_offset)
+    if A_offset:
+      print(">> starting offset: {}".format(A_offset))
+      A_dir = opt.start_img.replace(numz[0], "{:05d}")
+      print(A_dir)
+    else:
+      print("Sequence not found")
+  return A_offset, A_im, A_dir
+
+def process_image(im):
+  img = im[:, :, ::-1].copy()
+  processed = False
+
+  if data_opt.clahe is True:
+    processed = True
+    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
+    l, a, b = cv2.split(lab)
+    clahe = cv2.createCLAHE(clipLimit=data_opt.clip_limit, tileGridSize=(8,8))
+    l = clahe.apply(l)
+    limg = cv2.merge((l,a,b))
+    img = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
+  if data_opt.posterize is True:
+    processed = True
+    img = cv2.pyrMeanShiftFiltering(img, data_opt.spatial_window, data_opt.color_window)
+  if data_opt.grayscale is True:
+    processed = True
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+  if data_opt.blur is True:
+    processed = True
+    img = cv2.GaussianBlur(img, (data_opt.blur_radius, data_opt.blur_radius), data_opt.blur_sigma)
+  if data_opt.canny is True:
+    processed = True
+    img = cv2.Canny(img, data_opt.canny_lo, data_opt.canny_hi)
+    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+
+  if processed is False or data_opt.process_frac == 0:
+    return img
+
+  src_img = im[:, :, ::-1].copy()
+  frac_a = data_opt.process_frac
+  frac_b = 1.0 - frac_a
+  array_a = np.multiply(src_img.astype('float64'), frac_a)
+  array_b = np.multiply(img.astype('float64'), frac_b)
+  img = np.add(array_a, array_b).astype('uint8')
+  return img
+
+def process_live_input():
+  opt, data_opt = load_opt()
+
+  A_offset, A_im, A_dir = load_first_frame(opt, data_opt)
+
+  data_loader = CreateRecursiveDataLoader(opt)
+  dataset = data_loader.load_data()
+  model = create_model(opt)
+
+  last_im = None
+  for i, data in enumerate(data_loader):
+    if i >= opt.how_many:
+      break
+    model.set_input(data)
+    model.test()
+    visuals = model.get_current_visuals()
+    img_path = model.get_image_paths()
+
+    if (i % 100) == 0:
+      print('%04d: process image...' % (i))
+
+    im = visuals['fake_B']
+    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 data_opt.recursive and last_im is not None:
+      if data_opt.sequence and A_dir is not None:
+        A_img = Image.open(sequence_path).convert('RGB')
+        A_im = np.asarray(A_img)
+        t = lerp(math.sin(i / data_opt.transition_period * math.pi * 2.0 ) / 2.0 + 0.5, data_opt.transition_min, data_opt.transition_max)
+        frac_a = data_opt.recursive_frac * (1.0 - t)
+        frac_b = data_opt.sequence_frac * (1.0 - t)
+        frac_c = 1.0 - frac_a - frac_b
+        array_a = np.multiply(last_im.astype('float64'), frac_a)
+        array_b = np.multiply(A_im.astype('float64'), frac_b)
+        array_c = np.multiply(im.astype('float64'), frac_c)
+        array_ab = np.add(array_a, array_b)
+        array_abc = np.add(array_ab, array_c)
+        next_im = array_abc.astype('uint8')
+
+      else:
+        frac_a = data_opt.recursive_frac
+        frac_b = 1.0 - frac_a
+        array_a = np.multiply(last_im.astype('float64'), frac_a)
+        array_b = np.multiply(im.astype('float64'), frac_b)
+        next_im = np.add(array_a, array_b).astype('uint8')
+
+      if data_opt.recurse_roll != 0:
+        last_im = np.roll(im, data_opt.recurse_roll, axis=data_opt.recurse_roll_axis)
+      else:
+        last_im = next_im.copy().astype('uint8')
+
+    else:
+      last_im = im.copy().astype('uint8')
+      next_im = im
+
+    # image_pil = Image.fromarray(im, mode='RGB')
+    # im = np.asarray(image_pil).astype('uint8')
+    #print(im.shape, im.dtype)
+
+    next_img = process_image(next_im)
+
+    cv2.imwrite(tmp_path, next_img)
+    os.rename(tmp_path, next_path)
+
+if __name__ == '__main__':
+  process_live_input()
\ No newline at end of file
diff --git a/options/dataset_options.py b/options/dataset_options.py
index 368f2b2..356ce60 100644
--- a/options/dataset_options.py
+++ b/options/dataset_options.py
@@ -91,6 +91,13 @@ class DatasetOptions(BaseOptions):
         )
 
         self.parser.add_argument(
+          '--transition-max',
+          default=1e-3,
+          type=float,
+          help='maximum amount of stabilization to apply'
+        )
+
+        self.parser.add_argument(
           '--recursive',
           action='store_true',
           help='recurse on previous output'