[StyleGAN2-ADA] 機械学習で顔画像のモーフィング動画を作成 [Python]

2022年7月15日金曜日

Artificial Intelligence

本記事では、StyleGAN2-ADAを用いて2つの顔画像の中間画像を生成しモーフィング動画を作成する方法をご紹介します。

result

StyleGAN2-ADA

概要

StyleGAN2は、StyleGANの水滴状の歪みが発生するという問題を解消した画像生成モデルです。

StyleGAN issue
出典: Analyzing and Improving the Image Quality of StyleGAN

ここからさらに、少ないデータ数でも効率的に学習できるようdata augumentationを取り入れたモデルがStyleGAN2-ADAとなります。

StyleGAN2-ada
出典: Training Generative Adversarial Networks with Limited Data

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、2つの顔画像の中間画像を連続的に生成してみます。

デモ(Colaboratory)

それでは、実際に動かしながら2つの顔画像の中間画像を連続的に生成します。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。
Open In Colab

また、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

初めにGithubからソースコードを取得します。

%cd /content

!git clone https://github.com/NVlabs/stylegan2-ada-pytorch.git
# for align face
!git clone https://github.com/adamian98/pulse.git

次にライブラリをインストールします。

!pip install moviepy

# ninja
!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip > /dev/null
!sudo unzip ninja-linux.zip -d /usr/local/bin/ > /dev/null
!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force > /dev/null

最後にライブラリをインポートします。

%cd /content/stylegan2-ada-pytorch

import cv2
import numpy as np
from PIL import Image
import dlib
from matplotlib import pyplot as plt
import torch
import dnnlib
import legacy
import imageio
from tqdm.notebook import tqdm

from moviepy.video.fx.resize import resize
from moviepy.editor import VideoFileClip

以上で環境セットアップは完了です。

テスト画像のセットアップ

ここでは、2つの顔画像をGoogle Colaboratory上にダウンロードしていきます。
本記事ではぱくたそ様の画像を使用させていただきます。

%cd /content/stylegan2-ada-pytorch
!rm -rf inputs
!mkdir -p inputs/img inputs/align

!wget -c https://www.pakutaso.com/shared/img/thumb/model10211041_TP_V.jpg \
      -O ./inputs/img/src.jpg

!wget -c https://www.pakutaso.com/shared/img/thumb/kuchikomi725_TP_V.jpg \
      -O ./inputs/img/dst.jpg

ダウンロードした画像からそれぞれ顔画像を切り抜き位置合わせします。

%cd /content/pulse
!python align_face.py \
  -input_dir /content/stylegan2-ada-pytorch/inputs/img \
  -output_dir /content/stylegan2-ada-pytorch/inputs/align \
  -output_size 1024 \
  -seed 12 \

Invert

位置合わせした2枚の顔画像をそれぞれモデルの潜在空間に投影します。

%cd /content/stylegan2-ada-pytorch

!python projector.py \
  --save-video 0 \
  --num-steps 1000 \
  --outdir=outputs/src \
  --target=inputs/align/src_0.png \
  --network={NETWORK}
  
%cd /content/stylegan2-ada-pytorch

!python projector.py \
  --save-video 0 \
  --num-steps 1000 \
  --outdir=outputs/dst \
  --target=inputs/align/dst_0.png \
  --network={NETWORK}

モーフィング

それでは、Src画像からDst画像への徐々に画像を変化させていきモーフィング動画を生成します。

%cd /content/stylegan2-ada-pytorch

lvec1 = np.load('./outputs/src/projected_w.npz')['w']
lvec2 = np.load('./outputs/dst/projected_w.npz')['w']

network_pkl = "https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl"
device = torch.device('cuda')
with dnnlib.util.open_url(network_pkl) as fp:
    G = legacy.load_network_pkl(fp)['G_ema'].requires_grad_(False).to(device) # type: ignore

diff = lvec2 - lvec1
step = diff / STEPS
current = lvec1.copy()
target_uint8 = np.array([1024,1024,3], dtype=np.uint8)

video = imageio.get_writer('./movie.mp4', mode='I', fps=FPS, codec='libx264', bitrate='16M')

for j in tqdm(range(STEPS)):
  z = torch.from_numpy(current).to(device)
  synth_image = G.synthesis(z, noise_mode='const')
  synth_image = (synth_image + 1) * (255/2)
  synth_image = synth_image.permute(0, 2, 3, 1).clamp(0, 255).to(torch.uint8)[0].cpu().numpy()

  repeat = FREEZE_STEPS if j==0 or j==(STEPS-1) else 1
   
  for i in range(repeat):
    video.append_data(synth_image)
  current = current + step


video.close()

出力結果は以下の通りです。

result

src画像からdst画像へと徐々に変化させた画像を生成し、最後にそれらの画像をフレームとして動画に変換することでモーフィング動画を作成しています。
気づいたらdst画像へと変わっていると感じるほど自然なつながりです。

まとめ

本記事では、本記事では、StyleGAN2-ADAを用いて2つの顔画像の中間画像を生成しモーフィング動画を作成する方法をご紹介しました。
潜在空間への投影時に元画像と離れてしまう場合がありますが、投影後は非常に自然な中間画像が生成されています。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。


また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Analyzing and Improving the Image Quality of StyleGAN

2. GitHub - NVlabs/stylegan2-ada-pytorch

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology