【論文読解】Correspondence Networks with Adaptive Neighbourhood Consensus

# 概要
2 つの画像中の同一カテゴリの物体の密なポイント対応を得るための手法ACN-Netを提案した論文"Correspondence Networks with Adaptive Neighbourhood Consensus"を読んでみたので、紹介してみます。本論文は、CVPR2020 に accept されています。

密なポイントの対応づけによって、たとえば下図のようにソース画像中の物体をターゲット画像中の物体の姿勢や形に変換する、といったことも可能になります。
本論文で提案されている手法である ANC-Net の結果は、一番右側に示されています。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/bcbf05f2-2a37-66ab-13ab-7374e7f23a84.png)

このような異なる画像間のポイント対応を求めることは、ピクセルレベルの画像間の対応関係を求めることになります。これにより、異なる位置から同一の物体を撮影した際のピクセルの対応付けや、グローバルな幾何的変換、オプティカルフローの推定といった後続タスクに使用できます。

本手法では密なポイント対応を求めることができますが、モデルの訓練に必要なのは疎なポイント対応だけです。そのため、既存のキーポイント（ランドマーク）がアノテーションされているデータセットを使用できます。

# 書誌情報

- Li, Shuda, et al. "Correspondence Networks with Adaptive Neighbourhood Consensus." arXiv preprint arXiv:2003.12059 (2020).
- [プロジェクトページ](https://ancnet.avlcode.org)

# ネットワーク構造
まずはネットワーク構造をざっと確認しましょう。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/7d00419f-377c-35c5-b612-6c75eb1ee49b.png)

$\mathbf{I^s}, \mathbf{I^t}$という 2 つの入力画像に対して、特徴抽出を行う CNN$\mathcal{F}$を適用します。これにより、特徴マップ$\mathbf{F^s}, \mathbf{F^t}$が抽出されます。$s, t$はそれぞれ、ソース、ターゲットを表します。$\mathcal{F}$は ImageNet で学習済みの ResNet や ResNeXt で、重みはフリーズさせているようです。

それぞれの特徴マップをもとに、Self-Similarity$\mathbf{S^s}, \mathbf{S^t}$が抽出されます。Self-Similarity は、特徴マップ内の高次のパターンを捉えるための特徴です。

ソースとターゲットの 2 つの特徴マップから、CorrelationMap$\mathbf{C_f}$が得られます。これにより、ソース側の特徴マップのあるセルが、ターゲット側の特徴マップのどのセルと類似しているのかを表現します。画像中では、このマップは 2 次元であるように見えますが、実際には$H_s \times W_s \times H_t \times W_t$というサイズの 4 次元テンソルです。Self-Similarity についても同様の処理を行い、CorrelationMap$\mathbf{C_s}$を得ます。

最後に、 $\mathbf{C_f}$と$\mathbf{C_s}$は、それぞれ非対称なサイズのカーネルを持つ 4 次元畳み込み操作を行うANCモジュール$\mathcal{N}$によって洗練され、そのあと統合されます。これにより、最終的な CorrelationMap$\bar{\mathbf{C}}$が得られます。
これを用いることで、ソース側の座標$(i, j)$に対応する点は$\bar{\mathbf{C}}_{i, j, *, *}$ のうちもっとも大きい値になっているターゲット側の座標$(k, l)$だよね、ということが求められます。

# Self-Similarityの算出
本論文の提案の 1 つは、Self-Similarity をシンプルに算出する、というものです。Self-Similarityは、ある点が周囲の点とどれだけ似ているのかを算出し、その空間的な分布パターンを表現した特徴量です。下図にこの算出の過程を示します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/7d486117-5d7c-101c-2595-cc57c37e8635.png)

興味のあるセルの特徴量と、自身を含む近隣の 9 つのセルの特徴量との Cosine 距離を算出します。これをベクトル化し、9 次元の特徴量が得られます。これを特徴マップ中の全セルに対して計算することで $\mathbf{S_0} \in \mathbb{R}^{h_f \times w_f \times 9}$が算出されます。

それに対して 5×5Conv と ReLU を適用した$\mathbf{S_1}$、さらに 5×5Conv-ReLU を適用した$\mathbf{S_2}$を算出します。これらは$\mathbf{S_0}$より高次のパターンを表現していると言えます。これらを結合して最終的な$\mathbf{S}$を得ます。

既存手法である DCCNet[^1]では下図のように、一度算出した Self-Similarity を元の特徴マップと組み合わせて Context-aware Feature という高次の特徴を得るという処理を行っています。本手法では Self-Similarity は特徴マップから分離され、独立して処理されるので、よりシンプルな処理になっていると言えます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/bf37b753-edd5-2711-7aaa-98d372d79738.png)

[^1]: Shuaiyi Huang, Qiuyue Wang, Songyang Zhang, Shipeng Yan, and Xuming He. Dynamic Context Correspondence Network for Semantic Alignment. In Proceedings of Intl. Conf. on Computer Vision (ICCV), 2019.

このように算出された Self-Similarity はソース側とターゲット側で算出され、組み合わせることで CorrelationMap$\mathbf{C_s}$が算出されます。一方で特徴マップ側でも、$\mathbf{C_f}$が計算されます。$\mathbf{C_f}, \mathbf{C_s}$は、それぞれ 1 次の特徴の相関と、高次な特徴の相関を担当していると言えます。

# ANCモジュール
2 つの CorrelationMap$\mathbf{C_f}, \mathbf{C_s}$は、それぞれノイズ除去によって洗練された後、$\bar{\mathbf{C}}$に統合されます。CorrelationMap はノイズが多くなってしまいがちなので、このようなフィルター処理が既存の手法でも適用されています。

本論文の提案のもう 1 つは、このノイズ除去のための ANC(Adaptive neighbourhood consensus)モジュールで、下式では$\mathcal{N}$に対応します。

$\bar{\mathbf{C}}=\mathcal{N}\left(\mathbf{C}\_{s}\right)+\left(\mathcal{N}\left(\mathbf{C}\_{s}^{\top}\right)\right)^{\top}+\mathcal{N}\left(\mathbf{C}\_{f}\right)+\left(\mathcal{N}\left(\mathbf{C}\_{f}^{\top}\right)\right)^{\top}$

ここで、$\mathcal{N}$は非対称なサイズのカーネルを使った 4 次元畳み込み操作です。この非対称なカーネルサイズによって、同一カテゴリの物体でも異なる大きさや形で現れる、という問題を吸収できます。

また、$\top$は転置を表しますが、ここでは 4 次元テンソルの先頭 2 次元とその他の 2 次元を入れ替える操作を表し、$(\mathbf{C}^{\top})\_{i,j,k,l} = \mathbf{C}_{k,l, i, j}$となります。これは、ソース画像とターゲット画像が入れ替わったときに得られる CorrelationMap と同一のものです。

このような入れ替えは、4 次元畳み込みのカーネルが例えば$5 \times 5 \times 5 \times 5$のような対称なサイズであれば、あまり意味のない操作となってしまいます。しかし、$3 \times 3 \times 5 \times 5$のように非対称なカーネルであれば、ソース画像とターゲット画像を入れ替えることで初めて拾える対応関係があるので、意味のある操作になります。

本論文では、非対称なカーネルサイズを持つ 4 次元畳み込みは、non-isotropic 4D convolution と呼ばれています。イメージとしては下図の右のように、異なるカーネルサイズがソース側とターゲット側に適用される、という処理になります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/18e0dd95-5af3-67cd-9bff-43ccdcf700cd.png)

本論文では、このような 4 次元畳み込み操作として、上図の(b),(c),(d)の 4 パターンを試しています。実験の結果、(d)がもっとも良い結果となったと結論づけています。

ここまでで得られた$\bar{\mathbf{C}}$は、いわばソース側とターゲット側の特徴マップの対応の強さを表すスコアのヒートマップと言えます。詳細は省きますが、正規化と softmax 操作によって、ソース側の座標とターゲット側の座標の対応を確率で表すマップ$\mathbf{M}$を得ることができます。例えば、ソース側の座標$(i, j)$がターゲット側の座標$(k, l)$に対応する確率や、その逆を求めることができます。

# ターゲットと損失関数

## キーポイントの対応に関する損失
本手法では、ソース画像とターゲット画像の疎なキーポイントの対応が教師データとして与えられています。これをもとに、ソース側の座標とターゲット側の座標の対応を確率で表すマップ$\mathbf{M}$の目標となるマップを作る必要があります。

理想的には、ある 1 点のみが 1、その他は 0 となっている確率マップを目標にしたいと考えます。しかし、キーポイントの対応は画像解像度におけるピクセルレベルで表現されているので、確率マップ$\mathbf{M}$の解像度に変換する必要があります。

下図のように、あるキーポイント$(x, y)$は、特徴マップの解像度における座標$(x_n, y_n)$に変換されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/5112bf88-cb76-2dab-b019-3d17ed918189.png)

この座標は多くの場合サブピクセルになるため、(c)のように近傍 4 点$(x^i_n, y^i_n)$を求め、距離に応じた重み$t_i$を目標として設定します($\sum_{i} t_i = 1$)。実際には、多くのキーポイント検出系の手法で用いられるように、(c)のような目標マップに対して 2 次元ガウスぼかしを加えることで最終的な目標$\mathbf{M}_{gt}$としています。

確率マップ$\mathbf{M}$には、ソース座標からターゲット座標への対応$\mathbf{M}^{s}$と、その逆$\mathbf{M}^{t}$という 2 種類が存在します。損失は、それぞれに対する目標マップとのフロベニウスノルムを使って算出されます。

$\mathcal{L}_{k}=\left\|\mathbf{M}^{s}-\mathbf{M}\_{g t}^{s}\right\|\_{F}+\left\|\mathbf{M}^{t}-\mathbf{M}\_{g t}^{t}\right\|\_{F}$

## 正則化のための損失
本論文の最後の提案点として、確率マップ$\mathbf{M}$に関する以下のような損失を提案されています。

$\mathcal{L}^m\_{o}=\left\|\mathbf{M}^{s} {\mathbf{M}^{s}}^{\top} -\mathbf{M}\_{g t}^{s}{\mathbf{M}\_{g t}^{s}}^{\top} \right\|\_{F}+\left\|\mathbf{M}^{t} {\mathbf{M}^{t}}^{\top} -\mathbf{M}\_{g t}^{t}{\mathbf{M}\_{g t}^{t}}^{\top}\right\|\_{F}$

この損失は、ソース側の座標とターゲット側の座標が、明確に 1 対 1 に対応しているべきという信念を表現するものと言えます。

全体の損失関数は、以下のようになります。$\alpha$はバランスのための係数で、0.001 が使用されています。

$\mathcal{L} = \mathcal{L}\_k + \alpha \mathcal{L}^m_o$

# まとめ
以上、簡単にですが ACN-Net を紹介しました。

Self-Similarity や、CorrelationMap に対する 4 次元畳み込み操作など、あまり触れたことのない概念を学べたのが良かったです。実装は今後[github](https://github.com/ActiveVisionLab/ANCNet)で公開される予定のようです。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/4010/6b61e96e-088b-250a-b9b5-aa8c1a1a66de.png)