diff options
| -rw-r--r-- | live-mogrify.py | 185 | ||||
| -rw-r--r-- | options/dataset_options.py | 7 |
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' |
