In [ ]:

    

is_stylegan_v1 = False


    

In [ ]:

    

from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import sys
import os
from datetime import datetime
from tqdm import tqdm

# ffmpeg installation location, for creating videos
plt.rcParams['animation.ffmpeg_path'] = str('/usr/bin/ffmpeg')
import ipywidgets as widgets
from ipywidgets import interact, interact_manual
from IPython.display import display
from ipywidgets import Button

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

%load_ext autoreload
%autoreload 2

# StyleGAN2 Repo
sys.path.append('/tf/notebooks/stylegan2')

# StyleGAN Utils
from stylegan_utils import load_network, gen_image_fun, synth_image_fun, create_video
# v1 override
if is_stylegan_v1:
    from stylegan_utils import load_network_v1 as load_network
    from stylegan_utils import gen_image_fun_v1 as gen_image_fun
    from stylegan_utils import synth_image_fun_v1 as synth_image_fun

import run_projector
import projector
import training.dataset
import training.misc

# Data Science Utils
sys.path.append(os.path.join(*[os.pardir]*3, 'data-science-learning'))

from ds_utils import generative_utils


    

In [ ]:

    

res_dir = Path.home() / 'Documents/generated_data/stylegan'


    

In [ ]:

    

MODELS_DIR = Path.home() / 'Documents/models/stylegan2'
MODEL_NAME = 'original_ffhq'
SNAPSHOT_NAME = 'stylegan2-ffhq-config-f'

Gs, Gs_kwargs, noise_vars = load_network(str(MODELS_DIR / MODEL_NAME / SNAPSHOT_NAME) + '.pkl')

Z_SIZE = Gs.input_shape[1]
IMG_SIZE = Gs.output_shape[2:]
IMG_SIZE


    

In [ ]:

    

img = gen_image_fun(Gs, np.random.randn(1, Z_SIZE), Gs_kwargs, noise_vars)
plt.imshow(img)


    

In [ ]:

    

def plot_direction_grid(dlatent, direction, coeffs):
    fig, ax = plt.subplots(1, len(coeffs), figsize=(15, 10), dpi=100)
    
    for i, coeff in enumerate(coeffs):
        new_latent = (dlatent.copy() + coeff*direction)
        ax[i].imshow(synth_image_fun(Gs, new_latent, Gs_kwargs, randomize_noise=False))
        ax[i].set_title(f'Coeff: {coeff:0.1f}')
        ax[i].axis('off')
    plt.show()


    

In [ ]:

    

# load learned direction
direction = np.load('/tf/media/datasets/stylegan/learned_directions.npy')


    

In [ ]:

    

nb_latents = 5

# generate dlatents from mapping network
dlatents = Gs.components.mapping.run(np.random.rand(nb_latents, Z_SIZE), None, truncation_psi=1.)

for i in range(nb_latents):
    plot_direction_grid(dlatents[i:i+1], direction, np.linspace(-2, 2, 5))


    

In [ ]:

    

# Setup plot image
dpi = 100
fig, ax = plt.subplots(dpi=dpi, figsize=(7, 7))
fig.subplots_adjust(left=0, right=1, bottom=0, top=1, hspace=0, wspace=0)
plt.axis('off')
im = ax.imshow(gen_image_fun(Gs, np.random.randn(1, Z_SIZE),Gs_kwargs, noise_vars, truncation_psi=1))

#prevent any output for this cell
plt.close()


    

In [ ]:

    

# fetch attributes names
directions_dir = MODELS_DIR / MODEL_NAME / 'directions' / 'set01'
attributes = [e.stem for e in directions_dir.glob('*.npy')]

# get names or projected images
data_dir = res_dir / 'projection' / MODEL_NAME / SNAPSHOT_NAME / '20200215_192547'
entries = [p.name for p in data_dir.glob("*") if p.is_dir()]
entries.remove('tfrecords')

# set target latent to play with
dlatents = Gs.components.mapping.run(np.random.rand(1, Z_SIZE), None, truncation_psi=0.5)
target_latent = dlatents[0:1]
#target_latent = np.array([np.load("/out_4/image_latents2000.npy")])


    

In [ ]:

    

%matplotlib inline

@interact
def i_direction(attribute=attributes, 
                entry=entries,
                coeff=(-10., 10.)):
    direction = np.load(directions_dir / f'{attribute}.npy')
    target_latent = np.array([np.load(data_dir / entry / "image_latents1000.npy")])
    
    new_latent_vector = target_latent.copy() + coeff*direction
    im.set_data(synth_image_fun(Gs, new_latent_vector, Gs_kwargs, True))
    ax.set_title('Coeff: %0.1f' % coeff)
    display(fig)