🍣

Segment Anything のゼロショット学習による画像全体のセグメンテーション

2024/06/25に公開

こんにちは、HACARUS でインターンをしている朱です。
前回の記事Segment-Anything Model (SAM) の初歩的な使い方について解説しました。本記事ではその応用として、ゼロショットで画像全体のセグメンテーションを実施する方法を紹介します。本記事で紹介する方法は、https://github.com/segments-ai/panoptic-segment-anything で紹介されている手法なので、実装の詳細について知りたい方は、そちらをご覧ください。

ゼロショット学習とは

ゼロショット学習は、タスクの具体例を与えることなく、タスクを実行できるようにモデルを学習させる機械学習の手法です。具体例を与える代わりに、関連する異なるタスクから学習を行います。例えば、物体検出の検出対象を変える際に、新たなデータセットを用意して再学習することなく、検出対象のクラス名(車や信号)のテキストを与え直すだけで検知ができるようになります。この手法は、タスクに対するラベル付きデータが揃っていない場合や、データの入手にコストがかかる場合に有用です。詳しくはZero-shot learningとは?をご覧ください。

SAM で画像全体のセグメンテーションをするには?

SAM は、特定のオブジェクトの情報を点(クリックポイント)やバウンディングボックスのプロンプトとして与えられることを前提に作られています。そのため、画像全体に対してセグメンテーションを実施する場合に SAM を直接適用することは難しいです。

一方で、SAM は高いゼロショット性能を示し、見たことのないオブジェクトに対しても精度の高いセグメンテーションを得ることができます。そこで、別のタスクのゼロショット性能が高いモデル(Grounding DINO や CLIPSeg)で個々のオブジェクトに粗く分割し、 SAM と組み合わせることで画像全体に対する高品質なセグメンテーションを作ることができます。

SAM を使った画像全体のセグメンテーションアルゴリズム

ここでは、 SAM の他にゼロショット性能が高い物体検知モデルである Grounding DINO と、セグメンテーションモデルである CLIPSeg を活用します。

手順

  1. Grounding DINO で物体検出を行い、得られたバウンディングボックスを SAM に入力し、thing mask を作成
  2. CLIPSeg でセグメンテーションを行い、1で得られたマスクと重なっている部分を取り除く
  3. 手順 2 で得られたマスクからサンプリングした点をプロンプトとして SAM に入力し、stuff mask を作成
  4. 手順 1 の thing mask と手順 3 の stuff mask を組み合わせる

以上の手順を図で表すと以下のようになります。

それでは各手順について実行例を用いてみていきましょう。

本記事では、前回に引き続き以下の画像を使用します。

手順1: GroundingDINO + SAM (thing mask の作成)

まず、GroundingDINO を用いて “sushi” の物体検出を行います。

# use dino_detection function to perform object(thing) detection on the image
thing_boxes, thing_category_ids, visualization = dino_detection(
    dino_model,
    image,                  # Image.open("test_img/sushi.jpg").convert("RGB"),
    image_array,            # np.asarray(image),
    thing_category_names,   # ["sushi"]
    category_name_to_id,    # {"sushi": 0},
    dino_box_threshold,     # 0.35,
    dino_text_threshold,    # 0.25,
    device,                 # "cpu",
    visualize=True,
)

画像中の全ての寿司にバウンディングボックスが付与されました。ガリも寿司として検知されていますが、バウンディングボックス左上の confidence score に注目すると、他の寿司と比べてガリの score は 0.41 と小さめに出ていることがわかります。

今回は、画像全体のセグメンテーションが目的なので、これらの全てのバウンディングボックスを使用して thing mask を作成します。6 個のバウンディングボックスをそれぞれプロンプトとして SAM に入力した結果以下のようになりました。

それぞれの寿司(or ガリ)が、かなりいい精度でセグメンテーションされていることがわかります。

手順2: CLIPSeg を使ったセグメンテーション

CLIPSeg を用いてメインの物体(寿司)以外のセグメンテーションを行います。ここでは、寿司以外のオブジェクトである ”wood” , ”table” を検出対象として設定してみます。

# use the clipseg_segmentation function to perform object(stuff) detection
clipseg_preds, clipseg_semantic_inds = clipseg_segmentation(
    clipseg_processor,
    clipseg_model,
    image,                             # Image.open("test_img/sushi.jpg").convert("RGB"),
    stuff_category_names,              # ["wood", "table"],
    segmentation_background_threshold, # 0.1,
    device,                            # "cpu",
)

CLIPSeg の結果、寿司下駄が “wood” として、テーブルが “table” としてセグメンテーションされており、 境界が多少歪んでいるものの、うまくセグメンテーションできていることがわかります。

次に、手順1 で得られた thing mask 領域はすでにセグメンテーション済みなので、重なっている部分を取り除いておきます。

手順3: CLIPSeg + SAM (stuff mask の作成)

手順 2 で得られたセグメンテーションの結果から、それぞれ点をサンプリング(下図, 赤X印)し、プロンプトとして SAM に入力します。テーブルに対する SAM の結果は下図の通りです。CLIPSeg だけだと、”wood” と “table” の輪郭が歪んでいたのに対して、 SAM を活用することで、より綺麗な輪郭を得ることができました。

手順4: thing mask と stuff mask を組み合わせる

最後に手順1 の Grounding Dino + SAM で得られた thing mask と 手順3 の CLIPSeg + SAM で得られた stuff mask を組み合わせた結果を示します。

まとめ

以上のように、ゼロショット性能に優れた複数のモデルを組み合わせることで、半自動(検出対象の “sushi”, “wood”, “table” を指定しただけ)で画像全体に対する高品質なセグメンテーションの結果を得ることができました。

従来の AI では目的のタスクに対して、学習データを集めて再学習することで性能を高めるというアプローチが主流でしたが、強いモデルを複数組み合わせて新たなタスクを解決するというアプローチも現実味を帯びてきているように感じます。モデルの組み合わせ方は膨大なので、どのように AI を組み合わせるのかのノウハウも今後は大事になってくるのかもしれません。

参考文献

HACARUS Tech Blog

Discussion