本記事では、Thin-Plate Spline Motion Model for Image Animationと呼ばれる機械学習手法を用いて、顔写真を動画の顔の動きに合わせて動かす方法をご紹介します。

出典: yoyo-nb/Thin-Plate-Spline-Motion-Model

Thin-Plate Spline Motion Model for Image Animation

概要

Image animationタスクでは、ドライビングビデオに従って、ソース画像にモーションを転送することで、画像を動かします。

従来技術では、教師なしで任意のドライビングビデオのモーションを画像に転送しようと試みていますが、動画と画像の間に大きなポーズのギャップがある場合、転送は依然として困難です。

Thin-Plate Spline Motion Modelは、この問題を克服するための新しいEnd to Endの教師なしモーション転送フレームワークです。

本手法では、まずthin-plate spline motion estimationで柔軟なオプティカルフローを形成します。これにより、ソース画像の特徴マップがドライビングビデオのフレームであるドライビング画像の特徴領域にワープします。

次に、欠落した領域を自然に復元するため多重解像度のオクルージョンマスクを用いて、効果的に特徴の融合を実現します。

最後にadditional auxiliary loss functionsは、ネットワークモジュールのそれぞれが分業してトレーニングできるように設計されており、高品質な画像生成を促進します。
本手法は、話している顔、人の動き、ピクセルアニメーションなど、様々なオブジェクトのアニメーション化を実現しています。

出典: Thin-Plate Spline Motion Model for Image Animation

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

本記事では上記手法を用いて、顔画像とドライビング動画を用いてImage animationを行います。

デモ(Colaboratory)

それでは、実際に動かしながらImage animationを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

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

初めに、論文発表元のGithubからソースコードを取得します

%cd /content

!git clone https://github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model.git

次にライブラリをインストールします。
PULSEは、画像から位置調整しながら顔部分を切り抜くalign_face.pyを使用するためcloneしています。

%cd /content/

!pip install face_alignment > /dev/null
# face alignment用にclone
!git clone https://github.com/adamian98/pulse.git

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

%cd /content/Thin-Plate-Spline-Motion-Model

import torch

import imageio
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from skimage.transform import resize
from IPython.display import HTML
import warnings
import os

from demo import load_checkpoints
from demo import make_animation
from skimage import img_as_ubyte

from google.colab import files
from moviepy.editor import *

warnings.filterwarnings("ignore")

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

テストデータのセットアップ

ドライビング動画と、ソース画像をセットアップしていきます。
任意の動画と画像をGoogle Colabにアップロードしていきます。

#@markdown 動画の切り抜き範囲(秒)を指定してください。\
#@markdown 30秒以上の場合OOM発生の可能性が高いため注意
start_sec =  3#@param {type:"integer"}
end_sec =  7#@param {type:"integer"}

(start_pt, end_pt) = (start_sec, end_sec)

%cd /content/Thin-Plate-Spline-Motion-Model

!rm -rf test_data
!mkdir test_data
%cd test_data
!mkdir image aligned_image video frames aligned_video

%cd video
print("upload video...")
video = files.upload()
video = list(video.keys())
video_file = video[0]
# 指定区間切り抜き
with VideoFileClip(video_file) as video:
    subclip = video.subclip(start_pt, end_pt)
    subclip.write_videofile("./video.mp4")
# frameに分割
!ffmpeg -i video.mp4 ../frames/%02d.png

%cd ../image
print("upload image...")
image = files.upload()
image = list(image.keys())
image_file = image[0]

本記事では、以下の画像と動画を使用させていただきます。

使用画像

使用動画

次に、画像と動画から顔部分を切り出します。

%cd /content/Thin-Plate-Spline-Motion-Model/test_data


!python /content/pulse/align_face.py \
  -input_dir /content/Thin-Plate-Spline-Motion-Model/test_data/image \
  -output_dir /content/Thin-Plate-Spline-Motion-Model/test_data/aligned_image \
  -output_size 256 \
  -seed 1234
  
%cd /content/Thin-Plate-Spline-Motion-Model

# tedの場合はoutput_size 384
!python /content/pulse/align_face.py \
  -input_dir /content/Thin-Plate-Spline-Motion-Model/test_data/frames \
  -output_dir /content/Thin-Plate-Spline-Motion-Model/test_data/aligned_video \
  -output_size 256 \
  -seed 1234
  
!ffmpeg -i /content/Thin-Plate-Spline-Motion-Model/test_data/aligned_video/%02d_0.png -c:v libx264 -vf "fps=25,format=yuv420p" /content/Thin-Plate-Spline-Motion-Model/test_data/aligned_video/aligned.mp4

モデルのセットアップ

続いて、モデルのセットアップを行います。
今回はalignされた顔写真を使用するため"vox"を選択しています。

%cd /content/Thin-Plate-Spline-Motion-Model

# @markdown モデル選択
dataset_name = 'vox' #@param ["vox", "taichi", "ted", "mgif"]
# @markdown 入力画像
source_image_path = '/content/Thin-Plate-Spline-Motion-Model/test_data/aligned_image/nissinIMGL0823_TP_V_0.png' #@param {type:"string"}
# @markdown 入力動画
driving_video_path = '/content/Thin-Plate-Spline-Motion-Model/test_data/aligned_video/aligned.mp4' #@param {type:"string"}
# @markdown 出力先
output_video_path = './generated.mp4' #@param {type:"string"}
# @markdown predict mode
predict_mode = 'relative' #@param ['standard', 'relative', 'avd']
# "relative"の際にTrueにすると出力結果の品質が向上
find_best_frame = False  #@param {type:"boolean"}

%cd /content/Thin-Plate-Spline-Motion-Model

# edit the config
device = torch.device('cuda:0')

%cd /content/Thin-Plate-Spline-Motion-Model

!mkdir checkpoints

# dataset_name = 'vox' #@param ["vox", "taichi", "ted", "mgif"]

if dataset_name == 'vox':
  !wget -c https://cloud.tsinghua.edu.cn/f/da8d61d012014b12a9e4/?dl=1 -O checkpoints/vox.pth.tar
  config_path = 'config/vox-256.yaml'
  checkpoint_path = 'checkpoints/vox.pth.tar'
  pixel = 256
if dataset_name == 'taichi':
  !wget -c https://cloud.tsinghua.edu.cn/f/9ec01fa4aaef423c8c02/?dl=1 -O checkpoints/taichi.pth.tar
  config_path = 'config/taichi-256.yaml'
  checkpoint_path = 'checkpoints/taichi.pth.tar'
  pixel = 256
if dataset_name == 'ted':
  !wget -c https://cloud.tsinghua.edu.cn/f/483ef53650b14ac7ae70/?dl=1 -O checkpoints/ted.pth.tar
  config_path = 'config/ted-384.yaml'
  checkpoint_path = 'checkpoints/ted.pth.tar'
  pixel = 384
if dataset_name == 'mgif':
  !wget -c https://cloud.tsinghua.edu.cn/f/cd411b334a2e49cdb1e2/?dl=1 -O checkpoints/mgif.pth.tar
  config_path = 'config/mgif-256.yaml'
  checkpoint_path = 'checkpoints/mgif.pth.tar'
  pixel = 256

Image Animation

最後にImage Animationを実行します。

inpainting, kp_detector, dense_motion_network, avd_network = load_checkpoints(config_path = config_path, checkpoint_path = checkpoint_path, device = device)

if predict_mode=='relative' and find_best_frame:
    from demo import find_best_frame as _find
    i = _find(source_image, driving_video, device.type=='cpu')
    print ("Best frame: " + str(i))
    driving_forward = driving_video[i:]
    driving_backward = driving_video[:(i+1)][::-1]
    predictions_forward = make_animation(source_image, driving_forward, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = predict_mode)
    predictions_backward = make_animation(source_image, driving_backward, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = predict_mode)
    predictions = predictions_backward[::-1] + predictions_forward[1:]
else:
    predictions = make_animation(source_image, driving_video, inpainting, kp_detector, dense_motion_network, avd_network, device = device, mode = predict_mode)

#save resulting video
imageio.mimsave(output_video_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)

HTML(display(source_image, driving_video, predictions).to_html5_video())

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

口元はうまく転送できているもの目元や顔の動きはやや不自然な結果となりました。
別の組み合わせを試してみた結果は以下の通りです。

こちらは、非常に自然な仕上がりになりました。

まとめ

本記事では、Thin-Plate Spline Motion Modelを用いた顔画像のImage Animationをご紹介しました。
顔画像の他にも体全体の画像を使ったImage Animationの学習済みモデルなどをあるので色々試してみてください。

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

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

参考文献

1.  論文 - Thin-Plate Spline Motion Model for Image Animation

2. GitHub - yoyo-nb/Thin-Plate-Spline-Motion-Model