概要
劣化した古い写真を復元する深層生成モデルを提案している論文を紹介します。本論文には、ジャーナル版があり、そちらではさらに人間の顔領域に特化した修復手法も追加で提案されていますが、本稿ではそれには触れません。
本論文では、古い写真の復元を、古い写真$\mathcal{R}$、劣化のない写真$\mathcal{Y}$、それを人工的に劣化させた写真$\mathcal{X}$という3つのドメインの間のドメイン変換の問題として捉えます。
本手法では、2つのVAEと1つのドメイン変換モデル$T_{\mathcal{Z}}$を使用してこれを実現します。1つ目のVAEは、古い写真$r \in \mathcal{R}$および合成劣化写真$x \in \mathcal{X}$をエンコード・デコードするVAEです。つまり、古い写真と合成劣化写真は共通の潜在空間にマッピングされます。2つ目のVAEは、劣化のない写真$y \in \mathcal{Y}$をエンコード・デコードするVAEです。こちらの潜在空間は、1つ目のVAEとは異なる空間です。$\mathcal{X, Y}$は、合成画像と元画像の集合ですので、1対1の対応が存在しています。これを利用すると、2つのVAEの潜在空間を変換するマッピングネットワーク$T_{\mathcal{Z}}$を構築することが可能です。
下図は、これら3つのドメインとそれらを結ぶVAE、マッピングネットワークを示したものです。
以下、もう少し詳しくみていきます。まずはネットワーク構造を確認します。次に、モデルの訓練に必要なデータの作り方について簡単に触れます。最後に、モデルの訓練のための損失関数を確認します。
書誌情報
- Wan, Ziyu, et al. "Bringing Old Photos Back to Life." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.
- 論文(CVPR2020)
- ジャーナル版論文
- 公式実装(PyTorch)
- Colab
モデルの構造
冒頭で示したように、本手法は2つのVAEとマッピングネットワーク$\mathcal{T}$から成り立っています。以下の図はこれらの繋がりを表します。
2つのVAE
以下の表は、本手法で使用するネットワークの構造を示した物です。
VAEに含まれるエンコーダー（Encoder）とデコーダー（Generator）は、以下の表のEncoderとGeneratorに示される共通した構造になっています。本手法ではVAEに対してGANの損失も使用しますので、厳密には、以下の表に示されているネットワーク以外にも、Discriminatorが存在します。
マッピングネットワーク
マッピングネットワーク$\mathcal{T}$には、Conv層やResBlockといった見慣れたモジュールに混じって、Partial nonlocalというモジュールが含まれています。このモジュールは、本手法のために提案されたものです。
マッピングネットワークの目的は、エンコーダー$E_{\mathcal{R, X}}$によって得られる潜在変数$z_{\mathcal{R}}(z_{\mathcal{X}})$を、$G_{\mathcal{Y}}$の入力に使われる潜在変数$z_{\mathcal{R\rightarrow Y}}(z_{\mathcal{X\rightarrow Y}})$へと変換することです。しかし、このマッピングはそこまで簡単なことではありません。
マッピングを難しくしているのは、古い画像に対して生じる構造的な劣化と非構造的な劣化の混在です。まず、引っかき傷やホコリの斑点のような劣化は構造的な劣化と呼ばれ、大域的なコンテクストを考慮して修復する必要があります。また、空間的に比較的均一に生じる非構造的な劣化として、たとえばフィルムグレインや色あせがあり、これらは近隣のピクセルに対する均質化操作によって修復する必要があります。
そのため、本手法では、マッピングネットワークに構造的劣化の位置を示すマスクを入力することで、これら2種類の劣化を明示的に分離した上で補完しつつ、潜在変数間のマッピングを行えるようなネットワークを提案しています。その鍵となるのがPartial nonlocalです。
以下の図は、Partial nonlocalを説明した図です。左側は、構造的劣化が生じている領域（Hole）が、周囲から集約された特徴によって補完される様子を表しています。一方で右側は、Partial nonlocalのネットワーク構造を表しています。通常のNonlocal networkと異なるのは、構造的劣化が生じている領域を表すマスク画像$m$も入力されているという点です。
Nonlocal networkは、特徴マップの大域的な特徴を集約することができる構造です。一方、本論文で提案しているPartial nonlocalは、構造的な劣化の位置を表すマスク画像$m$を考慮することで、劣化が生じている領域を避けつつ大域的な特徴を集約し、劣化が生じている部分を補完するのに使用できる特徴を作成することに成功しています。
Partial nonlocalは、以下のように計算されます。
O_{i}=\nu\left(\sum_{\forall j} s_{i, j} \mu\left(F_{j}\right)\right)
$s_{i,j}$は、あるピクセルが別のどのピクセルを重視するかという対応の重みを表します。$s_{i,j}$は以下のように計算されますが、マスク領域を考慮することが通常のNonlocal blockと異なるところです。
s_{i, j}=\frac{\left(1-m_{j}\right) f_{i, j}}{\sum_{\forall k}\left(1-m_{k}\right) f_{i, k}}
ここで、$f_{i,j}$は以下のように求められます。
f_{i, j}=\exp \left(\theta\left(F_{i}\right)^{T} \cdot \phi\left(F_{j}\right)\right)
なお、以上の式中に現れる$\theta, \phi, \mu, \nu$はいずれも$1\times1$Conv層です。
入力される特徴マップ$F$、nonlocalブロックの出力$O$、そしてマスク画像$m$を用いると、構造的な劣化が生じている位置の特徴を補完することができます。これは、以下のような式で表されます。$\rho_{\text {local}}, \rho_{\text {global}}$は、個別に訓練される非線形変換を表します。
F_{fuse}=(1-m) \odot \rho_{\text {local}}(F)+m \odot \rho_{\text {global}}(O)
データの作成方法
モデルをどのように訓練するのかという話に入る前に、その際に使用するデータの作成方法を確認しておきます。
合成劣化写真
まず、合成劣化写真の作り方を確認しておきます。
本研究では、実際の写真に生じている劣化を完全にシミュレートすることはできませんが、以下のような比較的近い画像を作成することに成功しています。
まず、様々な傷や紙のテクスチャを収集します。これらは、ランダムな透明度で写真にブレンドされます。次に、大きな面積の劣化をシミュレートするために、ランダムな形状や毛羽立ちの穴を生成します。最後に、大域的な劣化である色あせやノイズを合成しています。
構造的劣化のマスク画像
Partial nonlocalに入力するマスク画像$m$は、U-Net風のセグメンテーションモデルによる推定結果が用いられますが、このモデルを訓練するためには構造的劣化の領域を表すマスク画像が必要です。
合成劣化写真に対するマスク画像は、合成過程で使用した傷や穴から作成することができます。この画像を利用すると、合成劣化写真からマスク画像を推定するセグメンテーションモデルを、一応は構築することができます。
しかし、合成データだけで訓練したモデルでは、実際の写真に適用したときに、あまり良い結果は得られません。そこで、本研究では、783枚の古い写真に対して人手でアノテーションを行い、実際の劣化写真に対するマスク画像を作成しています。合成データで訓練した後、実際の写真に関するデータ400枚を用いて、ファインチューニングを行うことで、最終的なセグメンテーションモデルを構築しています。
モデルの訓練と損失関数
訓練の手順としては、まず2つのVAEを訓練し、次にそれらの潜在空間をつなぐマッピングネットワーク$\mathcal{T}$を訓練します。それぞれの段階でどのような損失関数が必要かを確認します。
VAEの損失関数
VAEには、現実の古い写真$\mathcal{X}$と合成劣化写真$\mathcal{X}$に関するVAE1、現実のきれいな写真$\mathcal{Y}$に関するVAE2が存在しており、それぞれ少し異なる損失関数が設定されています。
VAE1
VAE1では、以下のような損失関数を用いて訓練が行われます。
\min _{E_{\mathcal{R}, \mathcal{X}}, G_{\mathcal{R}, \mathcal{X}}} \max _{D_{\mathcal{R}, \mathcal{X}}} \mathcal{L}_{\mathrm{VAE}_{1}}(r)+\mathcal{L}_{\mathrm{VAE}_{1}}(x)+\mathcal{L}_{\mathrm{VAE}_{1}, \mathrm{GAN}}^{\operatorname{latent}}(r, x)
上式を構成している3つの要素は、それぞれ、$\mathcal{R}$に関するVAE1の損失、$\mathcal{X}$に関するVAE1の損失、$\mathcal{R,X}$の潜在空間を近づけるためのGAN損失になっています。それぞれを詳しくみていきましょう。
まず、$\mathcal{R}$に関するVAE1の損失は、中身を見ると下式で表されます。
\begin{aligned} \mathcal{L}_{\mathrm{VAE}_{1}}(r) &=\mathrm{KL}\left(E_{\mathcal{R}, \mathcal{X}}\left(z_{r} \mid r\right)|| \mathcal{N}(0, I)\right) \\ &+\alpha \mathbb{E}_{z_{r} \sim E_{\mathcal{R, X}}\left(z_{r} \mid r\right)}\left[\left\|G_{\mathcal{R}, \mathcal{X}}\left(r_{\mathcal{R} \rightarrow \mathcal{R}} \mid z_{r}\right)-r\right\|_{1}\right] \\ &+\mathcal{L}_{\mathrm{VAE}_{1}, \mathrm{GAN}}(r) \end{aligned}
これは通常のVAEの損失（潜在変数の正則化項、再構成誤差）とGAN損失から成り立っています。GAN損失は、LSGANのleast-square lossが使用されています。GAN損失を加えることで、VAEが陥りがちなぼやけた生成結果になる問題を回避しています。
全く同様に、$\mathcal{X}$に関するVAE1の損失は、以下の式で表されます。
\begin{aligned} \mathcal{L}_{\mathrm{VAE}_{1}}(x) &=\mathrm{KL}\left(E_{\mathcal{R}, \mathcal{X}}\left(z_{x} \mid x\right)|| \mathcal{N}(0, I)\right) \\ &+\alpha \mathbb{E}_{z_{x} \sim E_{\mathcal{R, X}}\left(z_{x} \mid x\right)}\left[\left\|G_{\mathcal{R}, \mathcal{X}}\left(x_{\mathcal{X} \rightarrow \mathcal{X}} \mid z_{x}\right)-x\right\|_{1}\right] \\ &+\mathcal{L}_{\mathrm{VAE}_{1}, \mathrm{GAN}}(x) \end{aligned}
最後に、$\mathcal{R,X}$の潜在空間を近づけるためのGAN損失について確認します。合成劣化写真は、完全に古い写真の劣化を再現できるわけではないので、$\mathcal{X}$の潜在空間$\mathcal{Z}_{\mathcal{X}}$と$\mathcal{R}$の潜在空間$\mathcal{Z}_{\mathcal{R}}$は、一致しません。それどころか、全く別の部分空間にマッピングされてしまう危険性もあります。本手法が単純なAutoEncoderではなく、VAEを使用している理由は、正則化項によって2つのドメインの潜在空間が、$\mathcal{N}(0, I)$という同一の空間に収まることを促す効果を期待しているためです。
GAN損失は、潜在変数に対してGANを適用して、2つのドメインの潜在空間がさらに近づくように促進します。下式の$D_{\mathcal{R}, \mathcal{X}}$は、潜在変数が$\mathcal{R, X}$のどちらからきたのかを推定できるように訓練されるDiscriminatorです。一方でエンコーダー$E_{\mathcal{R}, \mathcal{X}}$は、入力画像が$\mathcal{R, X}$のどちらから来たのかを騙せるように訓練されます。これにより、潜在空間$\mathcal{Z}_{\mathcal{X}}, \mathcal{Z}_{\mathcal{R}}$をできるだけ近づけることができます。
\begin{aligned} \mathcal{L}_{\mathrm{VAE}_{1}, \mathrm{GAN}}^{\text {latent }}(r, x)=& \mathbb{E}_{x \sim \mathcal{X}}\left[D_{\mathcal{R}, \mathcal{X}}\left(E_{\mathcal{R}, \mathcal{X}}(x)\right)^{2}\right] \\ &+\mathbb{E}_{r \sim \mathcal{R}}\left[\left(1-D_{\mathcal{R}, \mathcal{X}}\left(E_{\mathcal{R}, \mathcal{X}}(r)\right)\right)^{2}\right] \end{aligned}
VAE2
VAE2に関する損失は、他のドメインに関する損失と同様のものが使われ、これ以外の損失はありません。
\begin{aligned} \mathcal{L}_{\mathrm{VAE}_{2}}(y) &=\mathrm{KL}\left(E_{\mathcal{Y}}\left(z_{y} \mid y\right)|| \mathcal{N}(0, I)\right) \\ &+\alpha \mathbb{E}_{z_{y} \sim E_{\mathcal{Y}}\left(z_{y} \mid y\right)}\left[\left\|G_{\mathcal{Y}}\left(y_{\mathcal{Y} \rightarrow \mathcal{Y}} \mid z_{y}\right)-y\right\|_{1}\right] \\ &+\mathcal{L}_{\mathrm{VAE}_{2}, \mathrm{GAN}}(y) \end{aligned}
マッピングネットワークの損失関数
潜在変数のマッピングネットワーク$\mathcal{T}$については、以下のような損失が使われています。
\mathcal{L}_{\mathcal{T}}(x, y)=\lambda_{1} \mathcal{L}_{\mathcal{T}, \ell_{1}}+\mathcal{L}_{\mathcal{T}, \mathrm{GAN}}+\lambda_{2} \mathcal{L}_{\mathrm{FM}}
左から順番に、合成劣化画像の潜在変数から元の画像の潜在変数への再構成誤差、マッピング結果が真の潜在変数に近づくようにするGAN損失、修復画像と元の画像の再構成との知覚的な距離を表すFeatureMatchnig損失となっています。
再構成誤差
もともと$\mathcal{X}$の画像は、$\mathcal{Y}$の画像を元に加工されている物なので、対応する画像が存在します。
これを利用して、マッピングされた合成画像についての潜在変数$\mathcal{T}(z_x)$と元画像の潜在変数$z_y$のL1距離による直接比較を行えます。
\left.\mathcal{L}_{\mathcal{T}, \ell_{1}}=\mathbb{E} \| \mathcal{T}\left(z_{x}\right)-z_{y}\right) \|_{1}]
GAN損失
修復画像$x_{\mathcal{X} \rightarrow \mathcal{Y}}$が、元画像の再構成$y_{\mathcal{Y} \rightarrow \mathcal{Y}}$と見分けがつかない品質になるように、GANの枠組みを使います。Discriminator$D_{\mathcal{T}}$を用いて、損失関数は以下のように定義されます。ここでも、LSGANが使われています。
\begin{aligned} \mathcal{L}_{\mathcal{T}, \mathrm{GAN}} =& \mathbb{E}_{x \sim \mathcal{X}}\left[D_{\mathcal{T}}\left(x_{\mathcal{X} \rightarrow \mathcal{Y}}\right)^{2}\right] \\ &+\mathbb{E}_{y \sim \mathcal{Y}}\left[\left(1-D_{\mathcal{T}}\left(y_{\mathcal{Y} \rightarrow \mathcal{Y}}\right)\right)^{2}\right] \end{aligned}
FeatureMatching損失
修復画像$x_{\mathcal{X} \rightarrow \mathcal{Y}}$は、元画像の再構成$y_{\mathcal{Y} \rightarrow \mathcal{Y}}$の見た目ができるだけ近づくようにしたいので、Discriminator$D_{\mathcal{T}}$とVGGの中間層の活性化された特徴マップのL1距離を利用して、その差異を評価します。
\begin{aligned} \mathcal{L}_{\mathrm{FM}} &=\mathbb{E}\left[\sum_{i} \frac{1}{n_{D_{\mathcal{T}}}^{i}}\left\|\phi_{D_{\mathcal{T}}}^{i}\left(x_{\mathcal{X} \rightarrow \mathcal{Y}}\right)-\phi_{D_{\mathcal{T}}}^{i}\left(y_{\mathcal{Y} \rightarrow \mathcal{Y}}\right)\right\|_{1}\right.\\ &\left.+\sum_{i} \frac{1}{n_{\mathrm{VGG}}^{i}}\left\|\phi_{\mathrm{VGG}}^{i}\left(x_{\mathcal{X} \rightarrow \mathcal{Y}}\right)-\phi_{\mathrm{VGG}}^{i}\left(y_{\mathcal{Y} \rightarrow \mathcal{Y}}\right)\right\|_{1}\right]\end{aligned}
まとめ
以上、古い写真の劣化を修復する手法を提案している論文について紹介しました。
ちゃんと読む前は、簡単な問題設定のように早とちりしていたのですが、3つのドメインを必要とする意外に複雑な問題設定と、Partial nonlocalのような丁寧なモデルの構造には驚きました。構造的劣化のマスク画像を人手で整備して不足しているデータをちゃんと自分たちで作るなど、労力をきっちりかけている良い研究だと思います。