SAM2 on Databricks with Databricks Notebooksの翻訳です。
本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
この記事では、Databricks上でMetaの最新のSegment-Anything-Model 2 (SAM2)のスタート方法をウォークスルーします。DatabricksノートブックでSAM2を用いた実験をカバーし、ビジュアライゼーションのためにカスタムのノートブックセルを用いることで、デフォルトのサンプルを拡張します。
コードに直行したいですか？始めるには完全なノートブックをチェックしてください。
SAM2とは？
SAM2は、画像と動画の両方においてオブジェクトをセグメント分けするための統合モデルです。いかなる画像や動画フレームのオブジェクトを選択するために、ポイント、ボックス、マスク座標をサポートしており、動画と画像タスクの両方において、パフォーマンス、精度、手動での介入において、SAMよりも劇的な改善を示しています。こちらで公開されているオリジナルの研究論文で詳細を確認することができます。
なぜセグメンテーション？
様々な業界において、セグメンテーションには数多くのアプリケーションが存在します。いくつかの例では、自動運転自動車の物体検知、疾病診断における医療画像分析、ソーシャルメディアプラットフォームにおけるコンテンツのモデレーションなどがあります。
ホスティングされる数多くのサービスでは、自分たちの機械学習オファリングの一部としてセグメンテーション機能を提供していますが、SAM2のように自己デプロイされるセグメンテーションモデルは、劇的な利点を提供することがあります: セグメンテーションプロセスに対するコントロール、特定ユースケースにモデルをカスタマイズする能力、大規模アプリケーションにおけるコスト削減の可能性などです。また、モデル自身がデプロイされることによって、センシティブな画像や動画を処理のためにサードパーティサービスに送信する必要がないため、優れたデータプライバシーとセキュリティによって、完全な自己ガバナンスを可能にします。
SAM2 on Databricks
スタートするには、GPUがサポートされている稼働中のクラスターが必要となります。ここでは、Databricksランタイムバージョン15.4 LTS MLでg2dn.xlarge [T4] を選択していますが、より大きなワークロードにおいては、より規模の大きいものを選択することも可能です。
次に、Databricks Dev Relサンプルレポジトリから、このノートブックをチェックアウトします。Databricksで直接ノートブックを実行することができますが、以下でキーとなるセクションのいくつかをハイライトします。
セットアップ
DatabricksでSAM2を実行する最も簡単な方法は、SAM2のGithub( https://github.com/facebookresearch/segment-anything-2 )で新規のDatabricks Gitフォルダを作成し、作成したGitフォルダのsegment-anything-2/notebooksにある既存のノートブックフォルダに上述のノートブックをインポートすることです。ノートブックのパスは、ノートブックがリポジトリ内の適切な場所にあることを前提としています。
Githubから直接SAM2ライブラリをインストールし、モデルのバイナリをダウンロードするためにdownload_ckpts.shを使用します。デフォルトでは、このスクリプトはモデルのすべてのサイズをダウンロードしますが、特定のサイズをダウンロードするようにファイルを直接変更することができます。利用できるサイズは、tiny、small、base_plus、largeです。
# install the sam2 pip package directly from the Databricks Git folder
%sh pip install ../../segment-anything-2
# make sure to build extensions
%sh cd ../ && python setup.py build_ext --inplace
# download the model binaries
%sh cd ../checkpoints && ./download_ckpts.sh
好きな動画を使うことができますが、この記事では2024 DATA & AI SummitのDay 1 Keynoteを使います。この動画を前処理し、個別のフレームに分割し、Databricksに直接アップロードする必要があります。この例では、FFmpegとDatabricks Volumeを使いますが、これは要件ではありません。
# use yt-dlp to get the keynote video from youtube
%pip install yt-dlp
%sh yt-dlp -o ./videos/keynote/keynote.mp4 -f "bestvideo[height<=480]" -u "username" -p "password" "https://www.youtube.com/watch?v=-6dt7eJ3cMs"
# use ffmpeg to split the video into frames and place them in a Volume
%sh ffmpeg -ss 00:00:15 -i ./videos/keynote/keynote.mp4 -t 00:00:10 -q:v 2 -start_number 0 /Volumes/sam/default/frames/'%05d.jpg'
セグメンテーション
ノートブックでフレームを直接表示するために、Python I/Oとmatplotlibの組み合わせを活用することができます。これは、MetaのSAM2のサンプルから直接引用したものです。
# `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg`
video_dir = "/Volumes/sam/default/frames"
# scan all the JPEG frame names in this directory
frame_names = [
p for p in os.listdir(video_dir)
if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
]
frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))
# take a look the first video frame
frame_idx = 0
plt.figure(figsize=(12, 8))
plt.title(f"frame {frame_idx}")
plt.imshow(Image.open(os.path.join(video_dir, frame_names[frame_idx])))
最初にpredictorを初期化する必要があります。これは、ffmpegで作成したボリューム内のすべての画像を処理します:
inference_state = redictor.init_state(video_path=video_dir)
セグメンテーションを開始するために、単一のポイントのみを使用することができます。上の画像では、matplotlibが座標システムを提供してくれているので、おおまかに見て、頭が(350,100)にあることが分かります。
セグメンテーションの結果を改善するために、好きなだけポイントを登録したり、さらに大きなポリゴンを登録することもできます。
points = np.array([[350, 100]], dtype=np.float32)
...
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
inference_state=inference_state,
points=points,
...
)
オリジナルの検知結果を残りの動画フレームに伝播させた後で、頭をうまく追従していることを確認することができます。
最後になりますが、フレームをウォークスルーするためのビューにカスタマイズするために、 IPython WidgetsとDatabricks NotebookのdisplayHTML機能を活用することができます。リンクされているノートブックでは、フレームを移動するたびにマウスの座標を表示するビルトインのJSスニペットをデモストレーションしており、これはオブジェクト座標の特定に役立ちます。
まとめ
この記事では、Databricksノートブックの機能を用いてセグメンテーションをスタートするための、DatabricksでのSAM2のサンプルを紹介しました。まさに手放しで利用できるセグメンテーションの能力を活用するために、オリジナルのオブジェクトの座標を特定するための物体検知モデルを組み合わせることで、これをさらに前進させることができます。Huggingface Modelsディレクトリで最新のモデルを探し出すことができます。
実践編で上のノートブックをウォークスルーしています。