WSL2でSAMURAIを試してみる

2024年11月23日 12:13

「視覚的物体追跡用に特別に設計されたSAM 2の強化版である」らしいSAMURAIを試してみます。
SAMURAIといえばSPIRITS。モズ落とし。話がそれました。

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

環境構築

venv環境構築。

python3 -m venv samurai
cd $_
source bin/activate

リポジトリをクローンします。

git clone https://github.com/yangchris11/samurai
cd samurai

パッケージをインストールします。

# sam2 強化版
cd sam2
pip install -e .
pip install -e ".[notebooks]"

# READMEより
pip install matplotlib==3.7 tikzplotlib jpeg4py opencv-python lmdb pandas scipy
# 追加分
pip install loguru

モデルのダウンロード

SAM2.1のモデルをダウンロードしておきます。

cd checkpoints && \
./download_ckpts.sh && \
cd ..

データセットのダウンロード

LaSOT形式のデータセットを用いて推論するようなので、Hugging Faceで公開されているデータをダウンロードしましょう。

https://huggingface.co/datasets/l-lt/LaSOT

とはいっても、これを全部ダウンロードすると容量がアレなので、今回は airplane.zipだけを対象にします。

mkdir ./data/LaSOT
cd $_

# download
wget https://huggingface.co/datasets/l-lt/LaSOT/resolve/main/airplane.zip

# extract
unzip -d airplane airplane.zip
cd ../../

このあと実行する scripts/main_inference.pyでは、./data/LaSOT/testing_set.txt というテスト対象のデータのディレクトリを示したファイルが必要になるので、testing_set.txt の記述を参考に

cat <<EOF > ./data/LaSOT/testing_set.txt
airplane-1
airplane-9
airplane-13
airplane-15
EOF

として4つのデータセットを指定したファイルを作ります。

2. 試してみる

(1) scripts/main_inference.py

では試してみましょう。

CUDA_VISIBLE_DEVICES=0 python scripts/main_inference.py

こんな感じで動いています。

Running video [2/4]: airplane-13
 with 2971 frames
SAMURAI mode: True
frame loading (JPEG): 100%|████████████████████████████████████████████████████████████| 2971/2971 [01:12<00:00, 40.85it/s]
propagate in video: 100%|██████████████████████████████████████████████████████████████| 2971/2971 [04:08<00:00, 11.97it/s]

CPUパワーに依りますが、3,000フレームで4分ほどです。

以下のファイルが生成されます。

(1) samuraiによる物体の追跡結果（座標データ）

$ ls -al results/samurai/samurai_base_plus/
total 340
drwxr-xr-x 2 user user   4096 Nov 23 11:06 .
drwxr-xr-x 3 user user   4096 Nov 22 23:40 ..
-rw-r--r-- 1 user user  41069 Nov 22 23:40 airplane-1.txt
-rw-r--r-- 1 user user  42184 Nov 22 23:44 airplane-13.txt
-rw-r--r-- 1 user user  62112 Nov 22 23:50 airplane-15.txt
-rw-r--r-- 1 user user 116565 Nov 23 00:01 airplane-9.txt
$

(2) LaSOTデータセットの座標（緑色）とsamuraiによる座標（青色）を長方形で示して合成された動画

$ ls -al visualization/samurai/base_plus/
total 150380
drwxr-xr-x 2 user user     4096 Nov 23 11:01 .
drwxr-xr-x 3 user user     4096 Nov 22 23:36 ..
-rw-r--r-- 1 user user 30559133 Nov 22 23:40 airplane-1.mp4
-rw-r--r-- 1 user user 21747084 Nov 22 23:44 airplane-13.mp4
-rw-r--r-- 1 user user  7969451 Nov 22 23:50 airplane-15.mp4
-rw-r--r-- 1 user user 72525661 Nov 23 00:01 airplane-9.mp4
$

(2) scripts/main_inference.py改

どうやって追跡しているのかしら？と思ってsamurai/scripts/main_inference.py を読んだところ、groundtrush.txtの１行を初期位置として読み込んで、追跡しているよう。

prompts = load_lasot_gt(osp.join(video_folder, cat_name, video.strip(), "groundtruth.txt"))

bbox, track_label = prompts[0]

prompts[0]のデータは、以下のtxtファイルの1行目に該当します。

$ head ./data/LaSOT/airplane/airplane-15/groundtruth.txt
197,216,106,39
183,211,106,39
173,195,105,39
162,179,104,40
145,170,103,41
116,163,101,45
85,162,106,45
75,169,103,45
69,177,105,44
68,181,102,44
(snip)

ということで。
以下のように、比較のためだけに座標データを取得している処理をコメントアウトし、

diff --git a/scripts/main_inference.py b/scripts/main_inference.py
index 542b5da..f13ed14 100644
--- a/scripts/main_inference.py
+++ b/scripts/main_inference.py
@@ -113,8 +113,8 @@ for vid, video in enumerate(test_videos):
                 for obj_id in bbox_to_vis.keys():
                     cv2.rectangle(img, (bbox_to_vis[obj_id][0], bbox_to_vis[obj_id][1]), (bbox_to_vis[obj_id][0]+bbox_to_vis[obj_id][2], bbox_to_vis[obj_id][1]+bbox_to_vis[obj_id][3]), color[(obj_id)%len(color)], 2)

-                x1, y1, x2, y2 = prompts[frame_idx][0]
-                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
+                #x1, y1, x2, y2 = prompts[frame_idx][0]
+                #cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                 out.write(img)

             predictions.append(bbox_to_vis)

groundtrush.txtを１行だけにしたファイルを作成しても動くよね？ね？ということで試してみます。

cat <<EOF > data/LaSOT/airplane/airplane-15/groundtruth.txt
197,216,106,39
EOF

では実行。

CUDA_VISIBLE_DEVICES=0 python scripts/main_inference.py

はい、予想通りできました。

(3) scripts/demo.py

これを実行すると、

$ CUDA_VISIBLE_DEVICES=0 python scripts/demo.py
--video_path ./data/LaSOT/airplane/airplane-15/img --txt_path ./data/LaSOT/airplane/airplane-150/groundtruth.txt
SAMURAI mode: False
Traceback (most recent call last):
  File "/mnt/data/shoji_noguchi/venv/samurai/samurai/scripts/demo.py", line 120, in <module>
    main(args)
  File "/mnt/data/shoji_noguchi/venv/samurai/samurai/scripts/demo.py", line 71, in main
    state = predictor.init_state(frames_or_path, offload_video_to_cpu=True)
  File "/mnt/data/shoji_noguchi/venv/samurai/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/data/shoji_noguchi/venv/samurai/samurai/./sam2/sam2/sam2_video_predictor.py", line 53, in init_state
    images, video_height, video_width = load_video_frames(
  File "/mnt/data/shoji_noguchi/venv/samurai/samurai/./sam2/sam2/utils/misc.py", line 198, in load_video_frames
    return load_video_frames_from_jpg_images(
  File "/mnt/data/shoji_noguchi/venv/samurai/samurai/./sam2/sam2/utils/misc.py", line 267, in load_video_frames_from_jpg_images
    images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)
RuntimeError: [enforce fail at alloc_cpu.cpp:117] err == 0. DefaultCPUAllocator: can't allocate memory: you tried to allocate 56623104000 bytes. Error code 12 (Cannot allocate memory)
$

メモリが足らん！となり、さらにSAMURAI modeがFALSEになるのですよね。うーん。issueで聞いてみるか。。。