【Segment Anything】ゼロショットセグメンテーションモデル
3つの要点
✔️ ゼロショットでセグメンテーションを行うことができるモデルを構築
✔️ 1100万枚の画像と11億個のマスクによる構築された大規模なセグメンテーションデータセットを提供
✔️ 実験では新しい画像分布や様々なタスクに対してもゼロショットで性能を発揮
Segment Anything
written by Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, Ross Girshick
(Submitted on 5 Apr 2023)
Comments: Project web-page: this https URL
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)
code:
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
概要
Segment Anything プロジェクトはMetaによる画像セグメンテーションにおける基盤モデルを構築することを目標としてタスク、モデル、データの3つの要素を扱っています。
このプロジェクトでは1100万枚の画像と11億個のマスクによる構築された大規模なセグメンテーションデータセット「SA-1B」を提供しています。またSegment Anything Model (SAM) はプロンプト可能なように設計、学習されており、新しい画像分布やタスクに対してもゼロショットで転移することができます。
後述の多くのタスクの実験を通して、その可能性を評価し、ゼロショット性能は素晴らしいものであるとしています。
タスク
NLPがトークン予測タスクを解くことで基礎モデルの事前学習を行い、プロンプトエンジニアリングを通して様々な下流タスクを解決するということを参考にして、画像セグメンテーションの基礎モデルを構築するために類似の機能を持つタスクをプロンプト可能なセグメンテーションタスクとして提案しています。
プロンプト可能なセグメンテーションタスクの目標はセグメンテーションプロンプトが与えられたら有効なセグメンテーションマスクを返すことです。ここでのプロンプトとは前景、背景の点の組み合わせ、バウンディングボックス、セグメンテーションマスク、自由形式のテキストといった画像内において何をセグメンテーションするかといった情報を指します。
ここで言うセグメンテーションマスクが有効であるための要件とは与えられたプロンプトが図1のように曖昧で、例えば全身、上半身、頭のように複数の対象を示すことができる場合であっても、それらの中から少なくとも1つは妥当なマスクを出力できるということです。
図1 単一の点からSAMにより出力されるマスク
プロンプトを通して既存もしくは新しいセグメンテーションタスクに幅広く対応できるモデルを作ることが、これまでのマルチタスクのシステムとの大きな違いです。マルチタスクシステムはジョイントセグメンテーション、インスタンスセグメンテーション、パノプティックセグメンテーションなど固定された一連のタスクを実行しますがトレーニングとテストタスクは同じです。プロンプト可能なセグメンテーションタスクのために訓練されたモデルが、より大きなシステムの要素として機能することで、推論時に新しい異なるタスクを実行することができます。
モデル
プロンプト可能なセグメンテーションのためのモデルとして図2に示すように画像エンコーダ、プロンプトエンコーダ、マスクデコーダの3つのコンポーネントを持つSegment Anything Model (SAM) を提案しています。
図2 Segment Anything Model の概要
最初の画像エンコーダは MAE (Masked autoencoders) を用いて事前学習したViTを使用しています。この画像エンコーダは画像ごとに1回実行され、モデルにプロンプトを与える前に適用します。
プロンプトエンコーダは点、バウンディングボックス、テキストに当たる疎なプロンプトとセグメンテーションマスクに当たる密なプロンプトを分けて考えます。点とボックスは位置エンコーティングを用いて、テキストにはCLIPのテキストエンコーダを利用しています。マスクは畳み込みを利用して埋め込まれ、画像エンベディングの要素ごとに合計されます。
マスクデコーダでは画像エンべディングとテキストエンべディング、出力トークンからマスクを生成します。画像エンべディングへのプロンプトとその逆の2つの方向でプロンプトセルフアテンションとクロスアテンションを使用して、すべての埋め込みを更新するようにトランスフォーマーデコーダプロックを改変しています。2つのブロックを実行した後、画像エンべディングをアップサンプリングし、MLP (Multilayer perceptron) は各画像場所でのマスクの前景確率を計算します。
出力が1つの場合、曖昧なプロンプトが与えられるとモデルは複数の有効なマスクを平均化してしまうため1つのプロンプトに対して複数のマスクの予測を行います。一般的なケースではマスクは全体、部分、より細かい部分の3つのマスク出力で十分であることが明らかになりました。(図1参照)
データセット
セグメンテーションマスクはインターネット上で豊富には無いので、データ収集のためにデータエンジンを構築しています。データエンジンは画像に対するセグメンテーションマスクのアノテーションのために3つのステージを経ており、それぞれ手動ステージ、半自動ステージ、全自動ステージとなっています。データセットSA-1Bは1100万枚の画像とこの全自動ステージによる11億のセグメンテーションマスクで構成されています。
最初の手動ステージではプロのアノテーターがSAMを搭載したブラウザベースのインタラクティブセグメンテーションツールを利用して、前景および背景のオブジェクトポイントをクリックしてマスクにラベルを付けます。さらにピクセル精度のブラシと消しゴムを使ってマスクをより良くします。
このステージの最初の段階ではSAMは一般的なセグメンテーションデータセットを用いて学習を行い、アノテーションが十分にできた後に、新たにアノテーションされたデータを使用して再学習を行い、モデルを改善していきました。
次の半自動ステージではセグメンテーションを行う能力を向上させるためにマスクの多様性を高めることを目標としています。最初にマスクの自動検出を行い、検出できなかったオブジェクトをアノテーターが追加でアノテーションを行なっていきます。このステージでも最初のステージと同様に定期的な再学習は行っています。
最後のステージではこれまでのステージで十分なマスクを集められている点および曖昧なケースでも有効なマスクを予測できるモデルを開発した点から自動でアノテーションを行なっていきました。32×32のポイントグリッドの1024点に対して、有効なオブジェクトに対応する可能性があるマスクを予測します。この点が部分もしくはより細かい部分の時はモデルは全体、部分、より細かい部分の3つのマスクを出力するようにします。モデルのIoU予測モジュールは確信度が高く、さらに安定したマスクを選ぶために用いる。最終的には NMS (non-maximal suppression) で重複を無くしている。
各ステージで使用した画像とマスクは以下の表の通りです。
表1 各ステージの画像枚数とアノテーション済みマスク数
画像枚数 | マスク数 | |
手動ステージ | 12万枚 | 430万個 |
半自動ステージ | 18万枚 | 590万個 |
全自動ステージ | 1100万枚 | 11億個 |
実験
5つの異なるタスクを用いて実験を行っています。
1点から推測するマスクの評価
多様な画像を持つ23種類のデータセットを用いて1つの前景点からオブジェクトのセグメンテーションを評価します。評価ではRITMと比較を行っています。ほとんどのデータセットのグランドトゥルースは検出可能なマスクを全て列挙していないため予測されたマスクとグランドトゥルースとの全てのIoUの平均であるmIoUからアノテーターがマスクの品質を評価することで補っています。SAMは23種のデータセットのうち16種でRITMを上回っています。
図3 23種類のデータセットにおけるSAMとRITMの比較
エッジ検出
BSDS500データセットを利用して、エッジ検出のタスクでSAMを評価します。16×16のグリッドからSAMで768個のマスクを予測し、NMSにより冗長なマスクを除去します。その後、SobelフィルタからエッジNMSを経てエッジマップを計算します。図4に示すようにSAMはエッジ検出のために学習を行っていないにも関わらず、グランドトゥルースにも無かったエッジ含めて検出できていることがわかります。
図4 ゼロショットエッジ検出の結果(入力画像、グランドトゥルース、SAM)
オブジェクトプロポーザル
物体検出における重要な役割を果たすオブジェクトの提案のタスクでSAMを評価します。自動マスク生成パイプラインを少し変更したバージョンを実行し、マスクを提案としてオブジェクトプロポーザルを出力します。LVISデータセットv1を用いて、平均再現率を計算し、ViTDet-Hとの結果を表2に比較します。ViTDet-HはSAMと異なりLVISで学習しているため全体的に高いパフォーマンスを発揮しています。しかし、SAMも中規模および大規模のオブジェクト、および珍しいオブジェクトではViTDet-Hを上回っています。
表2 LVISデータセットv1に対するオブジェクトプロポーザルの結果
インスタンスセグメンテーション
ViTDetによる物体検出のバウンディングボックスをプロンプトとすることでSAMでインスタンスセグメンテーションを行います。COCOデータセットとLVISデータセットを用いて、ViTDetが予測するマスクと比較を行います。表3よりSAMのマスクAPはViTDetより下回っていることが分かります。しかし図5のようにViTDetと比較して、SAMはより綺麗な境界を持つ高品質のマスクを生産する傾向があります。また図6では人間による評価では一貫してSAMがViTDetを上回っていることが確認できます。
表3 インスタンスセグメンテーションの結果
図5 セマンティックセグメンテーションの検出結果(グランドトゥルース、ViTDet、SAM)
図6 人間によるマスクの評価
テキストからマスク生成
自由形式のテキストからセグメンテーションを行うタスクを検討します。この実験以前では同じSAMを使用していましたが、このタスクでは新しいテキスト注釈を必要としない方法でテキスト認識ができるように学習を行っています。100ピクセルの2乗を超える領域を持つマスクごとにCLIPの画像エンべディングを抽出します。次に学習中に抽出された画像エンベディングを最初のインタラクションとしてSAMにプロンプトを出します。CLIPの画像エンべディングはテキストエンべディングと一致するように学習しているため、SAMの学習は画像エンべディングで行うが推論はテキストエンべディングを使用します。つまり推論時はCLIPのテキストエンコーダを介して、テキストを実行し、結果をテキストをSAMにプロンプトとして埋め込みます。
図7の結果では「ホイール」などの単純なテキストや「ビーバーの歯のグリル」などのフレーズに基づいてセグメンテーションができていることが分かります。またSAMがテキストプロンプトのみから正しくセグメンテーションできない場合でも、追加の点により予測の修正を行うことができます。
図7 テキストからマスクの生成結果
まとめ
画像セグメンテーションに対する飛躍を可能にする新しいタスク(プロンプト可能なセグメンテーション)、モデル(SAM)、データセット(SA-1B)がSegment Anythingプロジェクトの主な貢献です。SAMは多くのタスクをゼロショットで実行できますが、セマンティックセグメンテーションとパノプティックセグメンテーションを実装する簡単なプロンプトを設計する方法は明らかになっていません。
論文中ではSAMが基礎モデルの地位を獲得するかどうかは、それがコミュニティでどのように使われるかによって決まるとしていますが、発表以降様々なツールなどで組み込まれていっていることが確認でき、今後のこの流れは続くと思われます。
この記事に関するカテゴリー