AI
AWS

MetaのSAM2を試してみた。

Varunのアバター
Varun

背景や目的

・以前のプロジェクトで、CV（コンピュータビジョン）とYOLOを用いて物体検出を試みました。その時、解析に必要なセグメンテーションに興味を持ちました。
・2ヶ月前にリリースされたMetaのSAM2（Segment Anything Model 2）は、画像や動画内の任意の物体を輝度差では無くモデルでセグメントできるようなので、実際に試してみたくなりました。
・実際のプロダクトとして使えるかを確認するために、ゴルフの打球で試してみました。　

実験環境構築

まずは、実験環境を構築します。
今回はGoogle Colabo上で構築して実験します。
（Google colaboの基本的な説明はここでは割愛します。）
・ColaboのNotebookで今回はGPUランタイムを使用するため、右上にある「Connect」の右側にある「▼」をクリックし、「Change runtime type」を選択します。すると、画面中央にモーダルウィンドウが表示されます。

・そのモーダルウィンドウで「Hardware accelerator」を「T4 GPU」に変更し、保存します。

・コードセルでは、GitHubのリポジトリからSAM2のパッケージをクローンしてインストールします。

Google Colab
# パッケージのインストール
!git clone https://github.com/facebookresearch/segment-anything-2.git
%cd ./segment-anything-2
!pip install -e .

※セッションの再起動エラーが表示された場合は、モーダル画面の「Restart session」を押し、ランタイムを再起動し、「Shift+Enter」を押してコードを再度実行してください。

・次はモデルチェックポイントをダウンロードします。
・モデルチェックポイントとは、SAM2のトレーニングプロセス中に保存された、事前にトレーニングされたモデルの重みやパラメータを指します。「Shift+Enter」で実行します。　

Google Colab
# モデルチェックポイントのダウンロード
%cd checkpoints
!./download_ckpts.sh
%cd ..

・GoogleDriveもマウントしておきます。

Google Colab
from google.colab import drive
drive.mount('/content/drive')

それでは、動画セグメンテーションの準備に取り掛かります。

実験１：ゴルフの軌跡（ティーショット）

・ゴルフのティーショットを追跡できるか確認します。

動画１：ティーショット元動画

（※速くて見えない場合は、再生後に右下の「設定メニュー」で速度を落としてみて下さい。）

SAM2を使用した動画セグメンテーションの手順は次の通りです。
１．フレームに分割 ：動画を個別のJPEGフレームに変換。
２．初期化と設定 ：マスクと点群の表示、Predictorの準備、推論状態の初期化。
３．目標指定 ：追跡対象のオブジェクトを指定。
４．予測処理 ：各フレームで対象を予測。
５．動画で保存 : 予測結果付き各静止画をを動画に変換し保存。

上記の手順通りで動画セグメンテーションをしましょう。

１．フレームに分割 ：

・まずは、以下のコードを使って動画をJPEGフレームのリストに変換します。動画のパスを設定して、コードセルを「Shift+Enter」で実行します。

Google Colab
# 動画をJPEGフレームのリストに変換
!mkdir videos
!ffmpeg -i /content/drive/...../SampleData/tee_trim.mp4 -q:v 2 -start_number 0 ./videos/'%05d.jpg'

・次は、作成したJPEGファイルを集め、ファイル名を数値として解釈してソートしたフレーム名のリストを作成します。

Google Colab
import os

# JPEGフレーム名の準備
video_dir = "./videos"
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]))

　

２．初期化と設定 ：

・動画セグメンテーション予測（=predictor）に関連するコードを準備します。「shift+enter」で実行されます。

Google Colab
from sam2.build_sam import build_sam2_video_predictor

# Predictorの準備
sam2_checkpoint = "/content/segment-anything-2/checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)

・動画の推論状態を初期化します。「shift+enter」で「init_state」が実行され、video_path内のすべてのJPEGフレームを対象に、inference_state(Dict型変数)が作成される。これはこの後の「予測」で使われます。

Google Colab
# 動画の推論状態の初期化
inference_state = predictor.init_state(video_path=video_dir)

　

３．目標指定 ：

・次に、動画内の目的となるゴルフボールの座標「input_point」の（x,y）を指定します。
・座標設定の確認のため、確認用フレーム（＝フレーム0）を使用して、物体設定位置（赤色の丸の位置）を確認します。
・この際、ボールの座標は手動で取得し、試行錯誤を繰り返して正確な座標を確認し、完璧な座標に調整します。「Shift+Enter」で実行されます。

Google Colab
import numpy as np

# 目的のオブジェクトの指定
ann_frame_idx = 0
ann_obj_id = 1
input_point = np.array([[820, 1325]], dtype=np.float32) # 位置
input_label = np.array([1], np.int32) # ラベル (1:前景点、0:背景点)

# マスクの表示
def show_mask(mask, ax, obj_id=None, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        cmap = plt.get_cmap("tab10")
        cmap_idx = 0 if obj_id is None else obj_id
        color = np.array([*cmap(cmap_idx)[:3], 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

# 点群の表示
def show_points(coords, labels, ax, marker_size=50):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='o', s=marker_size, edgecolor='red', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='o', s=marker_size, edgecolor='white', linewidth=1.25)
Google Colab
import cv2
import matplotlib.pyplot as plt

# 指定フレームをPredictorで予測
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=input_point,
    labels=input_label,
)

# 確認
plt.figure(figsize=(12, 8))
plt.title(f"frame {ann_frame_idx}")
image = cv2.imread(video_dir + "/" + frame_names[ann_frame_idx])
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(image)
show_points(input_point, input_label, plt.gca())
show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0])
設定前フレーム（Before）設定後フレーム（After）

・ちゃんとボールにマークされたので大丈夫ですね。

　

４．フレーム予測 ：

・ゴルフボールの座標確認はOKなので、全フレームを予測します。「shift+enter」で実行されます。

Google Colab
# 全フレームをPredictorで予測
video_segments = {}
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
    video_segments[out_frame_idx] = {
        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
        for i, out_obj_id in enumerate(out_obj_ids)
    }

・全ての結果フレームをプロットして「output」ディレクトリに保存します。「Shift+Enter」で実行されます。

Google Colab
!mkdir output

# 結果画像フレームの保存
plt.close("all")
for out_frame_idx in range(len(frame_names)):
    plt.figure(figsize=(6, 4))
    plt.title(f"frame {out_frame_idx}")
    plt.axis("off")
    plt.tight_layout(pad=0)

    # 元画像の描画
    image = cv2.imread(video_dir + "/" + frame_names[out_frame_idx])
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    plt.imshow(image)

    # マスクの描画
    for out_obj_id, out_mask in video_segments[out_frame_idx].items():
        show_mask(out_mask, plt.gca(), obj_id=out_obj_id)

    # 結果画像フレームの保存
    basename = os.path.basename(frame_names[out_frame_idx])
    output_frame = os.path.join("output", basename)
    plt.savefig(output_frame)
    plt.close()

　

５．動画で保存 :

・「output」ディレクトリーで保存した結果フレーム画像を動画に変換します。「shift+enter」で実行されます。

Google Colab
# 結果画像フレームを動画に変換
!ffmpeg -framerate 30 -pattern_type glob -i './output/*.jpg' -c:v libx264 -r 30 -pix_fmt yuv420p /content/drive/MyDrive/output.mp4

・完了したら、predictor（予測）をリセットします。「shift+enter」で実行されます。

Google Colab
# 動画の推論状態のリセット
predictor.reset_state(inference_state)

ティーショット動画セグメンテーション結果：
・結果動画は以下の通り物体ハイライト動画です。
　動画内の物体はオレンジ色でハイライトされています。

　　　　　　　　　　　　　動画２：動画１の物体ハイライト動画

・結果はNGでした。最初の数フレームは適切にセグメント化されていますが、ボールが小さいため、ほとんど見えなくなります。その結果、35フレーム目からはゴルフボールではなくゴルフクラブがセグメント化され始めます。

（※速くて見えない場合は、再生後に右下の「設定メニュー」で速度を落としてみて下さい。）

・結果の動画が見難いので、SAM2デモのウェブサイトで視覚効果オーバーレイを使って対策しました。

動画3：視覚効果オーバーレイ

・ゴルフボールのスピードが少し遅い、アプローチショットのゴルフ動画を同じ手順を実行します。

実験２：ゴルフの軌跡（アプローチ）

・アプローチは以下の動画を使いました。

動画4：アプローチ元動画

・実験1と同様に、予測前にゴルフボールの座標を設定します。

設定前フレーム（Before）設定後フレーム（After）

・全フレームを予測し、保存した結果フレーム画像を動画に変換します。

アプローチ動画セグメンテーション結果：

　　　　　　　　　　　　　動画5：動画4の物体ハイライト動画

・実験１と同様に、視覚効果オーバーレイ結果動画を作成しました。

動画6：視覚効果オーバーレイ

・結果はまあまあでした。ボールは2回目のバウンド以降は認識されなくなるようです。主な原因は、ボールが小さすぎて見えなくなるためだと思われます。
・対策として、ビデオの前処理（特定の領域の拡大など）が有効かもしれません。

まとめ

・ティーショットでは、ボールが小さく、検出が難しいため、製品化には課題が残ります。SAMは現段階では、私たちのニーズに完全に合う形ではないかもしれません。
・アプローチショットの場合、ボールの位置を指定するだけで追跡ができました。

SAMのさらなる改善に注目していきます。

ご連絡フォーム

フィードバックを是非お願いします。
本記事の方法での問題点や、よりよい方法のアイデアを頂けると大変助かります。

・この記事に関して

・その他のご連絡

この記事をシェアする

DevAIsをもっと見る

今すぐ購読し、続きを読んで、すべてのアーカイブにアクセスしましょう。

続きを読む