StyleGAN2で顔ランドマーク座標を指定して画像生成【後編】

前回はキーポイントを指定した際に他の顔属性も変わってしまっていたので，顔属性を保持しながらキーポイントを変化させるために属性推定器を学習させます．

データ

顔属性

やはりラベル付きのデータが欲しいので，CelebAを使います．CelebAのラベルは40種類あります．全ラベルの意味は以下を参考に． cedro3.com

論文ではAzure Face APIを使ってラベルを増やして48種類の属性を推定しています．が，今回はCelebAでラベル付けされた40種類の属性のうち，①顔ランドマーク座標の指定で代替できそうな属性(目や口の開き方など)，②髪色に関係する属性　を除いた28種類を用いることにします．

髪色

CelebAでは髪色に関するラベルは「Black Hair」「Blond Hair」「Brown Hair」の3種類しかありません．そこで柔軟にかつ多様に髪色を変更できるように，RGBの値で髪色を指定することを考えます．幸いなことに，SGF(今回使用している手法)は原理的には属性の指定方法に制約がありません．

具体的には顔領域をパース出来るモデルを使って髪部分のみを抽出し，その平均色をSGFの入力に加えることとします．パーシングには以下のモデルを使用します． github.com

SGFの models/classifier.py に組み込みます．

out = self.bisenet(img_tensor)[0]           # パース結果を出力
out = out.argmax(1)                              # チャネル方向にargmaxしてクラスを推定
hair = img_tensor[:, :, out[0]==17].mean(2)[0]     # 17(髪)クラスだと推定した部分の平均を取る

モデル

論文と同じようにVGGFace2で訓練したSEResNet50を使います．以下の学習済みモデルを使用します．これを使って適当にfine-tuningします． github.com

SGF実装

前回までの実装は毎iterationの更新ごとに画像を生成して目的性質との差を計算し，その誤差が最小になったところでストップします．が，これは速度の面からみると画像生成の部分がボトルネックになります．そこで誤差が最小になるところでストップせず，決めたステップ数まで更新するようにします．精度は下がりますが，速度向上は見込めます．

translate関数

z += d_z
# 画像生成
image, _ = G([z], input_is_latent=True, randomize_noise=False)

# 性質を推定
c = C(image).unsqueeze(0).to(device)

# 2乗和誤差
c_diff_new = (c-c1).pow(2).mean()
            
print(c_diff_new)

# 今回の誤差(c_diff_new)が前回の誤差(c_diff)よりも大きい場合は打ち止め
if c_diff > 0 and c_diff < c_diff_new:
    break

c_diff = c_diff_new

translate_faster関数

z += d_z
# cは一様に変化すると仮定
c += d_c

結果

キーポイント

まずは顔キーポイント座標を指定して生成してみます．左列がソース画像，上行がターゲット座標です．前回と比べればかなりアイデンティティは保存されているように見えます．そもそもキーポイント座標を別の画像から持ってきているため，輪郭などの不整合が起こっているというのもあってアイデンティティが微妙に変化しているということも考えられます． f:id:kym384:20210809163559p:plain