SSAP [Proposal-freeなInstance Segmentation手法] の紹介と実験

こんにちは。先進技術部でアルバイトをしている河田です。
今回はNaiyu Gaoらによって提案されたInstance Segmentationの手法であるSSAP[1]について紹介します。本論文はICCV 2019で採択されました。
SSAPの概要は以下の通りです。

• Instance SegmentationタスクにおいてSingle-shotでProposal-freeな手法を提案し、SOTAを達成
• Semantic Segmentation と、ピクセル間の関係性を表すAffinity pyramidによりInstance Segmenationを実現
• 階層構造を用いることで、５倍のスピードアップと９%のAP精度向上を達成。

PyTorchによる再現実装も公開しているのでご覧ください。

はじめに

現在、Instance SegmentationといえばFaster R-CNN[2]を元にしたMask R-CNN[3]が最も有名です。Mask R-CNNでは、まず各物体を検出してBBoxを作成し(Proposal-based)、検出領域ごとにSegmentationを行うという２段構成で物体領域を見つけるという手法を取ります。この方法では２つのモジュールを通ることになるため、速度が遅いのが欠点です。また、Instance Segmentation の精度がBBoxの検出精度に大きく影響を受けてしまいます。

一方で、SSAPでは物体を検出しBBoxを生成することなく(Proposal-free)、１段構成で(Single-shot) Instance Segementationを行います。具体的には、モデル内部でSemantic Segmentationとピクセル間の関係性を表すAffinityを同時に作成し、この二つを組み合わせます。Semantic SegmentationとAffinityを同時に学習するため、これらの相互関係も踏まえた予測となり、精度向上も期待できます。 さらに、SSAPではInstance Segmentationを作成する際に、粗い解像度で作成したマスクを細かい解像度の特徴量を使って階層的に洗練し識別を行います。階層的なモデルを使用した結果、階層構造を使用しない場合に比べて５倍のスピードアップとAPで９%の精度向上につながったと述べられています。これらの手法を組み合わせることで、難しい都市の景観のデータセット[4]に対してSOTAを達成しました。

画像識別タスクの基礎知識

SSAPはInstance Segmentationを行う手法ですが、そもそもInstance Segmentationとはどのようなタスクなのでしょうか。
画像を入力として与え、それを機械学習で識別するタスクは様々あります。ここでは、簡単に４つの画像識別タスクについてご紹介します。

(1) 画像分類(classification)
与えられた画像が何であるか(クラス)を識別するタスクです。予めクラスの候補が指定されており([犬、猫、車、人…])、その中で最も事後確率が高いクラスを出力します。

(2) 物体検出(detection)
画像分類を行い、さらにその物体の場所も識別するタスクです。多くはBBoxと呼ばれる矩形を用いて各物体の位置を特定します。画像中に複数の物体がある場合は、それら全てについて物体のクラスおよび物体の位置を識別します。

(3) Semantic Segmentation
BBoxを使用せず、ピクセル単位で各クラスの領域を求めるタスクがSegmentationです。複数の物体がある場合でも、同じクラスであれば区別なく同じクラスとして識別を行います。

(4) Instance Segmentation
Semantic Segmentationとは異なり、各クラス領域をピクセル単位で識別し、さらに同じクラスでも別の物体(Instance)であればそれらを区別するタスクです。物体検出とSemantic Segmentationを合わせたタスクで、物体検出よりも精密に物体の形を捉える必要があります。 SSAPでは４つ目のInstance Segmentationのタスクを行います。画像認識タスクの中では細かい部分の識別が必要になるので難しいタスクとされています。

SSAPのモデル構造

「はじめに」で述べたとおり、SSAPではInstance Segmentationのタスクを行うために、モデル内でSemantic SegmentationとAffinityを作成します。Semantic Segmentationはピクセル毎のクラスを推測し、Affinityは近くのピクセルが同じ物体(Instance)に属しているかどうかを推測します。その後、この二つの推測結果を階層的なグラフ分割により合成することで、画像の中のInstanceを識別します。Affinityと階層的なグラフ分割について、それぞれ詳しく見ていきましょう。

Affinity

Affinityは、あるピクセルとその周囲の r×r の範囲にあるピクセルが同じInstanceに属する確率を示します。
Affinityは０〜１の値を取り、値が1に近いほど同じInstanceに属している確率が高いことを表します。 この中心のピクセルを１ピクセルずつ動かし、画像中のすべてのピクセルについてAffinityを計算するので、 画像の大きさが h×w だとすると、全部で r^2×h×w の大きさのAffinity mapが作成されます。 学習中は正解Affinityと推測したAffinityの二乗誤差を最小化します。 なお、aは推測したAffinityの値、yは正解のAffinityの値(ground truth)を示します。a^jはある中心ピクセルに対するj番目のピクセルのAffinityで、a^j\in(0, 1) です。同様にy^jはj番目の正解データで、中心のピクセルと同じInstanceに所属している場合は１、それ以外は０を取ります。

しかし、このlossをそのまま最小化してしまうと、バランスが非常に悪くなってしまいます。なぜなら、Affinityが０となるのは物体と物体の境目付近のみであり、ほとんどのAffinityの値は１だからです。そのため、学習中は正解が１、つまり同じ物体に属しているピクセルのAffinityを80%の確率でランダムにdropさせ、Lossを計算しないようにします。さらに、通常は背景の領域に比べ物体の領域が狭いため、いずれかのInstanceに属するピクセルについては３倍のLossを設定することで物体に注目させます。これらの工夫をすることでバランスをとります。

階層的なAffinity (Affinity Pyramid)

先ほど周りのr×rのピクセルについてAffinityを作成すると述べましたが、このままではあるピクセルに対し、非常に近くのピクセルしか推測を行えません。そのため、Affinityを計算しない遠くのピクセルについては同じInstanceに属しているのかどうかがわかりません。また、大きな物体についても同様の理由で識別が難しくなってしまいます。この問題の解決策としては、例えばrの値を大きくしたり、画像の解像度を小さくしたりすることが考えられます。しかし、rの値を大きくして遠くのピクセルまでAffinityを取ろうとするとメモリが大量に必要になると同時に、Semantic Segmentationにも悪影響を及ぼすことが論文内で示されています。また、画像の解像度を小さくすると長距離のAffinityを考慮することができる一方で、細かいピクセル単位での識別が困難になり、小さな物体を検出しなくなってしまいます。

そこで、SSAPでは複数の解像度でAffinityを作成する(Affinity Pyramid)ことで、遠近両方のAffinityを取ることを提案しています。先ほど述べたように、解像度を下げることでAffinityを取る範囲rを変えずに長距離のAffinityを考慮できます。例えば、解像度が\frac{1}{64}のAffinityで２ピクセル離れている場合、元の解像度では128ピクセル分の距離のAffinityを意味します。小さい物体は低解像度で発見することができませんが、高解像度で発見が可能です。複数の解像度でAffinityを計算することで、小さい物体も大きい物体も識別することができるようになるわけです。 さて、上で述べたような複数の解像度のAffinityを作成するためにはどのようなモデルを用いればよいでしょうか。SSAPでは、一般的にU-Net[5]と呼ばれる”U”の形をしたモデルを使用します。このモデルでは一度解像度を落とした後(Encoder)、再度解像度を上げる(Decoder)ことで細かい部分と大局的な部分の両方に注目できる構造になっています。また、解像度を上げる際に、解像度を落とす前の特徴量を足し合わせることにより、細かい部分の情報が消えてしまわないように工夫されています。

このモデルにより各解像度でAffinityを算出する際に、同時にSemantic SegmentationもAffinityと同じ解像度で算出します。AffinityとSemantic Segmentationを同時に学習することで相互関係を踏まえた予測が行えるようになり、精度が向上します。なお、論文では元画像の大きさの [\frac{1}{4}, \frac{1}{8}, \frac{1}{16}, \frac{1}{32}, \frac{1}{64}] の５つの解像度でAffinityとSemantic Segmentationを計算しています。

階層的なグラフ分割(Cascaded Graph Partition)

さて、Semantic SegmentationとAffinityの値は学習を進めることで取得できることが分かりました。それでは、この二つを使ってどのようにInstance Segmentationを作成していくのでしょうか。SSAPでは、「グラフ分割」の考え方に基づいてピクセルをInstance毎に分割する手法を取ります。
そもそもグラフは、頂点を意味する「ノード」と、ノード間をつなぐ「エッジ」により構成されています。グラフ分割では、これを何らかの条件でいくつかの部分グラフに分割します。ここで、ピクセルをノードとみなし、周囲 r×r との間にエッジが張られていると考えると、ピクセル全体はグラフを構成していると言えます。このグラフを、Affinityを考慮して分割することでInstanceを生成していきます。
グラフ分割について、もう少し細かく解説します。
通常のグラフ分割問題は、以下の式(3)を式(4)の制約の下で最適化する問題です。 w_eはエッジeの重みです。y_eは０か１のどちらかを取り、 y_e=1 はエッジeを切断することを、y_e=0 はエッジeを切断しないことを意味します。つまり、w_eが負の時に切断 (y_e=1) し、w_eが正の時に切断しない (y_e=0) ようにすることで、式(4)は最小になります。

では、重みが正のエッジは全て残し、重みが負のエッジは全て切れば良いかというと、実はそうではありません。例えば下図のように重みwが与えられた場合、w_{2, 3}=-0.9のエッジのみを切断してグラフを分割しようとしても、他のエッジがつながっているためグラフを分割することができません。この問題を回避するため、うまくグラフが分割される制約(式(4))が必要となります。
Cはサークル、つまりグラフのある閉路を示しており、e'はCに含まれるエッジの一つを示します。また、eはCに含まれるe'以外の全てのエッジを意味します。つまり、式(4)は、Cに含まれる一つのエッジe'を選んだ時に、y_eの合計よりもy_{e'}が大きくなってはいけないことを意味しています。この制約は簡単に言うと、「閉路上のエッジを切断するなら、必ず二箇所以上で切らなければならない」ということです。例えば輪ゴムを切って２つに分けることを想像すると、必ず２箇所は切らなければならないということがわかるかと思います。

具体例を見てみましょう。例えば下図でノード[２]とノード[３]の間のエッジをe'に選ぶとします。この時、図８のようにエッジを切断することを考えると、y_{e'}=1です。また、その他のエッジは \sum y_e=0+0+0=0 です。つまり、\sum y_e<y_{e'} となってしまい、式(4)の制約が満たされないことになります。

以上より、基本となるグラフ分割問題は、制約である式(4)を満たしつつ、式(3)を最適化する問題だと言えます。 さて、グラフ分割について理解できたところで、グラフ分割をInstance Segmentationに応用する方法を見ていきましょう。まず無向グラフG=(V, E)を考えます。Vはピクセルのセット、E\subseteq V^2はAffinityを計算する範囲内rにあるペアのエッジを示します。e_{u, v}\in Eはピクセル(u, v)の間のエッジを、a_{u, v}, a_{v, u}\in (0, 1)はピクセル(u, v)間のAffinityを示します。両方向からAffinityを計算するため、１つのエッジで２つのAffinityが計算されています。この２つのAffinityは式(5)により平均され、式(6)によりエッジの重みw_{u, v}を計算します。 式(6)を図にすると以下のようになり、Affinityの平均が0.5より大きければエッジの重みは正、小さければ負に変換されることがわかります。 計算したエッジの重みが正であればAffinityが大きく、二つのピクセルは同じInstanceに属している可能性が高いことを意味するのでエッジはそのまま残します。逆にエッジの重みが負であればAffinityが小さいため、エッジを切断します。このようにすることで、グラフ分割をInstance Segmentationに応用できるようになります。

しかし、グラフ分割は基本的に繰り返し処理が入るため、ノードが多くなるにつれ時間も長くかかります。ピクセルごとにノードを作成するとノード数が非常に大きくなってしまうため、実際のアプリケーションにグラフ分割をそのまま使うのは難しいという難点があります。また、大きな物体に関してはほとんどのピクセルが物体の内部に位置するため、簡単な問題であるにも関わらず非常に時間がかかってしまいます。

※具体的にグラフ分割を解く方法については長くなってしまうため末尾に掲載します(具体的なグラフ分割の解き方(GAEC法))。

階層構造

そこで、論文では先ほど解像度ごとに作成したSemantic SegmentationとAffinityを使って階層的にグラフ分割を行い、実行時間の短縮を図ることを提案しています。いきなり大きな解像度でグラフ分割をするとノード数が多くなり、非常に時間がかかってしまいます。しかし、解像度が低いSemantic SegmentationとAffinityを用いてグラフ分割を行い、その結果をもとに徐々に解像度の高いグラフを分割していくことで、速度の向上が期待できます。解像度が低いグラフ分割結果でも、識別したInstanceの内側のピクセルの部分は信頼できるため、内側の部分をup-samplingし、これを一つのノードとして次の高解像度のグラフ分割の際にエッジと共に渡していきます。このように階層的にグラフ分割を繰り返すことで、Instance segmentationの結果が徐々に洗練されていきます。さらに、ノードが増えるのも防げるため、推測速度が早くなります。

Segmentation Refinement

ここまでグラフ分割について述べてきましたが、仮にSemantic Segmentationの結果毎にグラフ分割を行うことを考えたとき、Semantic Segmentationのクラス識別が間違えてしまうと、グラフ分割がどれだけ正しくてもInstanceの作成に失敗してしまいます。例えば下図の(b)のようにトラックのクラス識別が３つのクラスに分かれてしまうと、Instance結果も(c)のように、多くのInstanceに分かれてしまいます。 そのため論文では、Segmentation Refinement(SR)という手法を提案しています。SRではクラス毎にはグラフ分割を行わず、背景ではない、いずれかのインスタンスに属すると判定された全てのピクセルに対してグラフ分割を行います。その後、作成したInstanceがどのクラスに属するのかを判定するという方法を取ることで、Semantic Segmentationの誤りがInstance Segmentationの誤りに直結しないようにしています。
さらに、Segmentation Refinementでは「同じ物体に属するピクセルのクラスの事後分布は似ている」という考え方を基に、Semantic Segmentationのクラス判別結果を用いてAffinityを修正することも提案しています。これらの処理を行うことで図の(d)が示すように、Segmentationを間違えてもInstanceの識別に成功する可能性が高くなります。

Affinityの修正は以下の式(７)〜式(９)で行われます。 \alpha_{u, v}'は修正したピクセル(u, v)間のAffinityを表しており、D_{JS}はJensen-Shannon divergenceを意味します。 S_uはピクセルuの各クラスの事後確率を示しています。これより、exp[-D_{JS}(S_u||S_v)] は２つのピクセル間におけるクラスの事後分布の隔たりを計算していることになります。クラスの事後確率の隔たりが大きいほど exp[-D_{JS}(S_u||S_v)] の値は小さくなり、結果としてAffinityを小さくします。 exp[-D_{JS}(S_u||S_v)] は０から１の値をとります。