object-separation

对象分离

从视频背景中分离出一个主体。所有操作都在用户的计算机上运行，不会上传任何内容。结果是同一段视频，但主体会以绿色高亮显示，让用户清楚看到分离出的内容；此外，还会将每一帧的蒙版保存到磁盘。

你可以边编写边运行 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。在配备 NVIDIA 显卡的 Windows 计算机上，请先从 pytorch.org 安装 CUDA 版 torch，因为直接使用 pip 会安装速度较慢的 CPU 版。下面所有 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 上可以运行，但速度较慢：每帧约需 5 秒，因此 30 fps 的 10 秒片段大约需要半小时。内存占用低于 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 帧并查看。选择主体上的两三个点；对于任何必须排除的内容（例如拿着主体的手指、包装或桌子），添加标签为 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。
• 只有首次运行需要联网下载模型。