StyleGanで福沢諭吉を混ぜてみる

前回StyleGanを少し触ってみた時に実際の実画像を潜在変数に変換するエンコーダーがあれば面白いのにと書きましたが、ググったら普通にありました。
https://github.com/Puzer/stylegan-encoder
なのでこのエンコーダーを試してみるために福沢諭吉の画像=>潜在変数に変換してみました。
先に結論をいうと福沢諭吉のスタイルミキシングは全然上手く行きませんでした。

①StyleGanエンコーダーのgit clone

git clone https://github.com/Puzer/stylegan-encoder

②福沢諭吉の画像=>潜在変数変換

最初に福沢諭吉の256×256の画像を用意します。
自分の場合、以下の画像を準備しました。(./img/yukiti.jpg)

python encode_images.py img/ generated_images/ latent/

次に潜在変数に変換するencode_images.pyを動かします。
ここで先ほどの用意した画像のフォルダパスと出力ファイルパスを渡す必要があります。

引数を毎回入れるのが面倒な場合はencode_images.pyに直接書いても良いです。

encode_images.py

    parser.add_argument('--src_dir', default='./img/', help='Directory with images for encoding')
    parser.add_argument('--generated_images_dir', default='./generated_images/', help='Directory for storing generated images')
    parser.add_argument('--dlatent_dir', default='./latent/', help='Directory for storing dlatent representations')

これを実行すると初期設定(--iterations=1000)だと3分ほどかかってlossが1.2=>0.6まで下がりgenerated_imagesフォルダにyukiti.png、latentフォルダにyukiti.npyが出力されます。
iterationsを増やすと潜在変数の精度が良くなるらしく、--iterations=50000の場合は二時間半ほどかかってlossが1.2=>0.12ぐらいまで下がりました。

結果はちらほらノイズが見受けられますが、遠目には良さそうに見えます。

③福沢諭吉の潜在変数確認

潜在変数はyukiti.npyファイルで与えられるのでこれを読み込んで確認してみます。
するとyukiti.npyは$(18,512)$の配列であることが分かります。512は潜在変数として18はどこから出てきたんだと最初思いましたがどうも18はstyleGANにおけるモデルの各層のことを指しているようです。
つまり一番最初に与えられるrandomseedで与えられる512次元の潜在変数とはちょっと種類が異なるようです。

test.py

dlatents = np.load('./latent/yukiti.npy')
print(dlatents)
print(dlatents.shape)

[[ 0.7956156  -0.3605363   0.09918146 ... -0.26614022  0.18791896
  -0.19557095]
 [ 0.6804809   0.4047316   0.2559202  ...  0.7069059   0.60288507
  -0.00266154]
 [ 0.47569656  2.1551647   1.0866272  ... -1.0127327   0.78042006
  -1.3761303 ]
 ...
 [-0.23147076 -0.02230923  0.52631843 ... -0.30792037 -0.14741592
   0.26553717]
 [ 0.7523304   0.0346575   0.26286486 ... -0.74985385  0.16475876
   0.16340141]
 [ 0.04558849 -0.02860947  0.25467917 ... -0.1821715  -0.06283879
   0.07273757]]
(18, 512)

④福沢諭吉のスタイルミキシングコード

さて福沢諭吉の潜在変数が分かったならこれを使って画像の合成を試してみます。
以下のように本来のdst_dlatentsを生成したyukiti.npyの値に置き換えてみます。

yukiti_mix.py

import os
import pickle
import numpy as np
import PIL.Image
import dnnlib
import dnnlib.tflib as tflib
import config

def load_Gs(url):
    fpath = './karras2019stylegan-ffhq-1024x1024.pkl'
    with open(fpath, mode='rb') as f:
        _G, _D, Gs = pickle.load(f)
    return Gs

def draw_style_mixing_figure(png, Gs, w, h, src_seeds, dst_seeds, style_ranges):
    print(png)
    src_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in src_seeds)
#    dst_latents = np.stack(np.random.RandomState(seed).randn(Gs.input_shape[1]) for seed in dst_seeds)
    src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]
#    dst_dlatents = Gs.components.mapping.run(dst_latents, None) # [seed, layer, component]

    dlatents = np.load('./latent/yukiti.npy')
    dst_dlatents = np.zeros((3,18,512))
    dst_dlatents[0] = dlatents
    dst_dlatents[1] = dlatents
    dst_dlatents[2] = dlatents
    print(dst_dlatents.shape)

    src_images = Gs.components.synthesis.run(src_dlatents, randomize_noise=False, **synthesis_kwargs)
    dst_images = Gs.components.synthesis.run(dst_dlatents, randomize_noise=False, **synthesis_kwargs)

    canvas = PIL.Image.new('RGB', (w * (len(src_seeds) + 1), h * (len(dst_seeds) + 1)), 'white')
    for col, src_image in enumerate(list(src_images)):
        canvas.paste(PIL.Image.fromarray(src_image, 'RGB'), ((col + 1) * w, 0))
    for row, dst_image in enumerate(list(dst_images)):
        canvas.paste(PIL.Image.fromarray(dst_image, 'RGB'), (0, (row + 1) * h))
        row_dlatents = np.stack([dst_dlatents[row]] * len(src_seeds))
        row_dlatents[:, style_ranges[row]] = src_dlatents[:, style_ranges[row]]
        row_images = Gs.components.synthesis.run(row_dlatents, randomize_noise=False, **synthesis_kwargs)
        for col, image in enumerate(list(row_images)):
            canvas.paste(PIL.Image.fromarray(image, 'RGB'), ((col + 1) * w, (row + 1) * h))
    canvas.save(png)

def main():
    tflib.init_tf()
    os.makedirs(config.result_dir, exist_ok=True)
    draw_style_mixing_figure(os.path.join(config.result_dir, 'yukiti-style-mixing.png'), 
                             load_Gs(url_ffhq), w=1024, h=1024, src_seeds=[639,701,687,615,2268], dst_seeds=[0,0,0],
                             style_ranges=[range(0,4)]+[range(4,8)]+[range(8,18)])

if __name__ == "__main__":
    main()

⑤スタイルミキシング結果

全然上手く行ってない…。

参考までにエンコーダーのffhq_dataset/latent_representations/に入ってるドナルド・トランプの潜在変数で試すとこちらは上手く行っている。福沢諭吉の潜在変数自体が駄目なのか…。

⑥福沢諭吉を笑わせる

Play_with_latent_directions.ipynbに書いてあるように福沢諭吉の潜在変数に笑顔を示す方向に潜在変数ベクトルを足してやる。

yukiti_smile.py

import os
import pickle
import numpy as np
import PIL.Image
import dnnlib
import dnnlib.tflib as tflib
import config
import sys

from encoder.generator_model import Generator

import matplotlib.pyplot as plt

tflib.init_tf()
fpath = './karras2019stylegan-ffhq-1024x1024.pkl'
with open(fpath, mode='rb') as f:
    generator_network, discriminator_network, Gs_network  = pickle.load(f)

generator = Generator(Gs_network, batch_size=1, randomize_noise=False)

def generate_image(latent_vector):
    latent_vector = latent_vector.reshape((1, 18, 512))
    generator.set_dlatents(latent_vector)
    img_array = generator.generate_images()[0]
    img = PIL.Image.fromarray(img_array, 'RGB')
    return img.resize((256, 256))

def move_and_show(latent_vector, direction, coeffs):
    fig,ax = plt.subplots(1, len(coeffs), figsize=(15, 10), dpi=80)
    for i, coeff in enumerate(coeffs):
        new_latent_vector = latent_vector.copy()
        new_latent_vector[:8] = (latent_vector + coeff*direction)[:8]
        ax[i].imshow(generate_image(new_latent_vector))
        ax[i].set_title('Coeff: %0.1f' % coeff)
    [x.axis('off') for x in ax]
    plt.show()

yukiti = np.load('./latent/yukiti.npy')
smile_direction = np.load('ffhq_dataset/latent_directions/smile.npy')

move_and_show(yukiti, smile_direction, [-1, 0, 1])

これもノイズがひどくてあまりよろしくない。（ちょっと口が開いてるように変化しているが）

参考までに添付されているドナルド・トランプの潜在変数なら上手く行く。

⑦駄目だった理由考察

・用意した画像が256x256で小さいから？
・潜在変数の最適化回数が少ないから？
・白黒写真なのが駄目だったのか？
・西洋人ではなく日本人なのが駄目だったのか？
などが考えられます。

追記：重要だったアライメント

コメントに指摘されたアライメント（目と口の位置を揃える)をやってみたところ結果がずいぶん改善されました。自分は単なる画像リサイズだと思って処理飛ばしていたんですよね…。
dlibが何だかインストールに詰まったのでpython2.7の仮想環境作ってpip install cmake、pip install dlibしました。そして新規にal_imgフォルダを作成し、python align_images.py img/ al_img/ で実行。

結果、al_img/フォルダに以下のようなアライメント画像が生成されました。
また、この画像の潜在変数をiteration=20000で求めました。実行時間1時間ほどでloss0.09になりました。
　
（左：アライメント画像、右：潜在変数へ変換後GAN生成画像）

この潜在変数を使った場合のスタイルミキシング結果は以下です。
追記前に比べれば潜在変数の質が良くなりましたね。でも、白黒写真の合成は向いてないかもしれません。

笑顔の方向ベクトルを加えた結果は以下です。

結論としてはアライメントは超重要です。