From 0c0c7553e6bc7b136296c328ecaf65488110311d Mon Sep 17 00:00:00 2001 From: Diego Porres Date: Mon, 28 Dec 2020 12:04:45 +0100 Subject: [PATCH 1/3] Added style mixing video and some minor code changes --- generate.py | 6 +- style_mixing.py | 223 ++++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 211 insertions(+), 18 deletions(-) diff --git a/generate.py b/generate.py index 7c4ec3b4..659d546e 100755 --- a/generate.py +++ b/generate.py @@ -86,7 +86,6 @@ def generate_images(network_pkl, seeds, truncation_psi, outdir, class_idx=None, fname = f'{outdir}/dlatent{i:02d}.png' print (f'Saved {fname}') PIL.Image.fromarray(img, 'RGB').save(fname) - return # Render images for dlatents initialized from random seeds. Gs_kwargs = { @@ -106,13 +105,14 @@ def generate_images(network_pkl, seeds, truncation_psi, outdir, class_idx=None, print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds))) rnd = np.random.RandomState(seed) z = rnd.randn(1, *Gs.input_shape[1:]) # [minibatch, component] - if(fixnoise): + if fixnoise: noise_rnd = np.random.RandomState(1) # fix noise tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width] else: tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width] image = Gs.run(z, label, **Gs_kwargs) # [minibatch, height, width, channel] - images.append(image[0]) + if grid: + images.append(image[0]) PIL.Image.fromarray(image[0], 'RGB').save(f'{outdir}/seed{seed:04d}.png') if(save_vector): np.save(f'{outdir}/vectors/seed{seed:04d}',z) diff --git a/style_mixing.py b/style_mixing.py index 7d183f85..c9becb44 100755 --- a/style_mixing.py +++ b/style_mixing.py @@ -12,13 +12,22 @@ import os import pickle import re +import sys import numpy as np import PIL.Image +import scipy import dnnlib import dnnlib.tflib as tflib +os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide" +import moviepy.editor + +import warnings # mostly numpy warnings for me +warnings.filterwarnings('ignore', category=FutureWarning) +warnings.filterwarnings('ignore', category=DeprecationWarning) + #---------------------------------------------------------------------------- def style_mixing_example(network_pkl, row_seeds, col_seeds, truncation_psi, col_styles, outdir, minibatch_size=4): @@ -73,17 +82,177 @@ def style_mixing_example(network_pkl, row_seeds, col_seeds, truncation_psi, col_ canvas.paste(PIL.Image.fromarray(image_dict[key], 'RGB'), (W * col_idx, H * row_idx)) canvas.save(f'{outdir}/grid.png') + +# ---------------------------------------------------------------------------- + + +def style_mixing_video(network_pkl, # Path to pretrained model pkl file + src_seed, # Seed of the source image style (row) + dst_seeds, # Seeds of the destination image styles (columns) + col_styles, # Styles to transfer from first row to first column + truncation_psi=1.0, # Truncation trick + only_stylemix=False, # True if user wishes to show only the style transferred result + outdir='out', + duration_sec=30.0, + smoothing_sec=3.0, + mp4_fps=30, + mp4_codec="libx264", + mp4_bitrate="16M", + minibatch_size=4): + # Calculate the number of frames: + num_frames = int(np.rint(duration_sec * mp4_fps)) + # Initialize TensorFlow + tflib.init_tf() + os.makedirs(outdir, exist_ok=True) + + print(f'Loading networks from "{network_pkl}"...') + with dnnlib.util.open_url(network_pkl) as fp: + _G, _D, Gs = pickle.load(fp) + w_avg = Gs.get_var('dlatent_avg') # [component] + + # Sanity check: styles are actually possible for generated image size + max_style = int(2 * np.log2(Gs.output_shape[-1])) - 3 + assert max(col_styles) <= max_style, f"Maximum col-style allowed: {max_style}" + + Gs_syn_kwargs = { + 'output_transform': dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True), + 'randomize_noise': False, + 'minibatch_size': minibatch_size + } + + # First column latents + print('Generating source W vectors...') + src_shape = [num_frames] + Gs.input_shape[1:] + src_z = np.random.RandomState(*src_seed).randn(*src_shape).astype(np.float32) # [frames, src, component] + src_z = scipy.ndimage.gaussian_filter( + src_z, + [smoothing_sec * mp4_fps] + [0] * (len(Gs.input_shape) - 1), + mode="wrap" + ) + src_z /= np.sqrt(np.mean(np.square(src_z))) + # Map into the detangled latent space W and do truncation trick + src_w = Gs.components.mapping.run(src_z, None) + src_w = w_avg + (src_w - w_avg) * truncation_psi + + # First row latents + print('Generating destination W vectors...') + dst_z = np.stack([np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in dst_seeds]) + dst_w = Gs.components.mapping.run(dst_z, None) + dst_w = w_avg + (dst_w - w_avg) * truncation_psi + # Get the width and height of each image: + _N, _C, H, W = Gs.output_shape + + # Generate ALL the source images: + src_images = Gs.components.synthesis.run(src_w, **Gs_syn_kwargs) + # Generate the column images: + dst_images = Gs.components.synthesis.run(dst_w, **Gs_syn_kwargs) + + # If the user wishes to show both the source and destination images + if not only_stylemix: + print('Generating full video (including source and destination images)') + # Generate our canvas where we will paste all the generated images: + canvas = PIL.Image.new("RGB", (W * (len(dst_seeds) + 1), H * (len(src_seed) + 1)), 'black') + + for col, dst_image in enumerate(list(dst_images)): + canvas.paste(PIL.Image.fromarray(dst_image, "RGB"), ((col + 1) * H, 0)) + + # Paste them using an aux function for moviepy frame generation + def make_frame(t): + # Get the frame number according to time t: + frame_idx = int(np.clip(np.round(t * mp4_fps), 0, num_frames - 1)) + # We wish the image belonging to the frame at time t: + src_image = src_images[frame_idx] + # Paste it to the lower left: + canvas.paste(PIL.Image.fromarray(src_image, "RGB"), (0, H)) + + # Now, for each of the column images: + for col, _ in enumerate(list(dst_images)): + # Select the pertinent latent w column: + w_col = np.stack([dst_w[col]]) # [18, 512] -> [1, 18, 512] for 1024x1024 images + # Replace the values defined by col_styles: + w_col[:, col_styles] = src_w[frame_idx, col_styles] + # Generate these synthesized images: + col_images = Gs.components.synthesis.run(w_col, **Gs_syn_kwargs) + # Paste them in their respective spot: + for row, image in enumerate(list(col_images)): + canvas.paste( + PIL.Image.fromarray(image, "RGB"), + ((col + 1) * H, (row + 1) * W), + ) + return np.array(canvas) + # Else, show only the style-transferred images (this is nice for the 1x1 case) + else: + print('Generating only the style-transferred images') + # Generate our canvas where we will paste all the generated images: + canvas = PIL.Image.new("RGB", (W * len(dst_seeds), H * len(src_seed)), "white") + + def make_frame(t): + # Get the frame number according to time t: + frame_idx = int(np.clip(np.round(t * mp4_fps), 0, num_frames - 1)) + # Now, for each of the column images: + for col, _ in enumerate(list(dst_images)): + # Select the pertinent latent w column: + w_col = np.stack([dst_w[col]]) # [18, 512] -> [1, 18, 512] + # Replace the values defined by col_styles: + w_col[:, col_styles] = src_w[frame_idx, col_styles] + # Generate these synthesized images: + col_images = Gs.components.synthesis.run(w_col, **Gs_syn_kwargs) + # Paste them in their respective spot: + for row, image in enumerate(list(col_images)): + canvas.paste( + PIL.Image.fromarray(image, "RGB"), + (col * H, row * W), + ) + return np.array(canvas) + # Generate video using make_frame: + print('Generating style-mixed video...') + videoclip = moviepy.editor.VideoClip(make_frame, duration=duration_sec) + grid_size = [len(dst_seeds), len(src_seed)] + mp4 = "{}x{}-style-mixing.mp4".format(*grid_size) + videoclip.write_videofile(os.path.join(outdir, mp4), + fps=mp4_fps, + codec=mp4_codec, + bitrate=mp4_bitrate) + #---------------------------------------------------------------------------- +# My extended version of this helper function: def _parse_num_range(s): - '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.''' - - range_re = re.compile(r'^(\d+)-(\d+)$') - m = range_re.match(s) - if m: - return list(range(int(m.group(1)), int(m.group(2))+1)) - vals = s.split(',') - return [int(x) for x in vals] + ''' + Input: + s (str): Comma separated string of numbers 'a,b,c', a range 'a-c', or + even a combination of both 'a,b-c', 'a-b,c', 'a,b-c,d,e-f,...' + Output: + nums (list): Ordered list of ascending ints in s, with repeating values + deleted (can be modified to not do either of this) + ''' + # Sanity check 0: + # In case there's a space between the numbers (impossible due to argparse, + # but hey, I am that paranoid): + s = s.replace(' ', '') + # Split w.r.t comma + str_list = s.split(',') + nums = [] + for el in str_list: + if '-' in el: + # The range will be 'a-b', so we wish to find both a and b using re: + range_re = re.compile(r'^(\d+)-(\d+)$') + match = range_re.match(el) + # We get the two numbers: + a = int(match.group(1)) + b = int(match.group(2)) + # Sanity check 1: accept 'a-b' or 'b-a', with a<=b: + if a <= b: r = [n for n in range(a, b + 1)] + else: r = [n for n in range(b, a + 1)] + # Use extend since r will also be an array: + nums.extend(r) + else: + # It's a single number, so just append it: + nums.append(int(el)) + # Sanity check 2: delete repeating numbers: + nums = list(set(nums)) + # Return the numbers in ascending order: + return sorted(nums) #---------------------------------------------------------------------------- @@ -102,15 +271,39 @@ def main(): formatter_class=argparse.RawDescriptionHelpFormatter ) - parser.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True) - parser.add_argument('--rows', dest='row_seeds', type=_parse_num_range, help='Random seeds to use for image rows', required=True) - parser.add_argument('--cols', dest='col_seeds', type=_parse_num_range, help='Random seeds to use for image columns', required=True) - parser.add_argument('--styles', dest='col_styles', type=_parse_num_range, help='Style layer range (default: %(default)s)', default='0-6') - parser.add_argument('--trunc', dest='truncation_psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5) - parser.add_argument('--outdir', help='Where to save the output images', required=True, metavar='DIR') + subparsers = parser.add_subparsers(help='Sub-commands', dest='command') + + parser_grid = subparsers.add_parser('grid', help='Generate a grid of style-mixed images') + parser_grid.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True) + parser_grid.add_argument('--row-seeds', dest='row_seeds', type=_parse_num_range, help='Random seeds to use for image rows', required=True) + parser_grid.add_argument('--col-seeds', dest='col_seeds', type=_parse_num_range, help='Random seeds to use for image columns', required=True) + parser_grid.add_argument('--col-styles', dest='col_styles', type=_parse_num_range, help='Style layer range (default: %(default)s)', default='0-6') + parser_grid.add_argument('--trunc', dest='truncation_psi', type=float, help='Truncation psi (default: %(default)s)', default=0.5) + parser_grid.add_argument('--outdir', help='Where to save the output images', required=True, metavar='DIR') + parser_grid.set_defaults(func=style_mixing_example) + + parser_video = subparsers.add_parser('video', help='Generate style-mixing video (using lerp)') + parser_video.add_argument('--network', help='Path to network pickle filename', dest='network_pkl', required=True) + parser_video.add_argument('--row-seed', type=int, help='Random seed to use for image source row (content)', dest='src_seed', required=True) + parser_video.add_argument('--col-seeds', type=_parse_num_range, help='Random seeds to use for image columns (styles)', dest='dst_seeds', required=True) + parser_video.add_argument('--col-styles', type=_parse_num_range, help='Style layer range (default: %(default)s)', default='0-6', dest='col_styles') + parser_video.add_argument('--only-stylemix', action='store_true', help='Add flag to only save the style-mixed images in the video', dest='only_stylemix') + parser_video.add_argument('--trunc', type=float, help='Truncation psi (default: %(default)s)', default=0.7, dest='truncation_psi') + parser_video.add_argument('--duration', type=float, help='Duration of video in seconds (default: %(default)s)', default=30, dest='duration_sec') + parser_video.add_argument('--fps', type=int, help='FPS of generated video (default: %(default)s)', default=30, dest='mp4_fps') + parser_video.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR') + parser_video.set_defaults(func=style_mixing_video) args = parser.parse_args() - style_mixing_example(**vars(args)) + kwargs = vars(args) + submd = kwargs.pop('command') + + if submd is None: + print('Error: missing subcommand. Re-run with --help for usage.') + sys.exit(1) + + func = kwargs.pop('func') + func(**kwargs) #---------------------------------------------------------------------------- From 2ea2a09c65c251e9d99545fbd3ebb8750babd75b Mon Sep 17 00:00:00 2001 From: Diego Porres Date: Tue, 29 Dec 2020 11:45:15 +0100 Subject: [PATCH 2/3] Minor bug fixes --- generate.py | 89 +++++++++++++++++++++++-------------------------- style_mixing.py | 23 +++++++------ 2 files changed, 55 insertions(+), 57 deletions(-) diff --git a/generate.py b/generate.py index 659d546e..4adde4ec 100755 --- a/generate.py +++ b/generate.py @@ -167,45 +167,45 @@ def truncation_traversal(network_pkl,npys,outdir,class_idx=None, seed=[0],start= #---------------------------------------------------------------------------- def valmap(value, istart, istop, ostart, ostop): - return ostart + (ostop - ostart) * ((value - istart) / (istop - istart)) + return ostart + (ostop - ostart) * ((value - istart) / (istop - istart)) class OSN(): - min=-1 - max= 1 + min=-1 + max= 1 - def __init__(self,seed,diameter): - self.tmp = OpenSimplex(seed) - self.d = diameter - self.x = 0 - self.y = 0 + def __init__(self,seed,diameter): + self.tmp = OpenSimplex(seed) + self.d = diameter + self.x = 0 + self.y = 0 - def get_val(self,angle): - self.xoff = valmap(np.cos(angle), -1, 1, self.x, self.x + self.d); - self.yoff = valmap(np.sin(angle), -1, 1, self.y, self.y + self.d); - return self.tmp.noise2d(self.xoff,self.yoff) + def get_val(self,angle): + self.xoff = valmap(np.cos(angle), -1, 1, self.x, self.x + self.d) + self.yoff = valmap(np.sin(angle), -1, 1, self.y, self.y + self.d) + return self.tmp.noise2d(self.xoff,self.yoff) def get_noiseloop(endpoints, nf, d, start_seed): features = [] zs = [] for i in range(512): - features.append(OSN(i+start_seed,d)) + features.append(OSN(i+start_seed,d)) inc = (np.pi*2)/nf for f in range(nf): - z = np.random.randn(1, 512) - for i in range(512): - z[0,i] = features[i].get_val(inc*f) - zs.append(z) + z = np.random.randn(1, 512) + for i in range(512): + z[0,i] = features[i].get_val(inc*f) + zs.append(z) return zs def line_interpolate(zs, steps): - out = [] - for i in range(len(zs)-1): - for index in range(steps): - fraction = index/float(steps) - out.append(zs[i+1]*fraction + zs[i]*(1-fraction)) - return out + out = [] + for i in range(len(zs)-1): + for index in range(steps): + fraction = index/float(steps) + out.append(zs[i+1]*fraction + zs[i]*(1-fraction)) + return out def generate_zs_from_seeds(seeds,Gs): zs = [] @@ -237,9 +237,9 @@ def generate_latent_images(zs, truncation_psi, outdir, save_npy,prefix,vidname,f for z_idx, z in enumerate(zs): if isinstance(z,list): - z = np.array(z).reshape(1,512) + z = np.array(z).reshape(1,512) elif isinstance(z,np.ndarray): - z.reshape(1,512) + z.reshape(1,512) print('Generating image for step %d/%d ...' % (z_idx, len(zs))) noise_rnd = np.random.RandomState(1) # fix noise tflib.set_vars({var: noise_rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width] @@ -249,7 +249,7 @@ def generate_latent_images(zs, truncation_psi, outdir, save_npy,prefix,vidname,f np.save(f'{outdir}/vectors/{prefix}{z_idx:05d}.npz',z) # np.savetxt(f'{outdir}/vectors/{prefix}{z_idx:05d}.txt',z) - cmd="ffmpeg -y -r {} -i {}/frames/{}%05d.png -vcodec libx264 -pix_fmt yuv420p {}/walk-{}-{}fps.mp4".format(framerate,outdir,prefix,outdir,vidname,framerate) + cmd=f"ffmpeg -y -r {framerate} -i {outdir}/frames/{prefix}%05d.png -vcodec libx264 -pix_fmt yuv420p {outdir}/walk-{vidname}-{framerate}fps.mp4" subprocess.call(cmd, shell=True) def generate_images_in_w_space(ws, truncation_psi,outdir,save_npy,prefix,vidname,framerate,class_idx=None): @@ -305,18 +305,16 @@ def generate_latent_walk(network_pkl, truncation_psi, outdir, walk_type, frames, zs = [] ws =[] - # npys specified, let's work with these instead of seeds # npys must be saved as W's (arrays of 18x512) if npys and (len(npys) > 0): ws = npys - wt = walk_type.split('-') if wt[0] == 'line': if seeds and (len(seeds) > 0): - zs = generate_zs_from_seeds(seeds,Gs) + zs = generate_zs_from_seeds(seeds, Gs) if ws == []: number_of_steps = int(frames/(len(zs)-1))+1 @@ -324,15 +322,15 @@ def generate_latent_walk(network_pkl, truncation_psi, outdir, walk_type, frames, number_of_steps = int(frames/(len(ws)-1))+1 if (len(wt)>1 and wt[1] == 'w'): - if ws == []: - for i in range(len(zs)): - ws.append(convertZtoW(zs[i],truncation_psi)) + if ws == []: + for i in range(len(zs)): + ws.append(convertZtoW(zs[i], truncation_psi)) - points = line_interpolate(ws,number_of_steps) - # zpoints = line_interpolate(zs,number_of_steps) + points = line_interpolate(ws,number_of_steps) + # zpoints = line_interpolate(zs,number_of_steps) else: - points = line_interpolate(zs,number_of_steps) + points = line_interpolate(zs,number_of_steps) # from Gene Kogan @@ -349,10 +347,10 @@ def generate_latent_walk(network_pkl, truncation_psi, outdir, walk_type, frames, # w.append(np.asarray(ws[i]).reshape(512,18)) # points = get_latent_interpolation_bspline(ws,frames,3, 20, shuffle=False) # else: - z = [] - for i in range(len(zs)): - z.append(np.asarray(zs[i]).reshape(512)) - points = get_latent_interpolation_bspline(z,frames,3, 20, shuffle=False) + z = [] + for i in range(len(zs)): + z.append(np.asarray(zs[i]).reshape(512)) + points = get_latent_interpolation_bspline(z,frames,3, 20, shuffle=False) # from Dan Shiffman: https://editor.p5js.org/dvs/sketches/Gb0xavYAR elif wt[0] == 'noiseloop': @@ -370,15 +368,15 @@ def generate_latent_walk(network_pkl, truncation_psi, outdir, walk_type, frames, else: seed_out = 'w-' + wt[0] + '-dlatents' - generate_images_in_w_space(points, truncation_psi,outdir,save_vector,'frame', seed_out, framerate) + generate_images_in_w_space(points, truncation_psi, outdir, save_vector,'frame', seed_out, framerate) elif (len(wt)>1 and wt[1] == 'w'): - print('%s is not currently supported in w space, please change your interpolation type' % (wt[0])) + print('%s is not currently supported in w space, please change your interpolation type' % (wt[0])) else: if(len(wt)>1): seed_out = 'z-' + wt[0] + ('-'.join([str(seed) for seed in seeds])) else: seed_out = 'z-' + walk_type + '-seed' +str(start_seed) - generate_latent_images(points, truncation_psi, outdir, save_vector,'frame', seed_out, framerate) + generate_latent_images(points, truncation_psi, outdir, save_vector, 'frame', seed_out, framerate) #---------------------------------------------------------------------------- @@ -603,7 +601,6 @@ def _parse_npy_files(files): file_list = files.split(",") - for f in file_list: # load numpy array arr = np.load(f) @@ -612,8 +609,6 @@ def _parse_npy_files(files): arr = arr['dlatents'] zs.append(arr) - - return zs #---------------------------------------------------------------------------- @@ -675,7 +670,7 @@ def main(): parser_generate_latent_walk.add_argument('--trunc', type=float, help='Truncation psi (default: %(default)s)', dest='truncation_psi', default=0.5) parser_generate_latent_walk.add_argument('--walk-type', help='Type of walk (default: %(default)s)', default='line') parser_generate_latent_walk.add_argument('--frames', type=int, help='Frame count (default: %(default)s', default=240) - parser_generate_latent_walk.add_argument('--fps', type=int, help='Starting value',default=24,dest='framerate') + parser_generate_latent_walk.add_argument('--fps', type=int, help='Starting value', default=24, dest='framerate') parser_generate_latent_walk.add_argument('--seeds', type=_parse_num_range, help='List of random seeds') parser_generate_latent_walk.add_argument('--npys', type=_parse_npy_files, help='List of .npy files') parser_generate_latent_walk.add_argument('--save_vector', dest='save_vector', action='store_true', help='also save vector in .npy format') @@ -695,7 +690,7 @@ def main(): parser_generate_neighbors.add_argument('--outdir', help='Root directory for run results (default: %(default)s)', default='out', metavar='DIR') parser_generate_neighbors.set_defaults(func=generate_neighbors) - parser_lerp_video = subparsers.add_parser('lerp-video', help='Generate interpolation video (lerp) between random vectors') + parser_lerp_video = subparsers.add_parser('lerp-video', help='Generate interpolation video (lerp; closed-loop) between random vectors') parser_lerp_video.add_argument('--network', help='Path to network pickle filename', dest='network_pkl', required=True) parser_lerp_video.add_argument('--seeds', type=_parse_num_range_ext, help='List of random seeds', dest='seeds', required=True) parser_lerp_video.add_argument('--grid-w', type=int, help='Video grid width/columns (default: %(default)s)', default=None, dest='grid_w') diff --git a/style_mixing.py b/style_mixing.py index c9becb44..78c4e87e 100755 --- a/style_mixing.py +++ b/style_mixing.py @@ -91,7 +91,7 @@ def style_mixing_video(network_pkl, # Path to pretrained model pkl f dst_seeds, # Seeds of the destination image styles (columns) col_styles, # Styles to transfer from first row to first column truncation_psi=1.0, # Truncation trick - only_stylemix=False, # True if user wishes to show only the style transferred result + only_stylemix=False, # Show only the style transferred result outdir='out', duration_sec=30.0, smoothing_sec=3.0, @@ -123,7 +123,7 @@ def style_mixing_video(network_pkl, # Path to pretrained model pkl f # First column latents print('Generating source W vectors...') src_shape = [num_frames] + Gs.input_shape[1:] - src_z = np.random.RandomState(*src_seed).randn(*src_shape).astype(np.float32) # [frames, src, component] + src_z = np.random.RandomState(*[src_seed]).randn(*src_shape).astype(np.float32) # [frames, src, component] src_z = scipy.ndimage.gaussian_filter( src_z, [smoothing_sec * mp4_fps] + [0] * (len(Gs.input_shape) - 1), @@ -151,7 +151,7 @@ def style_mixing_video(network_pkl, # Path to pretrained model pkl f if not only_stylemix: print('Generating full video (including source and destination images)') # Generate our canvas where we will paste all the generated images: - canvas = PIL.Image.new("RGB", (W * (len(dst_seeds) + 1), H * (len(src_seed) + 1)), 'black') + canvas = PIL.Image.new("RGB", (W * (len(dst_seeds) + 1), H * (len([src_seed]) + 1)), 'black') for col, dst_image in enumerate(list(dst_images)): canvas.paste(PIL.Image.fromarray(dst_image, "RGB"), ((col + 1) * H, 0)) @@ -180,11 +180,12 @@ def make_frame(t): ((col + 1) * H, (row + 1) * W), ) return np.array(canvas) + mp4_name = f'{len(dst_seeds)}x1-style-mixing.mp4' # Else, show only the style-transferred images (this is nice for the 1x1 case) else: print('Generating only the style-transferred images') # Generate our canvas where we will paste all the generated images: - canvas = PIL.Image.new("RGB", (W * len(dst_seeds), H * len(src_seed)), "white") + canvas = PIL.Image.new("RGB", (W * len(dst_seeds), H * len([src_seed])), 'black') def make_frame(t): # Get the frame number according to time t: @@ -204,12 +205,11 @@ def make_frame(t): (col * H, row * W), ) return np.array(canvas) + mp4_name = f'{len(dst_seeds)}x1-onlystylemix.mp4' # Generate video using make_frame: print('Generating style-mixed video...') videoclip = moviepy.editor.VideoClip(make_frame, duration=duration_sec) - grid_size = [len(dst_seeds), len(src_seed)] - mp4 = "{}x{}-style-mixing.mp4".format(*grid_size) - videoclip.write_videofile(os.path.join(outdir, mp4), + videoclip.write_videofile(os.path.join(outdir, mp4_name), fps=mp4_fps, codec=mp4_codec, bitrate=mp4_bitrate) @@ -257,9 +257,12 @@ def _parse_num_range(s): #---------------------------------------------------------------------------- _examples = '''examples: - - python %(prog)s --outdir=out --trunc=1 --rows=85,100,75,458,1500 --cols=55,821,1789,293 \\ - --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl + # Grid of style-transferred images, transferring coarse and middle layers of the model + python %(prog)s grid --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl \\ + --outdir=out --trunc=1 --row-seeds=85,95-100,1500 --col-seeds=55,821,1789,293 --col-styles=0-7 + # Video using lerp (60 seconds at 60 fps) and transferring coarse and fine layers of the model + python %(prog)s video --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada/pretrained/metfaces.pkl \\ + --outdir=out --trunc=1 --row-seed=1 --col-seeds=300-305 --col-styles=0-3,8-15 --duration=60 --fps=60 ''' #---------------------------------------------------------------------------- From b5ae20377656567cd273cff7c97d62582a657703 Mon Sep 17 00:00:00 2001 From: Diego Porres Date: Thu, 31 Dec 2020 11:29:07 +0100 Subject: [PATCH 3/3] Set number of steps + save every step dlatent/img --- projector.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/projector.py b/projector.py index 8f6be7e7..54ab72d4 100755 --- a/projector.py +++ b/projector.py @@ -22,8 +22,8 @@ import dnnlib.tflib as tflib class Projector: - def __init__(self): - self.num_steps = 1000 + def __init__(self, num_steps=1000): + self.num_steps = num_steps self.dlatent_avg_samples = 10000 self.initial_learning_rate = 0.1 self.initial_noise_factor = 0.05 @@ -202,7 +202,7 @@ def images_uint8(self): #---------------------------------------------------------------------------- -def project(network_pkl: str, target_fname: str, outdir: str, save_video: bool, seed: int): +def project(network_pkl: str, target_fname: str, outdir: str, save_video: bool, seed: int, num_steps: int, save_every_step: bool): # Load networks. tflib.init_tf({'rnd.np_random_seed': seed}) print('Loading networks from "%s"...' % network_pkl) @@ -220,12 +220,14 @@ def project(network_pkl: str, target_fname: str, outdir: str, save_video: bool, target_float = target_uint8.astype(np.float32).transpose([2, 0, 1]) * (2 / 255) - 1 # Initialize projector. - proj = Projector() + proj = Projector(num_steps) proj.set_network(Gs) proj.start([target_float]) # Setup output directory. os.makedirs(outdir, exist_ok=True) + if save_every_step: + os.makedirs(f'{outdir}/steps', exist_ok=True) target_pil.save(f'{outdir}/target.png') writer = None if save_video: @@ -237,6 +239,9 @@ def project(network_pkl: str, target_fname: str, outdir: str, save_video: bool, assert step == proj.cur_step if writer is not None: writer.append_data(np.concatenate([target_uint8, proj.images_uint8[0]], axis=1)) + if save_every_step: + PIL.Image.fromarray(proj.images_uint8[0], 'RGB').save(f'{outdir}/steps/step_{step:04d}.jpg') + np.save(f'{outdir}/steps/step_{step:04d}.npy', proj.dlatents) dist, loss = proj.step() t.set_postfix(dist=f'{dist[0]:.4f}', loss=f'{loss:.2f}') @@ -274,11 +279,13 @@ def main(): formatter_class=argparse.RawDescriptionHelpFormatter ) - parser.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True) - parser.add_argument('--target', help='Target image file to project to', dest='target_fname', required=True) - parser.add_argument('--save-video', help='Save an mp4 video of optimization progress (default: true)', type=_str_to_bool, default=True) - parser.add_argument('--seed', help='Random seed', type=int, default=303) - parser.add_argument('--outdir', help='Where to save the output images', required=True, metavar='DIR') + parser.add_argument('--network', help='Network pickle filename', dest='network_pkl', required=True) + parser.add_argument('--target', help='Target image file to project to', dest='target_fname', required=True) + parser.add_argument('--save-video', help='Save an mp4 video of optimization progress (default: true)', type=_str_to_bool, default=True) + parser.add_argument('--seed', help='Random seed', type=int, default=303) + parser.add_argument('--num-steps', help='Number of steps to take (default: %(default)s)', type=int, default=1000) + parser.add_argument('--save-every-step', help='Save all dlatents at every step', action='store_true') + parser.add_argument('--outdir', help='Where to save the output images', required=True, metavar='DIR') project(**vars(parser.parse_args())) #----------------------------------------------------------------------------