画像内の物体の数を数える手法の論文紹介: CACViT

皆さん、こんにちは。AnyTechの立浪と申します。先日、バンクーバーで開催されたAAAI-24に行ってまいりました。そこでみたポスターの中から気になった論文を一つご紹介したいと思います。[2305.04440] Vision Transformer Off-the-Shelf: A Surprising Baseline for Few-Shot Class-Agnostic Countingという論文です。この論文で対象にしているタスクは"Few-Shot Class-Agnostic Counting"です。その名の通り、数えたい物体の少量の画像を与えて、検索対象の画像の中にその物体がいくつ含まれているかを推定するタスクです。このタスクでは訓練時に特定のクラスで条件づけをすることはしません。推論時には訓練データに含まれていないようなクラスの物体であっても、数え上げがうまくなされることが期待されます。

(Vision) Transformer程度の知識は仮定させていただきますが、この論文のコンセプトをなるべく容易に理解できるようにご説明したいと思います。私自身も、数え上げ系のタスクはほとんど触れたことがなく（一度レビューしたことがあるくらい）、ほとんど初心者です。そんな私でも理解しやすいくらい、CACViTはシンプルなフレームワークでした。おそらく、読者の皆さんも理解しやすいのではないかと思います。

背景

物体数え上げのタスクは人を数える（Crowd countingといいます）タスクのように、特定のドメインの物体を数え上げる"Class-Specific"のタスクが主流です。それに対して、近年は Class-Agnostic Counting (CAC)というタスクも提案されています。これは、任意のカテゴリの少数の画像を与えて、その物体の数え上げるタスクです。つまり、CACは"Class-Specific"な設定より汎化されたタスクです。

extract-then-match

CACの従来手法は"extract-then-match"と呼ばれるタイプだと著者らは主張しています。CACの設定ではクエリ画像(Query)からいくつかの模範画像(Exampler)の数を数え上げますが、クエリ画像と模範画像をそれぞれ特徴抽出してから、類似度計算をします。最近では、CounTR: [2208.13721] CounTR: Transformer-based Generalised Visual CountingのようにTransformerベースの手法もありますが、こちらもクエリ画像のViT特徴と模範画像のCNN特徴を利用しています。そして、これらの特徴の間の類似度をcross-attentionを使って計算します。このように、CounTRまではVision Transformerに本来備わっているself-attentionをうまく使えておらず、特徴抽出器も冗長でした。

extract-and-match

この論文の手法は"extract-and-match"です。つまり、クエリ画像と模範画像のそれぞれの特徴抽出器を用意する必要はなく、end-to-endで訓練できます。提案手法であるCACViTでは、クエリ画像と模範画像に同じViTを使い、そのattention mapを利用して、密度マップを予測します（なお、この密度マップのGround truthは合計すると検索したい物体の数になっています）。推論して得た密度マップから、最終的な物体の数を得ます。次以降から、このアルゴリズムの詳細についてみていきたいと思います。

アルゴリズム

CACViTの全体の概要は以下の図で示されています。クエリ画像といくつかの模範画像を入力し、最終的な出力した密度マップをGround truthとL2 lossを計算します。このL2 lossが小さくなるように勾配降下法で訓練します。このように、概要は非常に単純かと思います。まずは最も重要な左下の部分から理解していきましょう。

self-attentionに対するデカップリングの観点

CACViTでは、新しい観点でself-attentionを解釈しています。CACViTにはクエリ画像と模範画像のトークンたちをひとまとまりにして入力します。これは何を意味するかというと、attention mapに4つの領域が現れることを意味します。クエリ画像同士 $A_{query}=Q_{query} \otimes K_{query}$ 、模範画像同士 $A_{exp}=Q_{exp} \otimes K_{exp}$ 、それからクエリ画像と模範画像の組み合わせが2通り $A_{match}=Q_{query} \otimes K_{exp}$ 、 $A_{class}=Q_{exp} \otimes K_{query}$ があります。

論文では、これらの意味について考察されていました。

$A_{query}$ の意味

$A_{query}$ は素直に考えて、クエリ画像の特徴抽出に相当します。

$A_{exp}$ の意味

$A_{exp}$ も模範画像の特徴抽出に相当します。

$A_{match}$ の意味

$A_{match}$ はクエリ画像をクエリ、模範画像をキーなので、cross-attentionに相当します。自然な方法での特徴マッチングです。

$A_{class}$ の意味

$A_{class}$ は $A_{match}$ と似てはいますが、意味合いが少し違います。著者らは浅い層・深い層でのときそれぞれについて、次のように解釈しています。浅い層のself-attentionはself-attentionの能力が弱いので、フォアグラウンドの物体を持つクエリ画像のトークンは、カテゴリの情報が引き出されます。深い層のself-attentionはself-attentionの能力が十分あります。そのため、フォアグラウンドの物体を持つクエリ画像のトークンは、クエリ画像内において、フォアグラウンドの物体の繰り返しによって背景情報を含むクエリートークンに比べて相対的にself-attentionが強くなります。クエリ画像に対するattentionと模範画像に対するattentionの合計は(softmaxをとるので)一定なので、 $A_{class}$ を転置すると背景領域に強く反応するとのことです。私の理解によると、模範画像はトークンは相対的に数が少なく、注意の大半がクエリ画像に持っていかれるということを主張しているのだと思います。

論文では、上記を可視化した図が例示されています。

$A_{class}$ の解釈には議論の余地があると思いますが、この論文では $A_{match}$ を明示的に使用します。

アスペクト比とマグニチュードの問題への対象

この手法を素直に適用しようとすると、アスペクト比とマグニチュードの問題があります。

アスペクト比の問題を理解するために、模範画像の扱い方から見ていきましょう。提案手法では模範画像は固定のサイズで使用します。スキーのような縦横比率が異なるようなものを模範画像とする場合、固定のサイズにリサイズすることによって、アスペクト比が変わってしまいます。つまり、情報の一部が失われてしまうのです。これがアスペクト比の問題です。また、元々が大きい画像の模範画像であっても小さい画像であっても、同じスケールにリサイズします。そのため、このリサイズによってスケールの情報も失われてしまいます。

次にマグニチュードの問題を見ていきます。単純にsoftmaxを使っただけのattention mapでは、別の設定で同様の正規化された分布が得られる可能性があります。例えば、以下の図のように各領域に同じ比率で物体があるような問題設定を考えます。左と右では物体の数は違うのですが、すべてがうまくいけば同じattention mapが得られてしまうことになります。この問題を回避するために、スケールの因子が必要になりそうです。

提案手法では2つの問題それぞれに対処するために工夫が施されています。

アスペクト比を考慮したスケール埋め込み

ますアスペクト比への対処です。オリジナルの画像の幅を $W_k$ 、高さを $H_k$ とします。これらを模範画像の幅 $W_z$ と高さ $H_z$ のマップに埋め込みます。まずは幅から考えます。そのために、 $0$ から $W_k$ までの線形補間を考えます。補間の点の数は始点と終点を含めて、 $W_z$ 点にします。これをブロードキャストしてマップ $\hat{W}_k \in \mathbb{R}^{W_z \times H_z}$ を作ります。高さも同様に、 $0$ から $H_k$ までの線形補間を $H_z$ 点になるようにとり、ブロードキャストしてマップ $\hat{H}_k \in \mathbb{R}^{W_z \times H_z}$ を作ります。これらを足した $S_k = \hat{W}_k + \hat{H}_k$ をスケール埋め込みとして、リサイズした模範画像 $z_k$ とチャンネルに沿って結合します。（これをやるとクエリ画像と模範画像のチャネル数が異なってしまい、入力周りで何らかの工夫が必要だと思うのですが、この点の工夫が書かれていません。クエリ画像と模範画像で別のパッチ埋め込みの行列を使うのでしょうかね？）

マグニチュードの埋め込み

大きさの問題への対処です。模範画像の幅 $W_k$ 、高さ $H_k$ とおき、パッチの幅 $W_p$ 、高さ $H_p$ とおくと、大きさのオーダーはざっくりと次のように表すことができます。

$ME_k = \dfrac{W_p \times H_p}{W_k \times H_k}$

模範画像が複数あれば、それぞれに対しての $ME_k$ を計算できます。最終的なマグニチュードの埋め込み $ME$ はこれらを平均します。 $ME$ を類似度スコアに掛け合わせて、最終的な類似度スコアを取得します。

アーキテクチャや設定など

ネットワークアーキテクチャは次のとおりです。まず、384x384の解像度のクエリ画像と64x64の模範画像を入力します。パッチサイズはいずれも16x16です。特徴抽出器としてMasked Autoencoer(MAE)を利用し、その後に追加のTransformerブロック(ただし次元を削減し、512とする）を3つ付け足します。回帰のデコーダーとして4つのアップサンプリングブロック（畳み込みと2倍のbilinear upsamplingを組み合わせた簡素なもの）を使います。訓練設定の詳細は省きますが、1個のNVIDIA GeForceRTX3090で12時間訓練すれば済む程度だそうです。

性能について

実験結果の詳細までは省きますが、論文では従来手法とFSC147というデータセットで比較しています。指標は密度マップのMAEとRMSEを比較です。模範画像が1個の場合と3個の場合の結果が載っていますが、模範画像が1個の場合でも、従来手法の模範画像が3個の場合を上回っています。また、アブレーション研究なども説得力のある結果です。

また、FSC147から車の画像を除いてモデルを訓練し、車の数え上げデータセットCARPKでの性能も調べています。その結果も先行手法を上回ることが確認されています。このことから、訓練に使用したデータセットに含まれていないようなカテゴリの画像に対しても、有効であることがわかります。

補足

密度マップのGround truthの作り方

この論文の範囲外ですが、データセットの密度マップのGround truthの作り方が気になりましたので、確認しました。まず、クエリ画像のオブジェクトの中心点にアノテーションします。その点を15x15のガウシアンカーネルで畳み込んでいるということでした（keypoint detectionなどで使われる方法と同様だと思います)。分散はどうするかは書かれてはいませんでしたが、おそらく一定なのだと思います。もし一定なのであれば、分散は画像に応じて一意に決まるかと思います。なぜならば、密度マップの合計はオブジェクトの数になっていなければならないからです。出典: [2104.08391] Learning To Count Everything

まとめ

この論文では、ViTベースのシンプルなCAC手法であるCACViTが提案されました。self-attentionにデカップリングの新しい観点を与え、また情報損失を補うために、アスペクト比やマグニチュードを考慮した埋め込みを導入しています。また、論文の結果では、CACViTは他を寄せ付けないSOTAの結果を達成しており、他カテゴリへの汎化性能も示されています。

私が個人的に気になったことが2つあります。1つは、CACViTではクエリ画像内の物体の大きさと、模範画像のオリジナル解像度の大きさがすべて均一であることが仮定されていると思いました。マグニチュードの埋め込みは、物体の大きさがまちまちになるような画像を扱う場合は、あまりこの手法は意味をなさない可能性があると思いました。この点で改善の余地があると思います。（そもそも、マグニチュードの埋め込みの本質的な意味が理解できない部分がありました。アブレーションでは意味があったとなっているものの、アスペクト比を考慮したスケール埋め込みには、すでに大きさの情報が含まれていると思います。より出力に近いところでその情報を注入することで、より意味をなすということなんでしょうかね？）いずれにせよ、実際の問題に適用する際には、このような仮定が含まれていることは意識しないといけないと思いました。もう1つは、コードが公開されておらず、大変残念だと思いました。Abstractには"Code will be available."とあり、近いうちに公開してくれることを期待します。公開されなかったとしても、時間を見つけて再現実装ができればと思います。最近は時間が取れておらずなかなか難しいですが...

AnyTech Engineer Blog

AnyTech Engineer Blogは、AnyTechのエンジニアたちによる調査や成果、Tipsなどを公開するブログです。