🔰

[Tips] Segment Anything Model 2 (SAM 2)を使うときに悩んだ細かい点: 動画の順方向と逆方向への処理

2024/08/14に公開

Python

機械学習

セグメンテーション

tech

背景

Metaが発表したセグメンテーションが注目されています。

実際に環境を構築すると、Google Colabなどの無料GPUでもギリギリ処理ができます (1画像1秒あたり)。CPUのみを使ったFork版を使った場合でも同じぐらいの処理速度になります。

私が確認できる範囲で確認したところ、RTX3090だと10-12画像/秒、RTX4090だと15-17画像/s、またRTA 6000 Adaだと25画像/秒ぐらいで動作する様子でした。

動画が1秒あたり24-25枚ぐらいのFPSだとすると、だいたい1分の動画を処理するのに60秒~120秒ぐらいということになります。なかなか有望なモデルですね。

補足: 環境設定ついて

適当なGPU環境かGoogle Colabなどで環境を作りました。このあたりを見ていただければ動くと思います（インストールに時間が多少かかります）。

処理の概要: トラッキングの処理について

SAM 2のモデルが画像の連番でデータを処理します。Google Colab該当する部分を見てみます。

入力動画の切り出し

入力動画 (basketball-1.mp4) の 100-300 の区間を画像に切り出し、00000.jpeg のような名前を付けます

前処理

frames_generator = sv.get_video_frames_generator(SOURCE_VIDEO, start=100, end=300)
images_sink = sv.ImageSink(
    target_dir_path=SOURCE_FRAMES.as_posix(),
    overwrite=True,
    image_name_pattern="{:05d}.jpeg"
)

with images_sink:
    for frame in frames_generator:
        frame = sv.scale_image(frame, SCALE_FACTOR)
        images_sink.save_image(frame)

生成された画像は次のようになります。

> tree basketball-1 
basketball-1
├── 00000.jpeg
├── 00001.jpeg
├── 00002.jpeg
├── 00003.jpeg
├── 00004.jpeg
├── 00005.jpeg
├── 00006.jpeg

オブジェクトのトラッキング

サンプルのコードでは00000.jpegを使って、オブジェクトを登録し、後に動画の各画像を処理します。ここでは手前の人を検地してみます。

オブジェクト

SOURCE_IMAGE = SOURCE_FRAMES / "00000.jpeg"
inference_state = sam2_model.init_state(video_path=SOURCE_FRAMES.as_posix())

# 手前の人 (背番号2の方)
default_box = [{'x': 550, 'y': 290, 'width': 0, 'height': 0, 'label': ''}]

boxes = widget.bboxes if widget.bboxes else default_box
points = np.array([
    [
        box['x'],
        box['y']
    ] for box in boxes
], dtype=np.float32)
labels = np.ones(len(points))

FRAME_IDX = 0
OBJECT_ID = 1

_, object_ids, mask_logits = sam2_model.add_new_points(
    inference_state=inference_state,
    frame_idx=FRAME_IDX,
    obj_id=OBJECT_ID,
    points=points,
    labels=labels,
)

ここまで処理すると、この人をトラッキングするんだぞ！という気持ちになります。

実際、00000.jpegに出てくる手前の人を、動画の後ろまでトラッキングし、色をつけて描画する処理がサンプルコードが実装しているものです。

塗りつぶした動画の生成する処理

後半の処理

# 入力動画の情報をコピー
video_info = sv.VideoInfo.from_video_path(SOURCE_VIDEO)
video_info.width = int(video_info.width * SCALE_FACTOR)
video_info.height = int(video_info.height * SCALE_FACTOR)

# 順番に処理し、1枚ずつ書き込む
frame_paths = sorted(sv.list_files_with_extensions(SOURCE_FRAMES.as_posix(), extensions=["jpeg"]))
with sv.VideoSink(TARGET_VIDEO.as_posix(), video_info=video_info) as sink:
    for frame_idx, object_ids, mask_logits in model.propagate_in_video(inference_state):
        frame_path = frame_paths[frame_idx]
        frame = cv2.imread(frame_path)
        masks = (mask_logits[0] > 0.0).cpu().numpy()

        detections = sv.Detections(
            xyxy=sv.mask_to_xyxy(masks=masks),
            mask=masks.astype(bool)
        )
        annotated_frame = mask_annotator.annotate(scene=frame.copy(), detections=detections)
        sink.write_frame(annotated_frame)

生成されるものです。

悩んだこと: 動画の途中 (00000.jpeg以外) のオブジェクトを扱うにはどうするのか

サンプルだと00000.jpegを使っていたのですが、実際処理したいときは必ずしも00000.jpegにオブジェクトがあるわけではありません。サンプルコードが

[0] [1] [2] [3] [4] ...
 |
 |--------------→ 順方向の処理 (= 出力動画)

というような順方向の処理だけを実装していましたが、実際は

[0] [1] [2] [3] [4] ...
             |
             |--------------→ 順方向の処理 (= 前側の動画)
  ←----------|                逆方向の処理 (= 後側の動画、逆再生)

というような処理をして、2つを結合して1つの動画を作りたいです。

やったこと

動画処理に使っている model.propagate_in_video(inference_state) の部分ですが、実は reverse=True という引数を渡すことができます。

これを使って、両方向を処理してから、作成された（マスクされた）画像列を処理したらよいです。すごく雑に処理すると、次のようにやります。

雑な実装

# F (→)
for frame_idx, object_ids, mask_logits in model.propagate_in_video(inference_state):
    frame_path = frame_paths[frame_idx]
    frame = cv2.imread(frame_path)
    masks = (mask_logits[0] > 0.0).cpu().numpy()

    detections = sv.Detections(
        xyxy=sv.mask_to_xyxy(masks=masks),
        mask=masks.astype(bool)
    )

    annotated_frame = mask_annotator.annotate(scene=frame.copy(), detections=detections)
    f_frames.append(annotated_frame)

# B (←)
for frame_idx, object_ids, mask_logits in model.propagate_in_video(inference_state, reverse=True):
    frame_path = frame_paths[frame_idx]
    frame = cv2.imread(frame_path)
    masks = (mask_logits[0] > 0.0).cpu().numpy()

    detections = sv.Detections(
        xyxy=sv.mask_to_xyxy(masks=masks),
        mask=masks.astype(bool)
    )

    annotated_frame = mask_annotator.annotate(scene=frame.copy(), detections=detections)
    b_frames.append(annotated_frame)

# 全体動画は、スタートのフレームが重複するのと、B (←) が逆再生になるので直す
frames = b_frames[::-1] + f_frames[1:]

# 保存する
with sv.VideoSink(TARGET_VIDEO.as_posix(), video_info=video_info) as sink:
    for f in frames:
        sink.write_frame(f)

結果

実際に処理をしてみます。50枚目の画像 (00050.jpeg) について

default_box = [{'x': 592, 'y': 349, 'width': 11, 'height': 14, 'label': ''}]

部分にあるボールを対象とし、両方向にpropagateした上で、1つの動画にしてみます。このような中間の場面のことです。

実装はいい感じにサンプルを直してください（後でgistなどに置きます）。このような処理をすると、100-300の範囲で200枚あった画像が、150枚と51枚 (1枚重複、おそらく) で処理されるようなプログレスバーが見て取れます（こちらはCPUで動かしたものなので、処理時間は遅いです）。

frame loading (JPEG): 100%|█| 200/200 [00:02<00:00, 68.12it/s
propagate in video: 100%|██| 150/150 [05:35<00:00,  2.24s/it]
propagate in video:  22%|▊   | 11/51 [00:23<01:29,  2.24s/it]

ここではTARGET_VIDEO_FとTARGET_VIDEO_Bを適当に定義し、順方向の処理・逆方向の処理・くっつけたやつ、をそれぞれGIFアニメにしました。

順方向 (150)	逆方向 (51,逆再生)	結合後

まとめ

なんとなく、やりたいことができた気がしています（このような順方向と逆方向にスキャンする様は、SAM 2のデモサイトでも見ることができます）。

どなたか悩んだ人の参考になれば幸いです＆もっといいやり方があれば教えて下さい。