SAM3 + Gradio で自動アノテーションツールを作ろう!

はじめに

皆さんはSAM3を触りましたか?
SAM3はテキストプロンプトのみで高精度なセグメントができます。
そして、SAM2に比べてBBOXが扱いやすくなっているので、アノテーション補助にとても相性が良いです。

早速SAM3とGradioを使って簡易的な自動アノテーションツールを作ってみました。
もし自動アノテーション機能の実装を検討してる方がいらっしゃいましたら是非参考にしてみてください。

参考リンク

• https://github.com/facebookresearch/sam3
• https://github.com/edgarGracia/gradio_image_annotator

実装

ソースコード
https://github.com/yama-is-bocchi/auto-annotation-demo

SAM3が使えるようになるまで

他のSAM3関連の記事でもすでに紹介されてると思いますが、重みファイルにアクセスするにはhuggingfaceのsam3のページ でリクエストを承認してもらう必要があります。
SAM3を使うには最低でもこのリポジトリの sam3.ptとgithubリポジトリのBPEファイルが必要です。

環境構築

今回作成したツールはGPU環境前提で、Dockerシステム上で起動します。
なのでnvidia container toolkitを事前に導入してください。

ツール利用手順

1. git clone https://github.com/yama-is-bocchi/auto-annotation-demo.git
2. configsディレクトリにBPEファイル, 重みファイルを配置
3. docker compose up -d

コーディング

実装はすごく単純です。
既存のアノテーションコンポーネントを利用して作成します。

main.py

import colorsys
from typing import Any

import gradio as gr
from PIL import Image
from gradio_image_annotation import image_annotator
from sam3.model.sam3_image_processor import Sam3Processor
from sam3.model_builder import build_sam3_image_model


def make_rainbow_color(i: int, total: int) -> tuple[int, int, int]:
    hue = i / max(total, 1)
    r, g, b = colorsys.hsv_to_rgb(hue, 1.0, 1.0)
    return int(r * 255), int(g * 255), int(b * 255)


def main() -> None:
    model = build_sam3_image_model(
        bpe_path="configs/bpe_simple_vocab_16e6.txt.gz",
        checkpoint_path="configs/sam3.pt",
        load_from_HF=False,
    )  # type: ignore
    processor = Sam3Processor(model=model)  # type: ignore

    with gr.Blocks() as app:
        with gr.Row():
            image_uploader = gr.Image(label="画像をアップロード", type="pil")
            with gr.Column():
                prompt_textarea = gr.TextArea(submit_btn=True, label="プロンプトを入力", placeholder="複数ある場合は改行区切りで入力")
        image_annotation = image_annotator(image_type="pil", label="アノテーションツール", visible=False)

        def submit_sam3_predictor(input_image: Image.Image, prompt: str) -> dict[str, Any]:
            # 画像セット
            inference_state = processor.set_image(input_image)  # type: ignore
            processor.reset_all_prompts(inference_state)
            bboxes: list[dict[str, Any]] = []
            # for:
            labels = [line for line in prompt.splitlines() if line.strip()]
            for i, label in enumerate(labels):
                # 推論
                result = processor.set_text_prompt(state=inference_state, prompt=label)
                boxes = result["boxes"].tolist()
                total = len(labels) * len(boxes)
                for j, (x_min, y_min, x_max, y_max) in enumerate(boxes):
                    color_index = i * len(boxes) + j
                    bboxes.append(
                        {
                            "label": label,
                            "color": make_rainbow_color(color_index, total),
                            "xmin": x_min,
                            "ymin": y_min,
                            "xmax": x_max,
                            "ymax": y_max,
                        },
                    )
            # アウトプット
            return gr.update(value={"image": input_image, "boxes": bboxes}, visible=True)

        prompt_textarea.submit(submit_sam3_predictor, inputs=[image_uploader, prompt_textarea], outputs=image_annotation)

        app.launch()


if __name__ == "__main__":
    main()

1ファイルに収まりました！
推論は数秒で終わります。
VRAMは8GBくらいあれば十分だと思います。

プレビュー

サーバーを起動したら、Web上で自動アノテーションを体験することができます。
試しに私のデバイス周りをアノテーションしようと思います。

こんな感じ複数のプロンプトを改行区切りで入力して推論！

結果　↓

まとめ

自然言語の指示だけで高精度の矩形を自動でアノテーションすることができました！
今回は作図の工程だけ自動化してみましたが、本格的にアノテーション作業をする場合は他にも時間が掛かるタスクがたくさんあります。
今後も自動化できそうなところを見つけて、デモアプリなどで実際に動くものをご紹介していきたいと思います~