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 subprocess from time import sleep if __name__ == '__main__': 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 # opt.experiment = data_opt.experiment # opt.start_img.split("/")[-1].split(".")[0] d = datetime.now() tag = "{}_{}_{}".format( opt.name, opt.experiment, d.strftime('%Y%m%d%H%M')) opt.tag = tag # = "pcfade___201805150250" opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/" A_offset = 0 A_im = None A_dir = None print("create render_dir: {}".format(render_dir)) if os.path.exists(render_dir): rmtree(render_dir) mkdirs(render_dir) # cmd = ("convert", opt.start_img, '-canny', '0x1+10%+30%', render_dir + "frame_00000.png") # process = subprocess.Popen(cmd, stdout=subprocess.PIPE) # output, error = process.communicate() def load_first_frame(): start_img_path = os.path.join(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)) if len(numz) > 0: A_offset = int(numz[0]) if A_offset: print(">> starting offset: {}".format(A_offset)) A_dir = opt.start_img.replace(numz[0], "{:05d}") else: print("Sequence not found") def process_image(im): img = im[:, :, ::-1].copy() if data_opt.clahe is 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: img = cv2.pyrMeanShiftFiltering(img, data_opt.spatial_window, data_opt.color_window) if data_opt.grayscale is True: img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if data_opt.blur is True: img = cv2.GaussianBlur(img, (data_opt.blur_radius, data_opt.blur_radius), data_opt.blur_sigma) if data_opt.canny is True: img = cv2.Canny(img, data_opt.canny_lo, data_opt.canny_hi) return img load_first_frame() data_loader = CreateRecursiveDataLoader(opt) dataset = data_loader.load_data() ds = dataset.dataset model = create_model(opt) visualizer = Visualizer(opt) # create website # web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch)) # webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch)) # test 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 % 20) == 0: print('%04d: process image... %s' % (i, img_path)) # ims = visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio) im = visuals['fake_B'] tmp_path = render_dir + "frame_{:05d}_tmp.png".format(i+1) edges_path = render_dir + "frame_{:05d}.png".format(i+1) render_path = render_dir + "ren_{:05d}.png".format(i+1) if A_dir is not None: sequence_path = A_dir.format(A_offset+i+1) # A_offset # save rendered image image_pil = Image.fromarray(im, mode='RGB') image_pil.save(tmp_path) os.rename(tmp_path, render_path) if dataset.name() == 'RecursiveDatasetDataLoader': if data_opt.recursive and last_im is not None: tmp_im = im.copy() if data_opt.sequence and A_dir is not None: A_img = Image.open(sequence_path).convert('RGB') A_im = np.asarray(A_img) frac_a = data_opt.recursive_frac frac_b = data_opt.sequence_frac 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) comb_ab = np.add(array_a, array_b) comb_abc = np.add(array_ab, array_c) im = comb_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) im = np.add(array_a, array_b).astype('uint8') if data_opt.recurse_roll != 0: last_im = np.roll(tmp_im, data_opt.recurse_roll, axis=data_opt.recurse_roll_axis) else: last_im = im.copy().astype('uint8') else: last_im = im.copy().astype('uint8') tmp_im = im.copy().astype('uint8') #print(im.shape, im.dtype) # image_pil = Image.fromarray(im, mode='RGB') # im = np.asarray(image_pil).astype('uint8') #print(im.shape, im.dtype) img = process_image(im) cv2.imwrite(tmp_path, img) os.rename(tmp_path, edges_path) # webpage.save() # os.remove(render_dir + "frame_00000.png") print(opt.render_dir) video_fn = tag + "_mogrify.mp4" cmd = ("ffmpeg", "-i", render_dir + "ren_%05d.png", "-y", "-c:v", "libx264", "-vf", "fps=30", "-pix_fmt", "yuv420p", "-s", "456x256", render_dir + video_fn) process = subprocess.Popen(cmd, stdout=subprocess.PIPE) output, error = process.communicate() print("________") cmd = ("scp", render_dir + video_fn, "jules@asdf.us:asdf/neural/") process = subprocess.Popen(cmd, stdout=subprocess.PIPE) output, error = process.communicate() print("https://asdf.us/neural/" + video_fn)