【論文読解】First Order Motion Model for Image Animation

# 概要

1 枚の静止画を、参照動画を元にリアルタイムで動画化する手法を提案した論文"First Order Motion Model for Image Animation"について解説します。

なにができるようになるのかは、公式実装のリポジトリを見ていただければ一発で理解できます。要するに、参照動画中の人物や物体の姿勢を静止画に適用し、動画化したいわけです。

本手法を元に構築されたアプリケーションとして、例えば以下のリポジトリで公開されている avatarify があります。Web カメラで撮影された自分の顔の動きに合わせて、モナリザやアインシュタインの顔がリアルタイムで動いてくれます。

https://github.com/alievk/avatarify

この記事では、本手法の基本的な枠組みを示した後、構成要素を少し深堀りして説明していきます。

# 書誌情報

- Siarohin, Aliaksandr, et al. "First Order Motion Model for Image Animation." Advances in Neural Information Processing Systems. 2019.
- https://papers.nips.cc/paper/8935-first-order-motion-model-for-image-animation
- [公式実装リポジトリ](https://github.com/AliaksandrSiarohin/first-order-model)

# 基本的な枠組み
以下の図は本手法を 1 枚で表したものです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/3baf8d1f-257e-4526-1f1b-9399b050774e.png)

本手法は、以下のような要素から構成されています。

- KeyPoint Detector: Source 画像$\mathbf{S}$と Driving Frame$\mathbf{D}$を入力すると、それらのキーポイント位置$\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{R}}(p_k), \mathcal{T}\_{\mathbf{D} \leftarrow \mathbf{R}}(p_k)$及び、ヤコビアン$J_k$を推定します。
- Dense Motion: キーポイント位置とヤコビアン、ソース画像$\mathbf{S}$を元に、マッピング関数$\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{D}}$とオクルージョンマスク$\hat{\mathcal{O}}\_{\mathbf{S} \leftarrow \mathbf{D}}$を推定します。
- Generation Module: 最後に、ソース画像$\mathbf{S}$と変換マップ$\hat{\mathcal{T}}\_{\mathbf{S} \leftarrow \mathbf{D}}$, オクルージョンマスク$\hat{\mathcal{O}}\_{\mathbf{S} \leftarrow \mathbf{D}}$を元に、最終的な出力結果を得ます。

これらの処理の第一目標は、Source 画像$\mathbf{S}$を Driving Frame$\mathbf{D}$の姿勢に変換する 2 次元マッピング関数$\hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}: \mathbb{R}^2 \rightarrow \mathbb{R}^2$を得ることです。
このマッピングは、出力画像の座標を入力すると、入力画像の座標を出力します。つまり、Driving Frame$\mathbf{D}$のような姿勢になることを期待されている出力画像の各画素を、ソース画像$\mathbf{S}$のどの画素から作ればいいのかを表すマップです。

$\hat{\mathcal{T}}\_{\mathbf{S} \leftarrow \mathbf{D}}$を得るために、KeyPoint Detector と Dense MOtion というモジュールが使われています。直感的な理解としては、KeyPoint DetectorによってSource 画像$\mathbf{S}$や Driving Frame$\mathbf{D}$に含まれるキーポイントを検出します。このキーポイントは$\mathbf{S, D}$間における疎なマッピングなので、これを参考にして密なマッピング関数$\hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}$を得る、というのが Dense Motion の役割です。

基本的には、$\hat{\mathcal{T}}\_{\mathbf{S} \leftarrow \mathbf{D}}$さえ得られれば、ソース画像に姿勢変化を加えた出力画像を得ることができます。しかし、より良くするためにオクルージョンを考慮して生成しましょう、というのがオクルージョンマスク$\hat{\mathcal{O}}_{\mathbf{S} \leftarrow \mathbf{D}}$と Generation Module の役割です。

訓練中は、Source 画像と Driving Frame は、同一の動画の中の任意のフレームから選択され、出力画像のターゲットは Driving Frame そのものとなります。なお、論文中に記述が無いように見えますが、実装を確認してみるとGANの枠組みも使用しているようですが、本稿では触れません。

推論時は、Source画像として任意の画像を指定し、Driving Frameとして任意の動画を入力することができます。これにより、Driving Frameのような姿勢へと変換されたSource画像が得られます。

# KeyPoint Detector
KeyPoint Detector の中身をもう少し詳しめに見ていきましょう。

まず、$\mathbf{S}$と$\mathbf{D}$は個別に処理され、それぞれ以下のような出力を得ます。これらの出力は独立に得ることができます。直感的には、1 つ目の出力はキーポイントの位置、2 つ目の出力はキーポイントにおける勾配を表現しています。

- $\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{R}}(p_k), \left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right)$ 
- $\mathcal{T}\_{\mathbf{D} \leftarrow \mathbf{R}}(p_k), \left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{D} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right)$

その後、以下のように計算されるヤコビアン$J_k$が得られます。

```math
J_k \equiv \left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)^{-1}
```


いずれにも下付きの$k$がついていますが、これはキーポイントのインデックスを表しています。つまり、全体としては、キーポイント数ぶんの上記出力が得られます。なお補足ですが、本手法においてはキーポイントは自動的に訓練されるもので、アノテーションは一切不要です。

上記でいきなり出てきた$\mathbf{R}$は Reference 画像を意味します。また、$p_k$は$\mathbf{R}$におけるキーポイントの座標を表します。$\mathbf{R}$はあくまでも概念的な基準画像とでも呼ぶべきもので、実体としては存在しておらず可視化はできません。$\mathbf{S, D}$に含まれるキーポイントを直接座標として得るのではなく、リファレンス画像$\mathbf{R}$を基準にした変換マッピング$\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{R}}(p_k)$, $\mathcal{T}\_{\mathbf{D} \leftarrow \mathbf{R}}(p_k)$によって表現している、というわけです。このReference画像を設けることによって、$\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{R}}(p_k)$と$\mathcal{T}\_{\mathbf{D} \leftarrow \mathbf{R}}(p_k)$を独立して計算できるという利点があります。

さて、いろいろな数式展開をすっ飛ばして結論だけ言えば、上記の出力を用いることでキーポイント$z=z_k$の周りでの$\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{D}}(z)$は以下のように近似できます。

```math
\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z) \approx \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}\left(p_{k}\right)+J_{k}\left(z-\mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}\left(p_{k}\right)\right)
```

以降はこの導出ですが、めんどくさければ読み飛ばしてください。

## 導出
$\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)$の$z=z_k$周りの 1 次までのテイラー展開は、以下のように計算できます。

```math
\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)=\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}\left(z_{k}\right)+\left(\left.\frac{d}{d z} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)\right|_{z=z_{k}}\right)\left(z-z_{k}\right)+o\left(\left\|z-z_{k}\right\|^2\right)
```

ここで$o$は無視し、$z_k = \mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}(p_k)$を代入すると、以下のような近似が得られます。

```math
\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)\approx \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}\left(z_{k}\right)+\left(\left.\frac{d}{d z} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)\right|_{z=z_{k}}\right)\left(z-\mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}(p_k)\right)
```

ここで$J_k = \left(\left.\frac{d}{d z} \mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{D}}(z)\right|\_{z=z_{k}}\right)$が得られれば、目的の式になるということがわかります。
このように 1 次の近似でマッピングの関数を得るので、本手法の名称に"First Order"が出てくるわけです。

もう少し$\left(\left.\frac{d}{d z} \mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{D}}(z)\right|\_{z=z_{k}}\right)$を詳しく見ていきます。連鎖律から以下のように分解できます。

```math
\left(\left.\frac{d}{d z} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)\right|_{z=z_{k}}\right)=\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}(p)\right|_{p=\mathcal{T}_{\mathbf{R} \leftarrow \mathbf{D}}\left(z_{k}\right)}\right)\left(\left.\frac{d}{d z} \mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}^{-1}(z)\right|_{z=z_{k}}\right)
```

逆関数$\mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}^{-1}(z)$が出てきてなかなか怖いのですが、必要なのは、そのヤコビアンだけです。逆関数のヤコビアンは、関数のヤコビアンの逆数であるという性質を利用すると、以下のように置き換えることができます。

```math
\left(\left.\frac{d}{d z} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{D}}(z)\right|_{z=z_{k}}\right)=\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)^{-1}
```

というわけで、$J_k$の定義にたどり着きました。

# 複数のマッピングの統合

キーポイント数ぶんのマッピング関数$\mathcal{T}\_{\mathbf{S} \leftarrow \mathbf{R}}\left(p_{k}\right)+J_{k}\left(z-\mathcal{T}\_{\mathbf{D} \leftarrow \mathbf{R}}\left(p_{k}\right)\right)$が得られたわけですが、それらを統合して 1 つのマッピング関数$\hat{\mathcal{T}}\_{\mathbf{S} \leftarrow \mathbf{D}}$にする必要があります。マッピング関数は、キーポイント位置の付近では正確ですが、そこから離れた領域では不正確になるため、位置によって重み付けしていいとこ取りをして統合します。

以下の 2 つのデータを作って合成します。

- キーポイントごとのマッピング関数を元に作られる合成画像$\mathbf{S}^{0}, \mathbf{S}^{1} \ldots \mathbf{S}^{K}$$\mathbf{S}^0$は背景画像としてSource 画像そのものが使用されます。
- キーポイント位置ほど大きくなるようなヒートマップ$\mathbf{H}_{1} \ldots \mathbf{H}\_{K}$。下式で作られる。

```math
\mathbf{H}_{k}(z)=\exp \left( - \frac{\left(\mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}\left(p_{k}\right)-z\right)^{2}}{2\sigma}\right)-\exp \left( - \frac{\left(\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}\left(p_{k}\right)-z\right)^{2}}{2\sigma}\right)
```

以上の合成画像$\mathbf{S}^{0}, \mathbf{S}^{1} \ldots \mathbf{S}^{K}$とヒートマップ$\mathbf{H}\_{1} \ldots \mathbf{H}\_{K}$を結合して、U-Net 構造を持つ Dense Motion に入力すると、出力として$K+1$チャネルのマスク$\mathbf{M}\_{0}, \mathbf{M}\_{1} \ldots \mathbf{M}\_{K}$が得られます。このマスクを使ってもともとのキーポイントごとのマッピング関数に重み付けを行い、最終的に以下のように統合されたマッピング関数$\hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}$が得られます。

```math
\hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}(z)=\mathbf{M}_{0} z+\sum_{k=1}^{K} \mathbf{M}_{k}\left(\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}\left(p_{k}\right)+J_{k}\left(z-\mathcal{T}_{\mathbf{D} \leftarrow \mathbf{R}}\left(p_{k}\right)\right)\right)
```


# オクルージョンの考慮

$\hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}$は Source 画像$\mathbf{S}$に直接適用されるのではなく、2 層の Conv 層からなる Encoder を通して得られた特徴マップ$\boldsymbol{\xi}$に対して適用されます。

Dense Motion は$\hat{\mathcal{T}}\_{\mathbf{S} \leftarrow \mathbf{D}}$以外にオクルージョンマップ$\hat{\mathcal{O}}\_{\mathbf{S} \leftarrow \mathbf{D}}$も出力します。この$\hat{\mathcal{O}}\_{\mathbf{S} \leftarrow \mathbf{D}}$と$\hat{\mathcal{T}}\_{\mathbf{S} \leftarrow \mathbf{D}}$をあわせて、特徴マップは以下のように変換されます。$f_{w}$は、特徴マップ$\boldsymbol{\xi}$にマッピング関数$\hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}$を適用することを表します。

```math
\boldsymbol{\xi}^{\prime}=\hat{\mathcal{O}}_{\mathbf{S} \leftarrow \mathbf{D}} \odot f_{w}\left(\boldsymbol{\xi}, \hat{\mathcal{T}}_{\mathbf{S} \leftarrow \mathbf{D}}\right)
```

オクルージョンマップ$\hat{\mathcal{O}}_{\mathbf{S} \leftarrow \mathbf{D}}$によってマッピング関数上で生じているオクルージョン領域を小さく重み付けし、新たな特徴マップ$\boldsymbol{\xi}^{\prime}$が得られています。

この特徴マップ$\boldsymbol{\xi}^{\prime}$を元に、2 層の Conv 層からなる Decoder を通して最終的な出力画像を得ます。

# 損失関数
損失関数は、出力画像が適切であるかを表す再構成誤差と、マッピング関数が極端な形になることを避ける正則化の損失から構成されます。

## 再構成誤差

再構成損失は、Driving Frame$\mathbf{D}$とそれに似せて作られた出力画像$\hat{\mathbf{D}}$との L1 誤差によって得られます。

$$
L_{r e c}(\hat{\mathbf{D}}, \mathbf{D})=\sum_{i=1}^{I}\left|N_{i}(\hat{\mathbf{D}})-N_{i}(\mathbf{D})\right|
$$

$N_i$は VGG-19 によって抽出された抽出特徴マップの$i$番目を表します。この損失は、様々に入力画像をリサイズして計算されます。$I=5$で、入力画像サイズは 256×256, 128×128, 64×64, 32×32 の 4 種類あるので、全部で 20 の損失が計算されます。

## マッピング関数の正則化

キーポイントの訓練が不安定になるのを避けるため、Keypoint Detector を通じて得られるマッピング関数が、単純な変換（thin-plate splines や affine 変換）に近づくような拘束条件を与えます。

適当にサンプルされた画像$\mathbf{X}$と、ランダムに作られた既知の単純な変換$\mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{Y}}$を使って変換された画像$\mathbf{Y}$を用意し、これら2枚の画像からKeypoint Detectorによって得られる$\mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{R}}\left(p_{k}\right), \mathcal{T}\_{\mathbf{Y} \leftarrow \mathbf{R}}\left(p_{k}\right), \left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right), \left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{Y} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right)$といった要素間で成立するはずの条件をつかって正則化を行います。

具体的には、以下の2つの等式が成り立ちます。これらの条件から、実際にKeypoint Detectorから得られた要素を代入して、左辺と右辺の差異をL1損失として算出します。

- $\mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{R}}\left(p_{k}\right) \equiv \mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{Y}} \circ \mathcal{T}\_{\mathbf{Y} \leftarrow \mathbf{R}}\left(p_{k}\right)$ 
- $\left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right) \equiv\left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{Y}}(p)\right|\_{p=\mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{R}}\left(p_{k}\right)}\right)\left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{Y} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right)$

なお後者の条件を使って損失を小さくしようとすると、2 つのヤコビアン$\left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{X} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right), \left(\left.\frac{d}{d p} \mathcal{T}\_{\mathbf{Y} \leftarrow \mathbf{R}}(p)\right|\_{p=p_{k}}\right)$の大きさが 0 に近づいてしまうという不安定性を抱えているので、以下のような拘束条件に変形してあげます。$\mathbb{1}$は単位行列です。

```math
\mathbb{1} \equiv\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{X} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)^{-1}\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{X} \leftarrow \mathbf{Y}}(p)\right|_{p=\mathcal{T}_{\mathbf{Y} \leftarrow \mathbf{R}}\left(p_{k}\right)}\right)\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{Y} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)
```


# 推論時の処理

実際に本手法を使って推論をするときに知りたくなるのは、出力画像$\mathbf{S}\_t$の各座標の値を$\mathbf{S}\_1$のどの座標から取得するのかを決めるマッピング関数$\mathcal{T}\_{\mathbf{S}\_{1} \leftarrow \mathbf{S}\_{t}}(z)$です。

しかし、この変換を$\mathbf{S}\_1, \mathbf{D}\_t$とを使って直接推論すると、グローバルな物体のジオメトリが崩壊する可能性があるそうです。そこで、$\mathbf{S}\_{1}$の姿勢にできるだけ似せた$\mathbf{D}\_{1}$を用意し、$\mathbf{D}\_{1}と\mathbf{D}\_{t}$の間で生じている相対的な動きを推定し、それが$\mathbf{S}\_{1}と\mathbf{S}\_{t}$でも生じていると考えることで、近似的なマッピングを得ます。

この考え方によって、マッピング関数は以下のように導出されます。詳しい導出は付録に記載されています。

```math
\mathcal{T}_{\mathbf{S}_{1} \leftarrow \mathbf{S}_{t}}(z) \approx \mathcal{T}_{\mathbf{S}_{1} \leftarrow \mathbf{R}}\left(p_{k}\right)+J_{k}\left(z-\mathcal{T}_{\mathbf{S} \leftarrow \mathbf{R}}\left(p_{k}\right)+\mathcal{T}_{\mathbf{D}_{1} \leftarrow \mathbf{R}}\left(p_{k}\right)-\mathcal{T}_{\mathbf{D}_{t} \leftarrow \mathbf{R}}\left(p_{k}\right)\right)
```

```math
J_{k}=\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{D}_{1} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)\left(\left.\frac{d}{d p} \mathcal{T}_{\mathbf{D}_{t} \leftarrow \mathbf{R}}(p)\right|_{p=p_{k}}\right)^{-1}
```

ここで注意すべき制約は、「$\mathbf{S}\_{1}$の姿勢にできるだけ似せた$\mathbf{D}\_{1}$を用意」するという点です。これは、avatarifyを上手に動かすときのコツでもあります。

# まとめ
以上、First Order Motion Model についてざっと解説してみました。

このような最先端の手法が簡単に動かせるというのは、それだけでワクワクしますね。
実際に使用する場合は頭の位置を最初にSource画像に合わせる必要があり、破綻の無い動画を得るためには少々コツが必要だったりします。また、Source画像としてイラスト画像を使用すると、明確に破綻するなどという限界もあります。そのあたりも、元論文を読むとなぜそのような破綻が生じるのか、わかるかと思います。