Skip to content

Instantly share code, notes, and snippets.

@woctezuma
Forked from rolux/age.py
Created Sep 7, 2020
Embed
What would you like to do?
# 1. Set up StyleGAN
import dnnlib
import dnnlib.tflib as tflib
import pretrained_networks
network_pkl = 'gdrive:networks/stylegan2-ffhq-config-f.pkl'
_G, _D, Gs = pretrained_networks.load_networks(network_pkl)
Gs_kwargs = dnnlib.EasyDict()
Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_kwargs.randomize_noise = False
Gs_syn_kwargs = dnnlib.EasyDict()
Gs_syn_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_syn_kwargs.randomize_noise = False
Gs_syn_kwargs.minibatch_size = 4
noise_vars = [
var for name, var in Gs.components.synthesis.vars.items()
if name.startswith('noise')
]
w_avg = Gs.get_var('dlatent_avg')
truncation_psi = 0.75
# 2. Get a vector
# https://github.com/a312863063/generators-with-stylegan2/blob/master/latent_directions/age.npy
# 3. Render results
import numpy as np
import PIL.Image
z = np.random.RandomState(5616).randn(1, 512)
w = Gs.components.mapping.run(z, None)
w = w_avg + (w - w_avg) * truncation_psi
v_age = np.load('age.npy')
n = 5
size = 256
canvas = PIL.Image.new('RGB', (n * size, size))
for i, v in enumerate(np.linspace(-10, 10, n)):
w_age = w + v * v_age
image = Gs.components.synthesis.run(w_age, **Gs_syn_kwargs)[0]
image = PIL.Image.fromarray(image)
image = image.resize((size, size), PIL.Image.LANCZOS)
canvas.paste(image, (i * size, 0))
canvas.save('age.png')
@woctezuma
Copy link
Author

woctezuma commented Sep 7, 2020

As mentioned here, to use your own image instead of a random one:

  1. find the latent representation of your image, via projection to latent space,
  2. replace:
z = np.random.RandomState(5616).randn(1, 512)
w = Gs.components.mapping.run(z, None)
w = w_avg + (w - w_avg) * truncation_psi

with:

w = np.load('path/to/my/latent_representations/filename.npy')

@woctezuma
Copy link
Author

woctezuma commented Sep 11, 2020

The code above is for:

  • expression transfer (adding a vector and a scaled difference vector),

but could be adapted for:

  • morphing (linear interpolation).

For style transfer (crossover), please refer to the code in the official repository.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment