1. はじめに ~画像レジストレーション~
画像レジストレーションとは、異なる画像間における点と点の対応関係から、画像間に存在する幾何学的な変換モデル（例えばアフィン変換など）を推定し、画像間の位置合わせを行うことをいいます。
画像レジストレーションの応用先としては、以下のような例があります。
- 複数の画像を組み合わせたモザイク画像（パノラマ画像）の生成
- 超解像画像の生成, etc.
今回は、このような画像間の位置合わせを、畳み込みニューラルネットワークで行った論文を紹介し、実際にコードを動かしてみます。
2. 今回紹介する論文
題名 : Convolutional Neural Network Architecture for Geometric Matching
発表 : Spotlight & poster @ CVPR'17
実装 : https://github.com/ignacio-rocco/cnngeometric_pytorch
arXiv : https://arxiv.org/abs/1703.05593
3. 手法
提案手法のポイントは、以下の三点です。
- 特徴抽出、特徴マッチング、モデルパラメータ推定という画像レジストレーションの一連のパイプラインを、単一のCNN（geometric-matching CNN）でend-to-endに学習
- 人為的に生成した、幾何変換が既知の画像ペアを教師データとすることにより、学習にアノテーションが不要
- 見た目が大きく異なる画像間でもある程度正確に位置合わせができる(例えば、下の図に示すような違う車種のバイクでも位置合わせができる)
（論文の手法を用いて、左下の画像をターゲットとして左上の画像の位置合わせを行った例）
ポイント1とポイント2については、続けて詳細を説明します。また、ポイント3については、次のセクションで提案手法を実際に動作させ、いくつかのテストデータに対する結果を見てみることにします。
なお、引用元の表記が無い図は全て論文からの引用です。
3.1 Geometric-matching CNNのアーキテクチャ
Geometric-matching CNNのアーキテクチャは、下図のとおりです。
上の図からもわかるとおり、大きく以下の三つのパートから構成されています。
- CNNに画像$I_A$および$I_B$を入力し、特徴マップ$f_A$と$f_B$を得る
- 特徴マップ$f_A$と$f_B$が保持する各特徴ベクトルに対して、特徴の類似度$f_{AB}$を計算する
- 特徴の類似度$f_{AB}$をCNNに入力し、画像の幾何変換パラメータ$\theta$を得る
以下では、それぞれのパートについて順番に解説を行います。
特徴抽出
特徴抽出（図のFeature extraction CNNのブロック）の仕組みは、わりとシンプルです。
位置合わせを行う画像$I_A$と$I_B$のそれぞれから特徴を抽出するため、ImageNetで学習済みのVGG16ネットワークにそれぞれの画像を入力します。
VGG16のpool4レイヤが出力した特徴マップに対して、特徴ベクトル毎にL2-normalizationを行います。このようにして得られた特徴マップを、$I_A$および$I_B$から抽出された特徴マップ$f_A$および$f_B$とします。
論文では$I_A$および$I_B$の画像サイズは縦横240ピクセルに設定されています。VGG16のpool4が出力する特徴マップは、チャンネル数512で、サイズは入力画像サイズの1/16となるので、$f_A$および$f_B$のサイズは、H×W×D = 15×15×512となります（前述のL2-normalizationは、512次元の特徴ベクトルごと独立に行われます）。
なお、学習時にはVGG16の重みパラメータも更新されます。
特徴マッチング
特徴マッチング（図のMatchingのブロック）では、特徴マップ$f_A$を構成する特徴ベクトルと、同じく特徴マップ$f_B$を構成する特徴ベクトルについて、特徴類似度$f_{AB}$を計算します。
その1段階目のステップでは、下図に示すcorrelation layerが使われます。
correlation layerでは、$f_B$を構成する512次元の特徴ベクトル$f_B(i, j)$と、$f_A$を構成する512次元の特徴ベクトル$f_A(i_k, j_k)$との内積$c_{AB}(i, j, k)$を計算します。
c_{AB}(i, j, k) = f_B(i, j)^T f_A(i_k, j_k)
$(i_k, j_k)$および$(i, j)$はそれぞれ、$f_A$および$f_B$を構成する特徴ベクトルの画素位置を表すインデックスです。また、$c_{AB}$の3つめの引数$k$と$(i_k, j_k)$には、以下のような関係があるものとします。
k = h \times j_k + i_k
$h$と$w$は、$f_A$および$f_B$のサイズ（それぞれ高さと幅）とします。
ここで重要なのは、$c_{AB}(i, j)$は、$f_B$の同じ画素位置にある特徴ベクトル$f_B(i, j)$と、$f_A$を構成するすべての特徴ベクトルとの類似度（内積）を保持している という点です。
すなわち、
- $c_{AB}(i, j, 0)$ は、 $f_B(i, j)$ と $f_A(0, 0)$ との内積,
- $c_{AB}(i, j, 1)$ は、 $f_B(i, j)$ と $f_A(1, 0)$ との内積,
- ...,
- $c_{AB}(i, j, w \times h - 1)$ は、 $f_B(i, j)$ と $f_A(h-1, w-1)$との内積
となります。
したがって今回の場合、$c_{AB}$の形状は、 H×W×D = h×w×(w×h) = 15×15×225 です。
次に、特徴マッチングの2段階目のステップとして、$c_{AB}$に対して以下のような操作を行い、最終的な特徴類似度$f_{AB}$を得ます。
- 各要素に対してReLU活性化関数を作用させ、負の類似度をゼロに置き換える
- $c_{AB}(i, j)$に対して、L2-normalizationによる正規化を行う
この2つ目の操作が特に重要で、ambiguous matchesを抑制する効果があります。
一般に、特徴のマッチングでは、類似度が一番高い特徴間を単純に対応づけるのではなく、一番高い類似度と二番目に高い類似度との差が一定以上のときにのみ、特徴間の対応付けを行うという操作を行います。これはsecond nearest neighbor testと呼ばれ、画像内に繰り返しパターンが存在する際などに、特徴間の誤った対応付けを行うことを防ぐ効果があります。
前述の2つ目の操作を改めて考えると、仮に$c_{AB}(i, j)$の要素が一つだけ突出して高い値（類似度）であった場合、L2-normalizationによってその類似度の値が強調されます。逆に、$c_{AB}(i, j)$の要素がどれも同じような値（類似度）であった場合、L2-normalizationによって類似度の値が縮小され、second nearest neighbor testと同じような効果が得られることが直感的にわかるかと思います。
最終的な特徴類似度$f_{AB}$の形状は、$c_{AB}$と同じで、H×W×D = h×w×(w×h) = 15×15×225 となります。
以上の特徴マッチングの処理は、文章ではイメージしづらいのですが、筆者が公開している動画を見ると、非常にわかりやすいです（特徴マッチングは、動画の1:20頃）。
モデルパラメータ推定
特徴の類似度$f_{AB}$を畳み込みニューラルネットワークに入力し、画像の位置合わせに必要な幾何変換パラメータ$\theta$を回帰します。
ネットワークの構造は上に示すとおりで、covolution→batch-normalization→ReLUを2セット行ったあと、全結合層を通じてパラメータを回帰します。
どの幾何変換モデルを使用するかは、予め決めておく必要があります。例えばアフィン変換をモデルとして採用する場合、P=6次元の独立なパラメータを回帰します。
3.2 Geometric-matching CNNの学習
教師データの生成
あるオリジナルの画像$I_A$に対して、ランダムな幾何変換を施し、画像$I_B$を生成します。
当然のことながらこのときの幾何変換パラメータ$\theta$というのは既知なので、色々な画像$I_A$に対してランダムな幾何変換を施して画像$I_B$を生成することにより、教師データ$(I_A, I_B, \theta)$を無数に生成できます。
生成された無数の画像ペア$I_A, I_B$および、変換パラメータ$\theta$を用いることで、人によるアノテーションを行うことなく、前述のgeometric-matching CNNを学習することが可能となります。
なお、幾何変換によって生じるborder effectを防ぐため、変換元の画像$I_A$には、画像周辺にpaddingが施されます（上の図を参照）。
変換される画像$I_A$ですが、論文では、東京のGoogleストリートビューの画像から製作された "Tokyo Time Machine dataset" と呼ばれるデータが使用されています。
損失関数の定式化
CNNの学習時に重要なのが損失関数の設計です。
普通に考えると、CNNが出力した幾何変換パラメータと、真の幾何変換パラメータとの差の絶対値を損失関数とすればよさそうですが、その場合、使用する幾何変換モデルごとに（例えば、アフィン変換やホモグラフィー変換などに対して別個に）ハイパーパラメータのチューニングが必要になってしまいます。
そこで論文では、使用する幾何変換モデルの種類に依存しない以下のような損失関数を提案しています。
- 画像$I_A$の上で適当な大きさのグリッドを定義し、それにしたがって等間隔に点をサンプリングする
- 真の幾何変換パラメータおよびCNNが推論した幾何変換パラメータのそれぞれを用いて、サンプリング点に対して幾何変換を行う
- サンプリング点の変換後の二点間の距離を計算し、その二乗和を損失関数とする
このような損失関数の定式化を行うことにより、幾何変換モデルの種類によらず一貫した方法で学習を行うことができるようになります。
3.3 異なる幾何変換の組み合わせ
画像レジストレーションでは、アフィン変換など比較的簡単なモデルによって画像の位置をおおまかに合わせた後、thin-plate-spline(TPS)など比較的自由度の高いモデルによって画像の細かな位置合わせを行うことが一般的です。
論文でもこのような考えにもとづいて、アフィン変換とthin-plate-spline変換（以下、TPS変換と書く）を組み合わせた、二段階構成の画像レジストレーションを行っています（上の図も参照）。
- アフィン変換を行うCNNとTPS変換を行うCNNの二つのネットワークを、あらかじめ別々に学習しておく
- まず、アフィン変換CNNに画像$I_A$と$I_B$を入力し、推論された結果を用いて$I_A$に対してアフィン変換を施す
- 次に、$I_A$に対してアフィン変換を施した画像と$I_B$とをTPS変換CNNに入力し、TPS変換によって$I_A$の細かな幾何変換を行う
このような二段構えの位置合わせを行うことで、位置合わせ精度の安定性を保ちつつ、様々なバリエーションの位置合わせに対応できる構成を実現しています。
4. 動作確認
4.1 環境構築
AnacondaでPython=3.7.0の仮想環境を構築し、そこに実行環境を作ることにしました。
今回は推論だけなので、GPUは使わずCPU環境で実行しています。
# Create python environment with conda
$ conda create -n cnn_geometric python=3.7
$ source activate cnn_geometric
# Clone repo
$ git clone https://github.com/ignacio-rocco/cnngeometric_pytorch
$ cd cnngeometric_pytorch
# Install Pytorch
$ conda install pytorch torchvision -c pytorch
# Install some Python modules
$ conda install jupyter scikit-image matplotlib pandas requests
# Launch jupyter notebook
$ jupyter notebook
# Then, access `http://localhost:8888` and open `demo_notebook.ipynb`
4.2 推論結果の確認
テストデータとして準備されている水鳥の画像を使って、位置合わせを行ってみます。
左から2番目の画像（ターゲット画像）に対して、左端の画像（ソース画像）の位置合わせを行っています。
左から3番目の画像がアフィン変換CNNのみを使って位置合わせを行った結果で、その右側がTPS変換CNNのみを使って位置合わせを行った結果です。
さらに、右端がアフィン変換CNNとTPS変換CNNの両方を使って位置合わせを行った結果です。
ソース画像とターゲット画像間で、水鳥の模様や水面の模様などに結構な違いがありますが、位置合わせを行った右端の画像とターゲット画像を見比べると水鳥の位置がほぼピッタリ一致しており、位置合わせが上手くいっていることがわかります。
学習時には、Googleストリートビューの画像とその画像に幾何変換を施しただけの画像を与えているだけにも関わらず、テクスチャなどに結構な違いのある画像の位置合わせが行えるというのは驚きです。物体の輪郭などをとらえるような学習が内部で行われているのだと考えられます。
他のテストデータに対する結果もいくつか載せておきます：
5. まとめと所感
今回、画像レジストレーションにおける特徴抽出、特徴マッチング、モデルパラメータ推定という一連のパイプラインを、end-to-endで学習するCNNについて紹介しました。
いくつかのサンプルに対する主観評価ではありますが、紹介した手法は、異なるドメインの画像間の位置合わせに対しても非常に有効であることが確認できました。
また、位置合わせの結果はさることながら、CNNのモデル化の隅々に人間（著者）の知恵がふんだんに織り込まれていることに面白さを感じました。特徴抽出などdeep learningが得意とするところは学習に任せる一方で、人間がモデル化できるところはモデル化してあげるのが、（画像位置合わせのタスクに限らず）よい機械学習モデルを構築する一つのポイントなのではないかと感じました。