object-separation

オブジェクトの分離

動画から被写体を1つ背景から分離します。すべての処理はユーザーのコンピューター上で実行され、何もアップロードされません。結果は被写体が緑色で強調表示された同じ動画なので、何が分離されたかを正確に確認できます。また、各フレームのマスクがディスクに保存されます。

Pythonコードはその都度作成して実行します。以下のコードブロックはテスト済みです。最初から作り直すのではなく、必要に応じて調整してください。

1. 削除可能なフォルダー内にセットアップする

Python 3.10以降とffmpegが必要です。プロジェクトフォルダー内の仮想環境にすべてをインストールすれば、そのフォルダーを削除するだけですべて削除できます。

python3 -m venv .object-separation
.object-separation/bin/pip install torch transformers opencv-python-headless pillow numpy

Windowsでは.object-separation\Scripts\pipを使用してください。WindowsでNVIDIAカードを使用する場合は、通常のpipでは低速なCPU版がインストールされるため、まずpytorch.orgからCUDA版のtorchをインストールしてください。以下のPythonコードはすべて、仮想環境のpythonで実行してください。

2. コンピューターをスキャンしてモデルを選ぶ

まずスキャンを実行し、検出されたデバイス、選択したモデル、ダウンロードサイズ、予想される速度をユーザーに伝えてください。

マシン	モデル	ダウンロード
12 GB以上のNVIDIA GPU	facebook/sam2.1-hiera-large	898 MB
12 GB未満のNVIDIA GPU、またはApple Silicon	facebook/sam2.1-hiera-base-plus	323 MB
GPUなし（CPUのみ）	facebook/sam2.1-hiera-tiny	156 MB

CPUでも動作しますが、低速です。1フレームあたり約5秒かかるため、30 fpsの10秒のクリップでは約30分かかります。使用するRAMは2 GB未満です。GPUならはるかに高速です。クリップが長く、GPUがない場合は、ユーザーに必要な部分だけを分離することを提案してください。

# Set these from the user's request. Frame numbers start at 0; the video must be constant frame rate.
VIDEO = "clip.mp4"
START, END = 0, 89              # frames to separate. Click points are picked on the START frame.
PTS = [[480, 270], [100, 60]]   # click points in 960-wide pixels: on the subject first, then on things to leave out
LBL = [1, 0]                    # 1 = this is the subject, 0 = this is NOT the subject

import os, platform, subprocess, numpy as np, torch, cv2
from transformers import Sam2VideoModel, Sam2VideoProcessor

print("system:", platform.system(), platform.machine())
if torch.cuda.is_available():
    gb = torch.cuda.get_device_properties(0).total_memory / 1e9
    print("gpu:", torch.cuda.get_device_name(0), f"{gb:.0f} GB")
    dev, MODEL = "cuda", "facebook/sam2.1-hiera-" + ("large" if gb >= 12 else "base-plus")
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
    dev, MODEL = "mps", "facebook/sam2.1-hiera-base-plus"
else:
    dev, MODEL = "cpu", "facebook/sam2.1-hiera-tiny"
print("device:", dev, "| model:", MODEL)
model = Sam2VideoModel.from_pretrained(MODEL).to(dev).eval()     # downloads on first run
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960   # working width. SAM 2 resizes to 1024 itself, so a smaller width is NOT faster.

3. クリックするポイントを選ぶ

座標グリッド付きのフレームSTARTを保存して確認します。被写体上のポイントを2つか3つ選び、ラベル0のポイントを、除外するもの（被写体を持つ指、パッケージ、テーブルなど）に追加します。選択内容が分かりにくい場合は、ユーザーに何を選んだか見せてください。

def read_frames(video, a, b):
    cap = cv2.VideoCapture(video); cap.set(cv2.CAP_PROP_POS_FRAMES, a); out = []
    for _ in range(a, b + 1):
        ok, im = cap.read()
        if not ok: break
        out.append(cv2.cvtColor(cv2.resize(im, (W, int(im.shape[0] * W / im.shape[1]))), cv2.COLOR_BGR2RGB))
    return out

def preview(frame, path="work/preview.jpg"):       # frame with a labelled grid, to pick click points
    os.makedirs("work", exist_ok=True)
    im = cv2.cvtColor(read_frames(VIDEO, frame, frame)[0], cv2.COLOR_RGB2BGR)
    for x in range(0, im.shape[1], 80):
        cv2.line(im, (x, 0), (x, im.shape[0]), (0, 255, 255), 1); cv2.putText(im, str(x), (x + 2, 12), 0, 0.4, (0, 255, 255), 1)
    for y in range(0, im.shape[0], 60):
        cv2.line(im, (0, y), (im.shape[1], y), (0, 255, 255), 1); cv2.putText(im, str(y), (2, y - 2), 0, 0.4, (0, 255, 255), 1)
    cv2.imwrite(path, im); return path

4. 被写体を分離する

メモリ使用量を抑えるため、動画は60フレーム単位で処理されます。各チャンクは、その直前のチャンクの最後のマスクから開始します。まず短い範囲（約20フレーム）を試して確認し、その後、全範囲を実行してください。

def separate_chunk(frames, points=None, labels=None, mask=None):
    H = frames[0].shape[0]
    s = proc.init_video_session(video=frames, inference_device=dev, video_storage_device="cpu",
                                inference_state_device="cpu", dtype=torch.float32)
    if mask is None:
        proc.add_inputs_to_inference_session(inference_session=s, frame_idx=0, obj_ids=1,
                                             input_points=[[points]], input_labels=[[labels]])
    else:
        proc.add_inputs_to_inference_session(inference_session=s, frame_idx=0, obj_ids=1,
                                             input_masks=mask.astype(np.uint8))
    res = {}
    with torch.inference_mode():
        model(inference_session=s, frame_idx=0)
        for o in model.propagate_in_video_iterator(s, start_frame_idx=0):
            m = proc.post_process_masks([o.pred_masks], original_sizes=[[H, W]], binarize=True)[0]
            res[o.frame_idx] = m.squeeze().cpu().numpy().astype(bool)
    return res

def separate(start, end, pts, lbl, ch=60):
    out, a, seed = {}, start, None
    while a < end:
        b = min(end, a + ch)
        fr = read_frames(VIDEO, a, b)
        r = separate_chunk(fr, points=pts, labels=lbl) if seed is None else separate_chunk(fr, mask=seed)
        for k, m in r.items():
            out[a + k] = m
        seed, a = r[len(fr) - 1], b
    return out

masks = separate(START, END, PTS, LBL)
os.makedirs("work/masks/subject", exist_ok=True)
for f, m in masks.items():
    cv2.imwrite(f"work/masks/subject/{f:05d}.png", m.astype(np.uint8) * 255)

5. 被写体を緑色で表示した動画を出力する

def green(f, img, fill=0.35):                       # green tint, glow and a bright edge on the subject
    m = cv2.GaussianBlur(cv2.resize(masks[f].astype(np.float32), (img.shape[1], img.shape[0])), (5, 5), 0)
    edge = cv2.morphologyEx((m > 0.5).astype(np.uint8), cv2.MORPH_GRADIENT, np.ones((7, 7), np.uint8)).astype(np.float32)
    glow = np.clip(cv2.GaussianBlur(edge, (37, 37), 0) * 2.5, 0, 1)[..., None]
    col = np.array([20, 255, 57], np.float32)          # BGR green
    img = img * (1 - m[..., None] * fill) + col * m[..., None] * fill + col * glow * 0.9
    return img * (1 - edge[..., None]) + np.array([220, 255, 200], np.float32) * edge[..., None]

def render(a, b, out, fn):                          # writes the frames a..b with the original audio
    cap = cv2.VideoCapture(VIDEO); cap.set(cv2.CAP_PROP_POS_FRAMES, a)
    fps = cap.get(cv2.CAP_PROP_FPS)
    ff = subprocess.Popen(["ffmpeg", "-v", "error", "-y", "-f", "rawvideo", "-pix_fmt", "bgr24",
                           "-s", f"{int(cap.get(3))}x{int(cap.get(4))}", "-r", str(fps), "-i", "-",
                           "-ss", f"{a / fps:.4f}", "-t", f"{(b - a + 1) / fps:.4f}", "-i", VIDEO,
                           "-map", "0:v", "-map", "1:a?", "-c:v", "libx264", "-crf", "18", "-pix_fmt", "yuv420p",
                           "-c:a", "aac", "-shortest", out], stdin=subprocess.PIPE)
    for f in range(a, b + 1):
        ok, img = cap.read()
        if not ok: break
        ff.stdin.write(np.clip(fn(f, img.astype(np.float32)), 0, 255).astype(np.uint8).tobytes())
    ff.stdin.close(); ff.wait()

render(START, END, "work/separated.mp4", green)

引き渡す前に、結果の開始、中間、および手や物体が前を横切る場面をいくつか確認してください。被写体以外のものに緑色がはみ出している場合は、その部分にラベル0のポイントを追加して、ステップ4をもう一度実行してください。被写体の一部が欠けている場合は、その部分にラベル1のポイントを追加してください。

ファイルの保存場所をユーザーに伝えてください。動画（被写体が緑色で表示されたもの）はwork/separated.mp4、フレームごとのマスク（白＝被写体）はwork/masks/subject/です。マスクはディスクに保存されるため、分離した被写体を使って次に行いたい処理は、再度分離を実行せずにマスクから開始できます。

6. クリーンアップ（任意）

ユーザーが作業を終えたら、すべて削除できます。

• 仮想環境：.object-separationフォルダーを削除します
• モデル：~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-*を削除します（Windowsの場合：C:\Users\<name>\.cache\huggingface\hub）

注意事項

• マスクの境界はくっきりしています。物体や手はきれいに分離できますが、柔らかい髪はやや苦手です。
• 被写体に付いているものは、すべて被写体の一部として扱われます（パッケージ、剥離フィルムなど）。その部分にラベル0のポイントを追加してください。
• 動画が可変フレームレートの場合は、先に変換してください：ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4。
• インターネット接続が必要なのは、モデルをダウンロードする初回実行時のみです。