StyleGAN2学習済みモデルを使ったmorphing、latent spaceの探求

学習データはまだまだ収集途中だし学習もまだ完了とは言えない状態なのだけど、なんとなくそれっぽい顔画像は生成できるくらいまでは出来てきているので、それを使った実験をしてみている。

学習した生成モデルを使って morphingを試してみてる #StyleGAN2 pic.twitter.com/NxlaOgp1RX
— すぎゃーん💯 (@sugyan) February 7, 2020

なるほどねー pic.twitter.com/7Qmf61G14t
— すぎゃーん💯 (@sugyan) February 9, 2020

主にこの記事を参考にしています。

1. `latents_in` の線形移動

まず最初に試したのは、通常の generator network を使った場合の latents_inを使ったもの。

generator は、(?, 512) の任意の入力を受けて画像を生成するようになっている。これが所謂潜在空間 (latent space) というやつでここの値をうまく選ぶことで所望の出力を得られるようになったりする、というわけだ。

例えば適当な乱数で (2, 512) を作ってそれを入力すると、2つの異なる画像が出力される。

その2つの (1, 512) の乱数ベクトルの、中間の値を入力として使えば生成される画像も2つの画像の中間のものになるだろう、という考え方。

import numpy as np
import tensorflow as tf
from PIL import Image

model_path = '...'
model = tf.saved_model.load(model_path)
generator = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

rnd = np.random.RandomState(0)
z = rnd.randn(2, 512)

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

for i, latents in enumerate(inputs):
    images = generator(latents=tf.constant([latents], tf.float32))['images']
    Image.fromarray(images.numpy()[0]).save(f'out_{i:02d}.png')

z[0] から z[1] までを 0.1 刻みで線形に変化させてそれぞれ画像を生成した結果。

f:id:sugyan:20200209192014p:plain

右端と左端は(一応)違う顔だが、その間で徐々に変化していっているのが分かる、と思う。 GIFアニメーションにするとこんな感じ。

f:id:sugyan:20200209192737g:plain

2. mapping networkの出力 `dlatents_in` を使う

StyleGAN の generator は、"mapping network" と "synthesis network" の2つの network によって作られている。 generator に入力された latents_in の (?, 512) の値は mapping network で (?, 14, 512) といったもう少し次元の大きい変数に変換される。(14 というのは 256x256 のときの数値で upsample の layer 数が増えるとまた 16 や 18 に変化したりする、のかな。) で、この Disentangled latents と呼ばれる出力が、後段の synthesis network への入力に使われ、実際の画像の合成が行われる、ということになる。 dlatents_in と呼ばれるこの値が画像生成のための入力としてはより直接的なものになるのかもしれない。

ザックリした理解では、 synthesis network は実に広範囲な生成の能力を持っているがより所望の(学習データに近い)主力を得るための dlatents_in を作り出すための前段として mapping network が作用している、という感じか。実際にはその出力からさらに truncation_psi trick といったより質の良い出力を得ることが出来る値にするための工夫がされているようだ。

ともかく、この 2つの network を分けて考えると、2つの画像間を補完する場合も mapping network への入力 latents_in を線形に変化させるよりも synthesis network への入力 dlatents_in を変化させた方が自然なものになりそう。

generator も2つに分けて各 network の入出力を扱えるようにした。 StyleGAN2 では return_dlatents = True と option を指定することで dlatents_in の tensor も得ることが出来る。

この dlatents_in をまず先に計算し、先程と同様にその値を 0.1 刻みで線形に変化させるようにする。

import numpy as np
import tensorflow as tf
from PIL import Image

model_path = '...'
model = tf.saved_model.load(model_path)
mapping = model.signatures['mapping']

rnd = np.random.RandomState(0)
z = rnd.randn(2, 512)
z = mapping(latents=tf.constant(z, tf.float32))['dlatents'].numpy()

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

synthesis = model.signatures['synthesis']
for i, latents in enumerate(inputs):
    images = synthesis(dlatents=tf.constant([latents], tf.float32))['images']
    Image.fromarray(images.numpy()[0]).save(f'out_{i:02d}.png')

f:id:sugyan:20200209200947p:plain

先述の latents_in の変化だと途中で顎のあたりに手のようなものが現れたりガチャガチャと忙しい変化になってしまっていたのに対し、今度はシームレスに2画像間を推移するようになった。

f:id:sugyan:20200209201028g:plain

latent space を探る

dlatents_in は本当に広い空間で、mapping network を通さずに本当に random な値を使って synthesis network に入力するとまったく汚い出力になってしまう。

f:id:sugyan:20200209201952p:plain

逆にここを上手く最適化してやることで学習データにまったく存在していない画像も生成することが出来る可能性があるようだ。これについてはまたこれから実験していきたい。

dlatents_in の値を mapping network の出力に絞ることでそれなりに学習データに近い出力が得られるようにはなるが、学習データの偏りや学習不足などのせいかどうしても印象が似たような顔ばかりになりやすい。

とはいえ、これも mapping network の出力として得られる dlatents_in が大きく異なっていればそれなりには異なる出力画像になるのではないか？

1000件ほどの random な latents_in を mapping network に入力し、得られた各 dlatents_in の間の距離を計算してみる。これが最も遠い組み合わせを選んで使ったら、出力画像も大きく異なるものになるだろうか？

rnd = np.random.RandomState(0)
z = rnd.randn(1000, 512)
z = mapping(latents=tf.constant(z, tf.float32))['dlatents'].numpy()

distances = []
for i in range(z.shape[0]):
    for j in range(i + 1, z.shape[0]):
        distances.append([np.linalg.norm(z[i] - z[j]), (i, j)])

_, (i, j) = sorted(distances, reverse=False)[0]
z = [z[i], z[j]]

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)